diff --git a/src/order_bak/cli.py b/src/order_bak/cli.py index af3745c..707e5e9 100644 --- a/src/order_bak/cli.py +++ b/src/order_bak/cli.py @@ -57,5 +57,33 @@ def backup(ctx, start, end, scan_path): ) +@main.command() +@click.option("--csv", "csv_path", required=True, type=click.Path(exists=True), help="Path to backup CSV to delete") +@click.pass_context +def delete(ctx, csv_path): + cfg = load_config(ctx.obj["config_path"]) + from azure.data.tables import TableClient + from order_bak.delete import delete_orders + + csv_file = Path(csv_path) + + import csv as csv_mod + with open(csv_file, newline="") as f: + count = sum(1 for _ in csv_mod.DictReader(f)) + + click.echo(f"Will delete {count} orders from {csv_file}") + if not click.confirm("Continue?"): + click.echo("Aborted") + return + + table = TableClient.from_connection_string(cfg.azure_connection_string, table_name="order") + delete_orders( + table_client=table, + csv_path=csv_file, + batch_size=cfg.delete_batch_size, + delay_ms=cfg.delete_delay_ms, + ) + + if __name__ == "__main__": main() diff --git a/src/order_bak/delete.py b/src/order_bak/delete.py new file mode 100644 index 0000000..ac11ebd --- /dev/null +++ b/src/order_bak/delete.py @@ -0,0 +1,39 @@ +import csv +import time +from collections import defaultdict +from itertools import islice +from pathlib import Path + + +def _chunked(iterable, size): + it = iter(iterable) + while chunk := list(islice(it, size)): + yield chunk + + +def delete_orders( + table_client, + csv_path: Path, + batch_size: int, + delay_ms: int, +) -> int: + by_partition: dict[str, list[dict]] = defaultdict(list) + + with open(csv_path, newline="") as f: + reader = csv.DictReader(f) + for row in reader: + pk = row["PartitionKey"] + rk = row["RowKey"] + by_partition[pk].append({"PartitionKey": pk, "RowKey": rk}) + + total_deleted = 0 + for pk, entities in by_partition.items(): + for batch in _chunked(entities, batch_size): + operations = [("delete", e) for e in batch] + table_client.submit_transaction(operations) + total_deleted += len(batch) + if delay_ms > 0: + time.sleep(delay_ms / 1000.0) + + print(f"Deleted {total_deleted} orders across {len(by_partition)} partitions") + return total_deleted diff --git a/tests/test_delete.py b/tests/test_delete.py new file mode 100644 index 0000000..f5098ba --- /dev/null +++ b/tests/test_delete.py @@ -0,0 +1,67 @@ +from datetime import datetime, timezone +from unittest.mock import MagicMock + +from order_bak.csv_utils import write_csv +from order_bak.delete import delete_orders + + +def _make_entity(pk, rk): + return { + "PartitionKey": pk, + "RowKey": rk, + "OrderTime": datetime(2025, 1, 1, tzinfo=timezone.utc), + "ReceiveTime": None, + "revenue": 10.0, + "platformOrder": "p", + "prodCount": 1, + "pfId": pk, + "prodId": "prod", + "prodPrice": 10.0, + "state": 1, + "ver": 3, + "customData": None, + } + + +def test_delete_groups_by_partition_and_batches(tmp_path): + entities = [] + for i in range(3): + entities.append(_make_entity("user1", f"order{i}")) + for i in range(3, 5): + entities.append(_make_entity("user2", f"order{i}")) + + csv_path = tmp_path / "2025-01-01.csv" + write_csv(csv_path, entities) + + mock_table = MagicMock() + + deleted_count = delete_orders( + table_client=mock_table, + csv_path=csv_path, + batch_size=2, + delay_ms=0, + ) + + assert deleted_count == 5 + # user1: 3 entities, batch_size=2 -> 2 batches (2 + 1) + # user2: 2 entities, batch_size=2 -> 1 batch + assert mock_table.submit_transaction.call_count == 3 + + +def test_delete_respects_batch_size(tmp_path): + entities = [_make_entity("user1", f"order{i}") for i in range(5)] + + csv_path = tmp_path / "2025-01-01.csv" + write_csv(csv_path, entities) + + mock_table = MagicMock() + + deleted_count = delete_orders( + table_client=mock_table, + csv_path=csv_path, + batch_size=2, + delay_ms=0, + ) + + assert deleted_count == 5 + assert mock_table.submit_transaction.call_count == 3 # 2+2+1