fix: use model_dump(mode="json") for serializing STAC objects (#131)

Author: hrodmn

lib/ingestor-api/runtime/src/collection.py

@@ -1,5 +1,6 @@
 import os
 
+from fastapi import HTTPException
 from pypgstac.db import PgstacDB
 from pypgstac.load import Methods
 
@@ -18,12 +19,14 @@ def ingest(collection: StacCollection):
         creds = get_db_credentials(os.environ["DB_SECRET_ARN"])
         with PgstacDB(dsn=creds.dsn_string, debug=True) as db:
             loader = Loader(db=db)
-            collection = [
-                collection.to_dict()
-            ]  # pypgstac wants either a string or an Iterable of dicts.
-            loader.load_collections(file=collection, insert_mode=Methods.upsert)
+            loader.load_collections(
+                file=[collection.model_dump(mode="json")], insert_mode=Methods.upsert
+            )
     except Exception as e:
-        print(f"Encountered failure loading collection into pgSTAC: {e}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Encountered failure loading collection into pgSTAC: {e}",
+        ) from e
 
 
 def delete(collection_id: str):

lib/ingestor-api/runtime/src/utils.py

@@ -2,7 +2,6 @@
 
 import boto3
 import pydantic
-from fastapi.encoders import jsonable_encoder
 from pypgstac.db import PgstacDB
 from pypgstac.load import Methods
 
@@ -41,8 +40,7 @@ def load_items(creds: DbCreds, ingestions: Sequence[Ingestion]):
     with PgstacDB(dsn=creds.dsn_string, debug=True) as db:
         loader = Loader(db=db)
 
-        # serialize to JSON-friendly dicts (won't be necessary in Pydantic v2, https://github.com/pydantic/pydantic/issues/1409#issuecomment-1423995424)
-        items = jsonable_encoder(i.item for i in ingestions)
+        items = [i.item.model_dump(mode="json") for i in ingestions]
         loading_result = loader.load_items(
             file=items,
             # use insert_ignore to avoid overwritting existing items or upsert to replace

lib/ingestor-api/runtime/tests/test_collection.py

@@ -3,6 +3,7 @@
 
 import pytest
 import src.collection as collection
+from fastapi import HTTPException
 from pypgstac.load import Methods
 from src.utils import DbCreds
 
@@ -31,6 +32,56 @@ def test_load_collections(stac_collection, loader, pgstacdb):
         collection.ingest(stac_collection)
 
     loader.return_value.load_collections.assert_called_once_with(
-        file=[stac_collection.to_dict()],
+        file=[stac_collection.model_dump(mode="json")],
         insert_mode=Methods.upsert,
     )
+
+
+def test_ingest_loader_error(stac_collection, pgstacdb):
+    """Test handling of errors during the loading process."""
+    with patch(
+        "src.collection.get_db_credentials",
+        return_value=DbCreds(
+            username="", password="", host="", port=1, dbname="", engine=""
+        ),
+    ):
+        os.environ["DB_SECRET_ARN"] = ""
+
+        with patch("src.collection.Loader") as mock_loader:
+            mock_loader_instance = Mock()
+            mock_loader.return_value = mock_loader_instance
+            mock_loader_instance.load_collections.side_effect = Exception(
+                "Invalid collection data"
+            )
+
+            with pytest.raises(HTTPException) as excinfo:
+                collection.ingest(stac_collection)
+
+            assert excinfo.value.status_code == 500
+            assert "Invalid collection data" in str(excinfo.value.detail)
+
+
+def test_ingest_missing_environment_variable(stac_collection):
+    """Test handling when the required environment variable is missing."""
+    if "DB_SECRET_ARN" in os.environ:
+        del os.environ["DB_SECRET_ARN"]
+
+    with pytest.raises(HTTPException) as excinfo:
+        collection.ingest(stac_collection)
+
+    assert "DB_SECRET_ARN" in str(excinfo.value)
+
+
+def test_ingest_credentials_error(stac_collection):
+    """Test handling of errors during credential retrieval."""
+    with patch(
+        "src.collection.get_db_credentials",
+        side_effect=Exception("Failed to retrieve credentials"),
+    ):
+        os.environ["DB_SECRET_ARN"] = ""
+
+        with pytest.raises(HTTPException) as excinfo:
+            collection.ingest(stac_collection)
+
+        assert excinfo.value.status_code == 500
+        assert "Failed to retrieve credentials" in str(excinfo.value.detail)

lib/ingestor-api/runtime/tests/test_utils.py

@@ -1,7 +1,6 @@
 from unittest.mock import Mock, patch
 
 import pytest
-from fastapi.encoders import jsonable_encoder
 from pypgstac.load import Methods
 from src.utils import DbCreds
 
@@ -28,8 +27,9 @@ def dbcreds():
 def test_load_items(loader, pgstacdb, example_ingestion, dbcreds):
     import src.utils as utils
 
-    utils.load_items(dbcreds, [example_ingestion])
+    ingestions = [example_ingestion]
+    utils.load_items(dbcreds, ingestions)
     loader.return_value.load_items.assert_called_once_with(
-        file=jsonable_encoder([example_ingestion.item]),
+        file=[i.item.model_dump(mode="json") for i in ingestions],
         insert_mode=Methods.upsert,
     )