feat: exception handling for file uploads

Author: tillywoodfield

.env.test

@@ -0,0 +1,5 @@
+ENABLE_UPLOAD=1
+BUCKET_REGION="test-region"
+BUCKET_NAME="test-bucket"
+BUCKET_ACCESS_KEY_ID="test-id"
+BUCKET_ACCESS_KEY_SECRET="test-secret"

oc4ids_datastore_pipeline/__init__.py

@@ -1,13 +1,9 @@
 import logging
 import time
 
-from dotenv import load_dotenv
-
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s:%(levelname)s:%(name)s:%(message)s",
     datefmt="%Y-%m-%dT%H:%M:%S",
 )
 logging.Formatter.converter = time.gmtime
-
-load_dotenv()

oc4ids_datastore_pipeline/pipeline.py

@@ -7,6 +7,7 @@
 
 import flattentool
 import requests
+from dotenv import load_dotenv
 from libcoveoc4ids.api import oc4ids_json_output
 
 from oc4ids_datastore_pipeline.database import (
@@ -116,7 +117,7 @@ def process_dataset(dataset_name: str, dataset_url: str) -> None:
         )
         csv_path, xlsx_path = transform_to_csv_and_xlsx(json_path)
         json_public_url, csv_public_url, xlsx_public_url = upload_files(
-            json_path=json_path, csv_path=csv_path, xlsx_path=xlsx_path
+            dataset_name, json_path=json_path, csv_path=csv_path, xlsx_path=xlsx_path
         )
         save_dataset_metadata(
             dataset_name=dataset_name,
@@ -148,4 +149,5 @@ def process_registry() -> None:
 
 
 def run() -> None:
+    load_dotenv()
     process_registry()

oc4ids_datastore_pipeline/storage.py

@@ -10,8 +10,6 @@
 logger = logging.getLogger(__name__)
 
 
-# TODO:
-ENABLE_UPLOAD = os.environ.get("ENABLE_UPLOAD", "0")
 BUCKET_REGION = os.environ.get("BUCKET_REGION")
 BUCKET_NAME = os.environ.get("BUCKET_NAME")
 BUCKET_ACCESS_KEY_ID = os.environ.get("BUCKET_ACCESS_KEY_ID")
@@ -30,8 +28,7 @@ def _get_client() -> Any:
     )
 
 
-def upload_file(local_path: str, content_type: str) -> str:
-    bucket_path = os.path.relpath(local_path, "data")
+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(
@@ -40,49 +37,65 @@ def upload_file(local_path: str, content_type: str) -> str:
         bucket_path,
         ExtraArgs={"ACL": "public-read", "ContentType": content_type},
     )
-    return (
+    public_url = (
         f"https://{BUCKET_NAME}.{BUCKET_REGION}.digitaloceanspaces.com/" + bucket_path
     )
-
-
-def upload_json(json_path: str) -> str:
-    return upload_file(json_path, content_type="application/json")
-
-
-def upload_csv(csv_path: str) -> str:
-    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(zip_file_path, content_type="application/zip")
-
-
-def upload_xlsx(xlsx_path: str) -> str:
-    return upload_file(
-        xlsx_path,
-        content_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",  # noqa: E501
-    )
+    logger.info(f"Uploaded to {public_url}")
+    return public_url
+
+
+def _upload_json(dataset_id: str, json_path: str) -> Optional[str]:
+    try:
+        return _upload_file(
+            local_path=json_path,
+            bucket_path=f"{dataset_id}/{dataset_id}.json",
+            content_type="application/json",
+        )
+    except Exception as e:
+        logger.warning(f"Failed to upload {json_path} with error {e}")
+        return None
+
+
+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}")
+        return None
+
+
+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
 
 
 def upload_files(
+    dataset_id: str,
     json_path: Optional[str] = None,
     csv_path: Optional[str] = None,
     xlsx_path: Optional[str] = None,
 ) -> tuple[Optional[str], Optional[str], Optional[str]]:
     # TODO: Option to delete local files once uploaded?
-    # TODO: Exception handling
-    if not bool(int(ENABLE_UPLOAD)):
+    if not bool(int(os.environ.get("ENABLE_UPLOAD", "0"))):
         logger.info("Upload is disabled, skipping")
         return None, None, None
-    logger.info("Uploading files")
-    if json_path:
-        json_public_url = upload_json(json_path)
-        logger.info(f"Uploaded JSON file to {json_public_url}")
-    if csv_path:
-        csv_public_url = upload_csv(csv_path)
-        logger.info(f"Uploaded CSV zip file to {csv_public_url}")
-    if xlsx_path:
-        xlsx_public_url = upload_xlsx(xlsx_path)
-        logger.info(f"Uploaded XLSX file to {xlsx_public_url}")
+    json_public_url = _upload_json(dataset_id, json_path) if json_path else None
+    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

tests/conftest.py

@@ -0,0 +1,3 @@
+from dotenv import find_dotenv, load_dotenv
+
+load_dotenv(find_dotenv(".env.test"))

tests/test_storage.py

@@ -1 +1,169 @@
-# TODO: Implement tests
+import os
+import tempfile
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+from pytest_mock import MockerFixture
+
+from oc4ids_datastore_pipeline.storage import upload_files
+
+
+@pytest.fixture(autouse=True)
+def mock_client(mocker: MockerFixture) -> Any:
+    os.environ["ENABLE_UPLOAD"] = "1"
+    mock_boto3_client = MagicMock()
+    patch_boto3_client = mocker.patch(
+        "oc4ids_datastore_pipeline.storage.boto3.session.Session.client"
+    )
+    patch_boto3_client.return_value = mock_boto3_client
+    return mock_boto3_client
+
+
+def test_upload_files_upload_disabled(mock_client: MagicMock) -> None:
+    os.environ["ENABLE_UPLOAD"] = "0"
+
+    json_public_url, csv_public_url, xlsx_public_url = upload_files(
+        "test_dataset",
+        json_path="dataset.json",
+        csv_path="dataset_csv.zip",
+        xlsx_path="dataset.xlsx",
+    )
+
+    mock_client.assert_not_called()
+    assert json_public_url is None
+    assert csv_public_url is None
+    assert xlsx_public_url is None
+
+
+def test_upload_files_nothing_to_upload(mock_client: MagicMock) -> None:
+    json_public_url, csv_public_url, xlsx_public_url = upload_files("test_dataset")
+
+    mock_client.assert_not_called()
+    assert json_public_url is None
+    assert csv_public_url is None
+    assert xlsx_public_url is None
+
+
+def test_upload_files_json(mock_client: MagicMock) -> None:
+    json_public_url, csv_public_url, xlsx_public_url = upload_files(
+        "test_dataset", json_path="data/test_dataset/test_dataset.json"
+    )
+
+    mock_client.upload_file.assert_called_once_with(
+        "data/test_dataset/test_dataset.json",
+        "test-bucket",
+        "test_dataset/test_dataset.json",
+        ExtraArgs={"ACL": "public-read", "ContentType": "application/json"},
+    )
+    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 is None
+
+
+def test_upload_files_json_catches_exception(mock_client: MagicMock) -> None:
+    mock_client.upload_file.side_effect = [Exception("Mock exception"), None, None]
+
+    with tempfile.TemporaryDirectory() as csv_dir:
+        json_public_url, csv_public_url, xlsx_public_url = upload_files(
+            "test_dataset",
+            json_path="data/test_dataset/test_dataset.json",
+            csv_path=csv_dir,
+            xlsx_path="data/test_dataset/test_dataset.xlsx",
+        )
+        assert json_public_url is None
+        assert (
+            csv_public_url
+            == "https://test-bucket.test-region.digitaloceanspaces.com/test_dataset/test_dataset_csv.zip"  # noqa: E501
+        )
+        assert (
+            xlsx_public_url
+            == "https://test-bucket.test-region.digitaloceanspaces.com/test_dataset/test_dataset.xlsx"  # noqa: E501
+        )
+
+
+def test_upload_files_csv(mock_client: MagicMock) -> None:
+    with tempfile.TemporaryDirectory() as csv_dir:
+        json_public_url, csv_public_url, xlsx_public_url = upload_files(
+            "test_dataset", csv_path=csv_dir
+        )
+
+    mock_client.upload_file.assert_called_once_with(
+        f"{csv_dir}_csv.zip",
+        "test-bucket",
+        "test_dataset/test_dataset_csv.zip",
+        ExtraArgs={"ACL": "public-read", "ContentType": "application/zip"},
+    )
+    assert json_public_url is None
+    assert (
+        csv_public_url
+        == "https://test-bucket.test-region.digitaloceanspaces.com/test_dataset/test_dataset_csv.zip"  # noqa: E501
+    )
+    assert xlsx_public_url is None
+
+
+def test_upload_files_csv_catches_exception(mock_client: MagicMock) -> None:
+    mock_client.upload_file.side_effect = [None, Exception("Mock exception"), None]
+
+    with tempfile.TemporaryDirectory() as csv_dir:
+        json_public_url, csv_public_url, xlsx_public_url = upload_files(
+            "test_dataset",
+            json_path="data/test_dataset/test_dataset.json",
+            csv_path=csv_dir,
+            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_xlsx(mock_client: MagicMock) -> None:
+    json_public_url, csv_public_url, xlsx_public_url = upload_files(
+        "test_dataset", xlsx_path="data/test_dataset/test_dataset.xlsx"
+    )
+
+    mock_client.upload_file.assert_called_once_with(
+        "data/test_dataset/test_dataset.xlsx",
+        "test-bucket",
+        "test_dataset/test_dataset.xlsx",
+        ExtraArgs={
+            "ACL": "public-read",
+            "ContentType": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",  # noqa: E501
+        },
+    )
+    assert json_public_url is None
+    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_xlsx_catches_exception(mock_client: MagicMock) -> None:
+    mock_client.upload_file.side_effect = [None, None, Exception("Mock exception")]
+
+    with tempfile.TemporaryDirectory() as csv_dir:
+        json_public_url, csv_public_url, xlsx_public_url = upload_files(
+            "test_dataset",
+            json_path="data/test_dataset/test_dataset.json",
+            csv_path=csv_dir,
+            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
+            == "https://test-bucket.test-region.digitaloceanspaces.com/test_dataset/test_dataset_csv.zip"  # noqa: E501
+        )
+        assert xlsx_public_url is None