feat: delete files for datasets removed from the registry

Author: tillywoodfield

oc4ids_datastore_pipeline/pipeline.py

@@ -19,7 +19,7 @@
     fetch_registered_datasets,
     get_license_name_from_url,
 )
-from oc4ids_datastore_pipeline.storage import upload_files
+from oc4ids_datastore_pipeline.storage import delete_files_for_dataset, upload_files
 
 logger = logging.getLogger(__name__)
 
@@ -137,7 +137,7 @@ def process_deleted_datasets(registered_datasets: dict[str, str]) -> None:
     for dataset_id in deleted_datasets:
         logger.info(f"Dataset {dataset_id} is no longer in the registry, deleting")
         delete_dataset(dataset_id)
-        # TODO: Delete stored files
+        delete_files_for_dataset(dataset_id)
 
 
 def process_registry() -> None:

oc4ids_datastore_pipeline/storage.py

@@ -96,3 +96,17 @@ def upload_files(
     csv_public_url = _upload_csv(dataset_id, csv_path) if csv_path else None
     xlsx_public_url = _upload_xlsx(dataset_id, xlsx_path) if xlsx_path else None
     return json_public_url, csv_public_url, xlsx_public_url
+
+
+def delete_files_for_dataset(dataset_id: str) -> None:
+    logger.info(f"Deleting files for dataset {dataset_id}")
+    try:
+        client = _get_client()
+        response = client.list_objects_v2(Bucket=BUCKET_NAME, Prefix=dataset_id)
+        if "Contents" in response:
+            objects_to_delete = [{"Key": obj["Key"]} for obj in response["Contents"]]
+            client.delete_objects(
+                Bucket=BUCKET_NAME, Delete={"Objects": objects_to_delete}
+            )
+    except Exception as e:
+        logger.warning(f"Failed to delete files with error {e}")

tests/test_pipeline.py

@@ -111,11 +111,15 @@ def test_process_deleted_datasets(mocker: MockerFixture) -> None:
     patch_delete_dataset = mocker.patch(
         "oc4ids_datastore_pipeline.pipeline.delete_dataset"
     )
+    patch_delete_files_for_dataset = mocker.patch(
+        "oc4ids_datastore_pipeline.pipeline.delete_files_for_dataset"
+    )
 
     registered_datasets = {"test_dataset": "https://test_dataset.json"}
     process_deleted_datasets(registered_datasets)
 
     patch_delete_dataset.assert_called_once_with("old_dataset")
+    patch_delete_files_for_dataset.assert_called_once_with("old_dataset")
 
 
 def test_process_dataset_catches_exception(mocker: MockerFixture) -> None:

tests/test_storage.py

@@ -6,7 +6,7 @@
 import pytest
 from pytest_mock import MockerFixture
 
-from oc4ids_datastore_pipeline.storage import upload_files
+from oc4ids_datastore_pipeline.storage import delete_files_for_dataset, upload_files
 
 
 @pytest.fixture(autouse=True)
@@ -185,3 +185,32 @@ def test_upload_files_xlsx_catches_upload_exception(mock_client: MagicMock) -> N
             == "https://test-bucket.test-region.digitaloceanspaces.com/test_dataset/test_dataset_csv.zip"  # noqa: E501
         )
         assert xlsx_public_url is None
+
+
+def test_delete_files_for_dataset(mock_client: MagicMock) -> None:
+    mock_client.list_objects_v2.return_value = {
+        "Contents": [
+            {"Key": "test_dataset/test_dataset.json"},
+            {"Key": "test_dataset/test_dataset_csv.zip"},
+            {"Key": "test_dataset/test_dataset.xlsx"},
+        ]
+    }
+
+    delete_files_for_dataset("test_dataset")
+
+    mock_client.delete_objects.assert_called_once_with(
+        Bucket="test-bucket",
+        Delete={
+            "Objects": [
+                {"Key": "test_dataset/test_dataset.json"},
+                {"Key": "test_dataset/test_dataset_csv.zip"},
+                {"Key": "test_dataset/test_dataset.xlsx"},
+            ]
+        },
+    )
+
+
+def test_delete_files_for_dataset_catches_exception(mock_client: MagicMock) -> None:
+    mock_client.list_objects_v2.side_effect = Exception("Mock exception")
+
+    delete_files_for_dataset("test_dataset")