chore: refactor stac-loader handler into component functions

Author: hrodmn

lib/stac-loader/runtime/src/stac_loader/handler.py

@@ -3,6 +3,7 @@
 import logging
 import os
 from collections import defaultdict
+from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     Annotated,
@@ -157,6 +158,211 @@ def process_s3_event(message_str: str) -> Dict[str, Any]:
         raise
 
 
+def parse_message_data(message_id: str, message_str: str) -> Dict[str, Any]:
+    """Parse message data, handling both S3 events and direct STAC JSON."""
+    if is_s3_event(message_str):
+        logger.debug(f"[{message_id}] Processing S3 event notification")
+        return process_s3_event(message_str)
+    else:
+        return json.loads(message_str)
+
+
+def store_item_if_newer(
+    items_by_collection: DefaultDict[
+        str, Dict[str, Tuple[Dict[str, Any], str, datetime]]
+    ],
+    item: Item,
+    message_id: str,
+    sns_timestamp: datetime,
+) -> None:
+    """Store item if it's newer than existing version."""
+    if not item.collection:
+        raise KeyError(f"item {item.id} is missing a collection id!")
+
+    existing = items_by_collection[item.collection].get(item.id)
+    if existing is None or sns_timestamp > existing[2]:
+        if existing:
+            logger.debug(
+                f"[{message_id}] Replacing older version of item {item.id} "
+                f"(old timestamp: {existing[2]}, new timestamp: {sns_timestamp})"
+            )
+        items_by_collection[item.collection][item.id] = (
+            item.model_dump(mode="json"),
+            message_id,
+            sns_timestamp,
+        )
+    else:
+        logger.debug(
+            f"[{message_id}] Skipping older version of item {item.id} "
+            f"(existing timestamp: {existing[2]}, message timestamp: {sns_timestamp})"
+        )
+
+
+def store_collection_if_newer(
+    collections_dict: DefaultDict[str, Tuple[Dict[str, Any], str, datetime]],
+    collection: Collection,
+    message_id: str,
+    sns_timestamp: datetime,
+) -> None:
+    """Store collection if it's newer than existing version."""
+    existing = collections_dict.get(collection.id)
+    if existing is None or sns_timestamp > existing[2]:
+        if existing:
+            logger.debug(
+                f"[{message_id}] Replacing older version of collection {collection.id} "
+                f"(old timestamp: {existing[2]}, new timestamp: {sns_timestamp})"
+            )
+        collections_dict[collection.id] = (
+            collection.model_dump(mode="json"),
+            message_id,
+            sns_timestamp,
+        )
+    else:
+        logger.debug(
+            f"[{message_id}] Skipping older version of collection {collection.id} "
+            f"(existing timestamp: {existing[2]}, message timestamp: {sns_timestamp})"
+        )
+
+
+def process_record(
+    record: Dict[str, Any],
+    collections_dict: DefaultDict[str, Tuple[Dict[str, Any], str, datetime]],
+    items_by_collection: DefaultDict[
+        str, Dict[str, Tuple[Dict[str, Any], str, datetime]]
+    ],
+) -> Optional[BatchItemFailure]:
+    """Process a single SQS record and return failure if processing fails."""
+    message_id = record.get("messageId")
+    if not message_id:
+        logger.warning("Record missing messageId, cannot report failure for it.")
+        return None
+
+    try:
+        sqs_body_str = record["body"]
+        logger.debug(f"[{message_id}] SQS message body: {sqs_body_str}")
+        sns_notification = json.loads(sqs_body_str)
+
+        message_str = sns_notification["Message"]
+        logger.debug(f"[{message_id}] SNS Message content: {message_str}")
+
+        sns_timestamp_str = sns_notification["Timestamp"]
+        sns_timestamp = datetime.fromisoformat(sns_timestamp_str.replace("Z", "+00:00"))
+        logger.debug(f"[{message_id}] SNS Timestamp: {sns_timestamp}")
+
+        message_data = parse_message_data(message_id, message_str)
+
+        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)
+            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']}"
+            )
+
+        logger.debug(f"[{message_id}] Successfully processed.")
+        return None
+
+    except (ValueError, KeyError, ValidationError, json.JSONDecodeError) as e:
+        logger.error(f"[{message_id}] Failed with error: {e}", extra=record)
+        return {"itemIdentifier": message_id}
+    except Exception as e:
+        logger.error(f"[{message_id}] Unexpected error: {e}", extra=record)
+        return {"itemIdentifier": message_id}
+
+
+def load_collections_to_db(
+    collections_dict: DefaultDict[str, Tuple[Dict[str, Any], str, datetime]],
+    pgstac_dsn: str,
+) -> List[BatchItemFailure]:
+    """Load collections to database and return failures."""
+    if not collections_dict:
+        return []
+
+    collections = [collection for collection, _, _ in collections_dict.values()]
+    message_ids = [msg_id for _, msg_id, _ in collections_dict.values()]
+
+    try:
+        with PgstacDB(dsn=pgstac_dsn) as db:
+            loader = Loader(db=db)
+            logger.info("loading collections into database.")
+            loader.load_collections(
+                file=collections,  # type: ignore
+                insert_mode=Methods.upsert,
+            )
+            logger.info(f"successfully loaded {len(collections)} collections.")
+        return []
+    except Exception as e:
+        logger.error(f"failed to load collections: {str(e)}")
+        return [{"itemIdentifier": message_id} for message_id in message_ids]
+
+
+def ensure_collection_exists(
+    db: PgstacDB, loader: Loader, collection_id: str, items: List[Dict[str, Any]]
+) -> None:
+    """Create a placeholder collection if it doesn't exist and environment allows."""
+    if not os.getenv("CREATE_COLLECTIONS_IF_MISSING"):
+        return
+
+    collection_exists = db.query_one(
+        f"SELECT count(*) as count from collections where id = '{collection_id}'"
+    )
+    if not collection_exists:
+        logger.info(
+            f"[{collection_id}] loading collection into database because it is missing."
+        )
+        collection = Collection(
+            id=collection_id,
+            description=collection_id,
+            links=Links([Link(href="placeholder", rel="self")]),
+            type="Collection",
+            license="proprietary",
+            extent=Extent(
+                spatial=SpatialExtent(bbox=[(-180, -90, 180, 90)]),
+                temporal=TimeInterval(interval=[[None, None]]),
+            ),
+            stac_version=items[0]["stac_version"],
+        )
+        loader.load_collections(
+            [collection.model_dump()],  # type: ignore
+            insert_mode=Methods.upsert,
+        )
+
+
+def load_items_for_collection(
+    collection_id: str,
+    items_dict: Dict[str, Tuple[Dict[str, Any], str, datetime]],
+    pgstac_dsn: str,
+) -> List[BatchItemFailure]:
+    """Load items for a single collection to database and return failures."""
+    items = [item_data for item_data, _, _ in items_dict.values()]
+    message_ids = [msg_id for _, msg_id, _ in items_dict.values()]
+
+    logger.debug(
+        f"[{collection_id}] Processing {len(items)} unique items from {len(items_dict)} dict entries. Item IDs: {list(items_dict.keys())}"
+    )
+
+    try:
+        with PgstacDB(dsn=pgstac_dsn) as db:
+            loader = Loader(db=db)
+            ensure_collection_exists(db, loader, collection_id, items)
+
+            logger.info(f"[{collection_id}] loading items into database.")
+            loader.load_items(
+                file=items,  # type: ignore
+                insert_mode=Methods.upsert,
+            )
+            logger.info(f"[{collection_id}] successfully loaded {len(items)} items.")
+        return []
+    except Exception as e:
+        logger.error(f"[{collection_id}] failed to load items: {str(e)}")
+        return [{"itemIdentifier": msg_id} for msg_id in message_ids]
+
+
 def handler(
     event: Dict[str, Any], context: Context
 ) -> Optional[PartialBatchFailureResponse]:
@@ -171,131 +377,25 @@ def handler(
     pgstac_dsn = get_pgstac_dsn()
 
     batch_failures: List[BatchItemFailure] = []
-
-    collections_dict: DefaultDict[str, Tuple[Dict[str, Any], str]] = defaultdict()
-    # Track items by collection and item id to deduplicate within a batch
-    # Maps: collection_id -> item_id -> (item_data, message_id)
-    items_by_collection: DefaultDict[str, Dict[str, Tuple[Dict[str, Any], str]]] = (
-        defaultdict(dict)
+    collections_dict: DefaultDict[str, Tuple[Dict[str, Any], str, datetime]] = (
+        defaultdict()
     )
+    items_by_collection: DefaultDict[
+        str, Dict[str, Tuple[Dict[str, Any], str, datetime]]
+    ] = defaultdict(dict)
 
     for record in records:
-        message_id = record.get("messageId")
-        if not message_id:
-            logger.warning("Record missing messageId, cannot report failure for it.")
-            continue
+        failure = process_record(record, collections_dict, items_by_collection)
+        if failure:
+            batch_failures.append(failure)
 
-        try:
-            sqs_body_str = record["body"]
-            logger.debug(f"[{message_id}] SQS message body: {sqs_body_str}")
-            sns_notification = json.loads(sqs_body_str)
-
-            message_str = sns_notification["Message"]
-            logger.debug(f"[{message_id}] SNS Message content: {message_str}")
-
-            if is_s3_event(message_str):
-                logger.debug(f"[{message_id}] Processing S3 event notification")
-                message_data = process_s3_event(message_str)
-            else:
-                message_data = json.loads(message_str)
-
-            if message_data["type"] == "Feature":
-                item = Item(**message_data)
-
-                if not item.collection:
-                    raise KeyError(f"item {item.id} is missing a collection id!")
-
-                # Store item by id, replacing any previous version in this batch
-                items_by_collection[item.collection][item.id] = (
-                    item.model_dump(mode="json"),
-                    message_id,
-                )
-            elif message_data["type"] == "Collection":
-                collection = Collection(**message_data)
-                collections_dict[collection.id] = (
-                    collection.model_dump(mode="json"),
-                    message_id,
-                )
-            else:
-                raise ValueError(
-                    f"expected either a 'Feature' or a 'Collection', received a {message_data['type']}"
-                )
-
-            logger.debug(f"[{message_id}] Successfully processed.")
-
-        except (ValueError, KeyError, ValidationError, json.JSONDecodeError) as e:
-            logger.error(f"[{message_id}] Failed with error: {e}", extra=record)
-            batch_failures.append({"itemIdentifier": message_id})
-        except Exception as e:
-            logger.error(f"[{message_id}] Unexpected error: {e}", extra=record)
-            batch_failures.append({"itemIdentifier": message_id})
-
-    if collections_dict:
-        collections = [collection for collection, _ in collections_dict.values()]
-        message_ids = [msg_id for _, msg_id in collections_dict.values()]
-        try:
-            with PgstacDB(dsn=pgstac_dsn) as db:
-                loader = Loader(db=db)
-                logger.info("loading collections into database.")
-                loader.load_collections(
-                    file=collections,  # type: ignore
-                    insert_mode=Methods.upsert,
-                )
-                logger.info(f"successfully loaded {len(collections)} collections.")
-        except Exception as e:
-            logger.error(f"failed to load collections: {str(e)}")
-            batch_failures.extend(
-                [{"itemIdentifier": message_id} for message_id in message_ids]
-            )
+    batch_failures.extend(load_collections_to_db(collections_dict, pgstac_dsn))
 
     for collection_id, items_dict in items_by_collection.items():
-        # Extract items and message_ids from the dict structure
-        items = [item_data for item_data, _ in items_dict.values()]
-        message_ids = [msg_id for _, msg_id in items_dict.values()]
-
-        logger.debug(
-            f"[{collection_id}] Processing {len(items)} unique items from {len(items_dict)} dict entries. Item IDs: {list(items_dict.keys())}"
+        batch_failures.extend(
+            load_items_for_collection(collection_id, items_dict, pgstac_dsn)
         )
 
-        try:
-            with PgstacDB(dsn=pgstac_dsn) as db:
-                loader = Loader(db=db)
-                if os.getenv("CREATE_COLLECTIONS_IF_MISSING"):
-                    collection_exists = db.query_one(
-                        f"SELECT count(*) as count from collections where id = '{collection_id}'"
-                    )
-                    if not collection_exists:
-                        logger.info(
-                            f"[{collection_id}] loading collection into database because it is missing."
-                        )
-                        collection = Collection(
-                            id=collection_id,
-                            description=collection_id,
-                            links=Links([Link(href="placeholder", rel="self")]),
-                            type="Collection",
-                            license="proprietary",
-                            extent=Extent(
-                                spatial=SpatialExtent(bbox=[[-180, -90, 180, 90]]),
-                                temporal=TimeInterval(interval=[[None, None]]),
-                            ),
-                            stac_version=items[0]["stac_version"],
-                        )
-                        loader.load_collections(
-                            [collection.model_dump()],  # type: ignore
-                            insert_mode=Methods.upsert,
-                        )
-
-                logger.info(f"[{collection_id}] loading items into database.")
-                loader.load_items(
-                    file=items,  # type: ignore
-                    insert_mode=Methods.upsert,
-                )
-                logger.info(f"[{collection_id}] successfully loaded {len(items)} items.")
-        except Exception as e:
-            logger.error(f"[{collection_id}] failed to load items: {str(e)}")
-
-            batch_failures.extend([{"itemIdentifier": msg_id} for msg_id in message_ids])
-
     if batch_failures:
         logger.warning(
             f"Finished processing batch. {len(batch_failures)} failure(s) reported."