feat: refactor exception handling

Author: tillywoodfield

oc4ids_datastore_pipeline/storage.py

@@ -28,61 +28,58 @@ def _get_client() -> Any:
     )
 
 
-def _upload_file(local_path: str, bucket_path: str, content_type: str) -> str:
-    logger.info(f"Uploading file {local_path}")
-    client = _get_client()
-    client.upload_file(
-        local_path,
-        BUCKET_NAME,
-        bucket_path,
-        ExtraArgs={"ACL": "public-read", "ContentType": content_type},
-    )
-    public_url = (
-        f"https://{BUCKET_NAME}.{BUCKET_REGION}.digitaloceanspaces.com/" + bucket_path
-    )
-    logger.info(f"Uploaded to {public_url}")
-    return public_url
-
-
-def _upload_json(dataset_id: str, json_path: str) -> Optional[str]:
+def _upload_file(local_path: str, bucket_path: str, content_type: str) -> Optional[str]:
     try:
-        return _upload_file(
-            local_path=json_path,
-            bucket_path=f"{dataset_id}/{dataset_id}.json",
-            content_type="application/json",
+        logger.info(f"Uploading file {local_path}")
+        client = _get_client()
+        client.upload_file(
+            local_path,
+            BUCKET_NAME,
+            bucket_path,
+            ExtraArgs={"ACL": "public-read", "ContentType": content_type},
         )
+        public_url = (
+            f"https://{BUCKET_NAME}.{BUCKET_REGION}.digitaloceanspaces.com/"
+            + bucket_path
+        )
+        logger.info(f"Uploaded to {public_url}")
+        return public_url
     except Exception as e:
-        logger.warning(f"Failed to upload {json_path} with error {e}")
+        logger.warning(f"Failed to upload {local_path} with error {e}")
         return None
 
 
+def _upload_json(dataset_id: str, json_path: str) -> Optional[str]:
+    return _upload_file(
+        local_path=json_path,
+        bucket_path=f"{dataset_id}/{dataset_id}.json",
+        content_type="application/json",
+    )
+
+
 def _upload_csv(dataset_id: str, csv_path: str) -> Optional[str]:
     try:
         directory = Path(csv_path)
         zip_file_path = f"{csv_path}_csv.zip"
         with zipfile.ZipFile(zip_file_path, mode="w") as archive:
             for file_path in directory.rglob("*"):
                 archive.write(file_path, arcname=file_path.relative_to(directory))
-        return _upload_file(
-            local_path=zip_file_path,
-            bucket_path=f"{dataset_id}/{dataset_id}_csv.zip",
-            content_type="application/zip",
-        )
     except Exception as e:
-        logger.warning(f"Failed to upload {csv_path} with error {e}")
+        logger.warning(f"Failed to zip {csv_path} with error {e}")
         return None
+    return _upload_file(
+        local_path=zip_file_path,
+        bucket_path=f"{dataset_id}/{dataset_id}_csv.zip",
+        content_type="application/zip",
+    )
 
 
 def _upload_xlsx(dataset_id: str, xlsx_path: str) -> Optional[str]:
-    try:
-        return _upload_file(
-            local_path=xlsx_path,
-            bucket_path=f"{dataset_id}/{dataset_id}.xlsx",
-            content_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",  # noqa: E501
-        )
-    except Exception as e:
-        logger.warning(f"Failed to upload {xlsx_path} with error {e}")
-        return None
+    return _upload_file(
+        local_path=xlsx_path,
+        bucket_path=f"{dataset_id}/{dataset_id}.xlsx",
+        content_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",  # noqa: E501
+    )
 
 
 def upload_files(

tests/test_storage.py

@@ -64,7 +64,7 @@ def test_upload_files_json(mock_client: MagicMock) -> None:
     assert xlsx_public_url is None
 
 
-def test_upload_files_json_catches_exception(mock_client: MagicMock) -> None:
+def test_upload_files_json_catches_upload_exception(mock_client: MagicMock) -> None:
     mock_client.upload_file.side_effect = [Exception("Mock exception"), None, None]
 
     with tempfile.TemporaryDirectory() as csv_dir:
@@ -105,7 +105,25 @@ def test_upload_files_csv(mock_client: MagicMock) -> None:
     assert xlsx_public_url is None
 
 
-def test_upload_files_csv_catches_exception(mock_client: MagicMock) -> None:
+def test_upload_files_csv_catches_zip_exception() -> None:
+    json_public_url, csv_public_url, xlsx_public_url = upload_files(
+        "test_dataset",
+        json_path="data/test_dataset/test_dataset.json",
+        csv_path="non/existent/directory",
+        xlsx_path="data/test_dataset/test_dataset.xlsx",
+    )
+    assert (
+        json_public_url
+        == "https://test-bucket.test-region.digitaloceanspaces.com/test_dataset/test_dataset.json"  # noqa: E501
+    )
+    assert csv_public_url is None
+    assert (
+        xlsx_public_url
+        == "https://test-bucket.test-region.digitaloceanspaces.com/test_dataset/test_dataset.xlsx"  # noqa: E501
+    )
+
+
+def test_upload_files_csv_catches_upload_exception(mock_client: MagicMock) -> None:
     mock_client.upload_file.side_effect = [None, Exception("Mock exception"), None]
 
     with tempfile.TemporaryDirectory() as csv_dir:
@@ -148,7 +166,7 @@ def test_upload_files_xlsx(mock_client: MagicMock) -> None:
     )
 
 
-def test_upload_files_xlsx_catches_exception(mock_client: MagicMock) -> None:
+def test_upload_files_xlsx_catches_upload_exception(mock_client: MagicMock) -> None:
     mock_client.upload_file.side_effect = [None, None, Exception("Mock exception")]
 
     with tempfile.TemporaryDirectory() as csv_dir: