diff --git a/lib/stac-loader/runtime/src/stac_loader/handler.py b/lib/stac-loader/runtime/src/stac_loader/handler.py index 127dfe5..4bf2c86 100644 --- a/lib/stac-loader/runtime/src/stac_loader/handler.py +++ b/lib/stac-loader/runtime/src/stac_loader/handler.py @@ -19,6 +19,7 @@ from pydantic import ValidationError from pypgstac.db import PgstacDB from pypgstac.load import Loader, Methods +from stac_pydantic.catalog import Catalog from stac_pydantic.collection import Collection, Extent, SpatialExtent, TimeInterval from stac_pydantic.item import Item from stac_pydantic.links import Link, Links @@ -180,7 +181,7 @@ def store_item_if_newer( def store_collection_if_newer( collections_dict: CollectionRecords, - collection: Collection, + collection: Collection | Catalog, message_id: str, sns_timestamp: datetime, ) -> None: @@ -232,14 +233,17 @@ def process_record( if message_data["type"] == "Feature": item = Item(**message_data) store_item_if_newer(items_by_collection, item, message_id, sns_timestamp) - elif message_data["type"] == "Collection": - collection = Collection(**message_data) + elif message_data["type"] in ("Collection", "Catalog"): + collection_type = ( + Collection if message_data["type"] == "Collection" else Catalog + ) + collection = collection_type(**message_data) store_collection_if_newer( collections_dict, collection, message_id, sns_timestamp ) else: raise ValueError( - f"expected either a 'Feature' or a 'Collection', received a {message_data['type']}" + f"expected a 'Feature', 'Collection', or 'Catalog', received a {message_data['type']}" ) logger.debug(f"[{message_id}] Successfully processed.") diff --git a/lib/stac-loader/runtime/tests/test_stac_loader_handler.py b/lib/stac-loader/runtime/tests/test_stac_loader_handler.py index ab793c2..f51795b 100644 --- a/lib/stac-loader/runtime/tests/test_stac_loader_handler.py +++ b/lib/stac-loader/runtime/tests/test_stac_loader_handler.py @@ -63,6 +63,17 @@ def create_valid_stac_collection(collection_id="test-collection"): } +def create_valid_stac_catalog(catalog_id="test-catalog"): + """Create a valid STAC catalog.""" + return { + "id": catalog_id, + "type": "Catalog", + "description": f"A test catalog with ID {catalog_id}", + "links": [], + "stac_version": "1.1.0", + } + + def test_get_pgstac_dsn_missing_env_var(): """Test get_pgstac_dsn when environment variable is missing""" # Save current env var if it exists @@ -1053,13 +1064,30 @@ def test_handler_with_collection_load_error( assert any(f["itemIdentifier"] == message_id for f in result["batchItemFailures"]) +def test_handler_with_valid_catalog(mock_aws_context, mock_pgstac_dsn, database_url): + """Test handler with a valid STAC catalog.""" + catalog_id = "test-catalog" + valid_catalog = create_valid_stac_catalog(catalog_id=catalog_id) + event = { + "Records": [create_sqs_record(valid_catalog, message_id="test-catalog-message-1")] + } + + result = handler(event, mock_aws_context) + + assert result is None + stored_catalog = get_collection(database_url, catalog_id) + assert stored_catalog is not None + assert stored_catalog["content"]["type"] == "Catalog" + assert "extent" not in stored_catalog["content"] + assert "license" not in stored_catalog["content"] + + def test_handler_with_unknown_type(mock_aws_context, mock_pgstac_dsn): - """Test handler with unknown STAC type (neither Feature nor Collection)""" - # Create an object with unknown type + """Test handler with an unsupported STAC type.""" unknown_object = { "id": "test-unknown", - "type": "Catalog", # Neither Feature nor Collection - "description": "A test catalog", + "type": "Unsupported", + "description": "An unsupported STAC object", "stac_version": "1.1.0", }