Fix ingestor regressions introduced by /pull/18. (#19)

* Refactor use of load_into_pgstac and move collection load logic.

* Fix regression not included in be40af7

* Fix uncaught regression accessing collection property from an item dict.

* Linting fixes.

Author: sharkinsspatial

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

@@ -1,14 +1,10 @@
 import os
 
 from pypgstac.db import PgstacDB
+from pypgstac.load import Methods
 
 from .schemas import StacCollection
-from .utils import (
-    IngestionType,
-    convert_decimals_to_float,
-    get_db_credentials,
-    load_into_pgstac,
-)
+from .utils import get_db_credentials
 from .vedaloader import VEDALoader
 
 
@@ -18,12 +14,13 @@ def ingest(collection: StacCollection):
     does necessary preprocessing,
     and loads into the PgSTAC collection table
     """
-    creds = get_db_credentials(os.environ["DB_SECRET_ARN"])
-    collection = [
-        convert_decimals_to_float(collection.dict(by_alias=True, exclude_unset=True))
-    ]
-    with PgstacDB(dsn=creds.dsn_string, debug=True) as db:
-        load_into_pgstac(db=db, ingestions=collection, table=IngestionType.collections)
+    try:
+        creds = get_db_credentials(os.environ["DB_SECRET_ARN"])
+        with PgstacDB(dsn=creds.dsn_string, debug=True) as db:
+            loader = VEDALoader(db=db)
+            loader.load_collection(file=collection, insert_mode=Methods.upsert)
+    except Exception as e:
+        print(f"Encountered failure loading collection into pgSTAC: {e}")
 
 
 def delete(collection_id: str):

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

@@ -3,16 +3,11 @@
 from typing import TYPE_CHECKING, Iterator, List, Optional, Sequence
 
 from boto3.dynamodb.types import TypeDeserializer
-from pypgstac.db import PgstacDB
 
-from .dependencies import get_settings, get_table
+from .config import settings
+from .dependencies import get_table
 from .schemas import Ingestion, Status
-from .utils import (
-    IngestionType,
-    convert_decimals_to_float,
-    get_db_credentials,
-    load_into_pgstac,
-)
+from .utils import get_db_credentials, load_items
 
 if TYPE_CHECKING:
     from aws_lambda_typing import context as context_
@@ -43,7 +38,7 @@ def update_dynamodb(
     """
     # Update records in DynamoDB
     print(f"Updating ingested items status in DynamoDB, marking as {status}...")
-    table = get_table(get_settings())
+    table = get_table(settings)
     with table.batch_writer(overwrite_by_pkeys=["created_by", "id"]) as batch:
         for ingestion in ingestions:
             batch.put_item(
@@ -64,24 +59,14 @@ def handler(event: "events.DynamoDBStreamEvent", context: "context_.Context"):
         print("No queued ingestions to process")
         return
 
-    items = [
-        # NOTE: Important to deserialize values to convert decimals to floats
-        convert_decimals_to_float(ingestion.item)
-        for ingestion in ingestions
-    ]
-
-    creds = get_db_credentials(os.environ["DB_SECRET_ARN"])
-
     # Insert into PgSTAC DB
     outcome = Status.succeeded
     message = None
     try:
-        with PgstacDB(dsn=creds.dsn_string, debug=True) as db:
-            load_into_pgstac(
-                db=db,
-                ingestions=items,
-                table=IngestionType.items,
-            )
+        load_items(
+            creds=get_db_credentials(os.environ["DB_SECRET_ARN"]),
+            ingestions=ingestions,
+        )
     except Exception as e:
         print(f"Encountered failure loading items into pgSTAC: {e}")
         outcome = Status.failed

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

@@ -1,23 +1,17 @@
 import decimal
 import json
-from enum import Enum
-from typing import Any, Dict, Sequence, Union
+from typing import Any, Dict, Sequence
 
 import boto3
 import orjson
 import pydantic
 from pypgstac.db import PgstacDB
 from pypgstac.load import Methods
 
-from .schemas import AccessibleItem, StacCollection
+from .schemas import Ingestion
 from .vedaloader import VEDALoader
 
 
-class IngestionType(str, Enum):
-    collections = "collections"
-    items = "items"
-
-
 class DbCreds(pydantic.BaseModel):
     username: str
     password: str
@@ -62,47 +56,29 @@ def decimal_to_float(obj):
     )
 
 
-def load_items(items: Sequence[AccessibleItem], loader):
+def load_items(creds: DbCreds, ingestions: Sequence[Ingestion]):
     """
-    Loads items into the PgSTAC database and
-    updates the summaries and extent for the collections involved
-    """
-    loading_result = loader.load_items(
-        file=items,
-        # use insert_ignore to avoid overwritting existing items or upsert to replace
-        insert_mode=Methods.upsert,
-    )
-
-    # Trigger update on summaries and extents
-    collections = set([item["collection"] for item in items])
-    for collection in collections:
-        loader.update_collection_summaries(collection)
-
-    return loading_result
-
-
-def load_collection(collection: Sequence[StacCollection], loader):
-    """
-    Loads the collection to the PgSTAC database
+    Bulk insert STAC records into pgSTAC.
     """
-    return loader.load_collections(
-        file=collection,
-        # use insert_ignore to avoid overwritting existing items or upsert to replace
-        insert_mode=Methods.upsert,
-    )
+    with PgstacDB(dsn=creds.dsn_string, debug=True) as db:
+        loader = VEDALoader(db=db)
+
+        items = [
+            # NOTE: Important to deserialize values to convert decimals to floats
+            convert_decimals_to_float(i.item)
+            for i in ingestions
+        ]
+
+        print(f"Ingesting {len(items)} items")
+        loading_result = loader.load_items(
+            file=items,
+            # use insert_ignore to avoid overwritting existing items or upsert to replace
+            insert_mode=Methods.upsert,
+        )
 
+        # Trigger update on summaries and extents
+        collections = set([item["collection"] for item in items])
+        for collection in collections:
+            loader.update_collection_summaries(collection)
 
-def load_into_pgstac(
-    db: "PgstacDB",
-    ingestions: Union[Sequence[AccessibleItem], Sequence[StacCollection]],
-    table: IngestionType,
-):
-    """
-    Bulk insert STAC records into pgSTAC.
-    The ingestion can be items or collection, determined by the `table` arg.
-    """
-    loader = VEDALoader(db=db)
-    loading_function = load_items
-    if table == IngestionType.collections:
-        loading_function = load_collection
-    return loading_function(ingestions, loader)
+        return loading_result