[core]: Models, containers for items & events (#169)

* [core]: Models, containers for items & events

In preparation for low-latency, we're adding / updating a handful of
models, some of which are persisted to Cosmos DB.

This adds the pydantic models for that.

Author: Tom Augspurger

deployment/terraform/resources/cosmosdb.tf

@@ -79,4 +79,24 @@ resource "azurerm_cosmosdb_sql_container" "leases" {
   database_name         = azurerm_cosmosdb_sql_database.pctasks.name
   partition_key_path    = "/id"
   partition_key_version = 1
-}
+}
+
+## Storage Events (for low-latency ingest)
+
+resource "azurerm_cosmosdb_sql_container" "storage-events" {
+  name                = "storage-events"
+  resource_group_name = data.azurerm_cosmosdb_account.pctasks.resource_group_name
+  account_name        = data.azurerm_cosmosdb_account.pctasks.name
+  database_name       = azurerm_cosmosdb_sql_database.pctasks.name
+  partition_key_path  = "/id"
+}
+
+## Items (for source of truth and low-latency ingest)
+
+resource "azurerm_cosmosdb_sql_container" "items" {
+  name                = "items"
+  resource_group_name = data.azurerm_cosmosdb_account.pctasks.resource_group_name
+  account_name        = data.azurerm_cosmosdb_account.pctasks.name
+  database_name       = azurerm_cosmosdb_sql_database.pctasks.name
+  partition_key_path  = "/stac_id"
+}

pctasks/core/pctasks/core/cosmos/container.py

@@ -35,7 +35,7 @@
 )
 from pctasks.core.cosmos.page import Page
 from pctasks.core.cosmos.settings import CosmosDBSettings
-from pctasks.core.models.run import Record
+from pctasks.core.models.record import Record
 from pctasks.core.models.utils import tzutc_now
 from pctasks.core.utils import grouped
 from pctasks.core.utils.backoff import BackoffStrategy, with_backoff, with_backoff_async

pctasks/core/pctasks/core/cosmos/containers/items.py

@@ -0,0 +1,46 @@
+from typing import Dict, Optional, Type, TypeVar
+
+from pydantic import BaseModel
+
+from pctasks.core.cosmos.container import (
+    ContainerOperation,
+    CosmosDBContainer,
+    CosmosDBDatabase,
+    TriggerType,
+)
+from pctasks.core.cosmos.settings import CosmosDBSettings
+from pctasks.core.models.item import ItemUpdatedRecord, StacItemRecord
+
+# Records that this container can hold
+T = TypeVar(
+    "T",
+    StacItemRecord,
+    ItemUpdatedRecord,
+)
+
+PARTITION_KEY = "/stac_id"
+
+STORED_PROCEDURES: Dict[ContainerOperation, Dict[Type[BaseModel], str]] = {}
+
+TRIGGERS: Dict[ContainerOperation, Dict[TriggerType, str]] = {}
+
+
+class ItemsContainer(CosmosDBContainer[T]):
+    def __init__(
+        self,
+        model_type: Type[T],
+        db: Optional[CosmosDBDatabase] = None,
+        settings: Optional[CosmosDBSettings] = None,
+    ) -> None:
+        super().__init__(
+            lambda settings: settings.get_items_container_name(),
+            PARTITION_KEY,
+            model_type=model_type,  # type: ignore[arg-type]
+            db=db,
+            settings=settings,
+            stored_procedures=STORED_PROCEDURES,
+            triggers=TRIGGERS,
+        )
+
+    def get_partition_key(self, model: T) -> str:
+        return model.stac_id

pctasks/core/pctasks/core/cosmos/containers/storage_events.py

@@ -0,0 +1,41 @@
+from typing import Dict, Optional, Type
+
+from pydantic import BaseModel
+
+from pctasks.core.cosmos.container import (
+    ContainerOperation,
+    CosmosDBContainer,
+    CosmosDBDatabase,
+    TriggerType,
+)
+from pctasks.core.cosmos.settings import CosmosDBSettings
+from pctasks.core.models.event import StorageEventRecord
+
+T = StorageEventRecord
+
+PARTITION_KEY = "/id"
+
+STORED_PROCEDURES: Dict[ContainerOperation, Dict[Type[BaseModel], str]] = {}
+
+TRIGGERS: Dict[ContainerOperation, Dict[TriggerType, str]] = {}
+
+
+class StorageEventsContainer(CosmosDBContainer[T]):
+    def __init__(
+        self,
+        model_type: Type[T],
+        db: Optional[CosmosDBDatabase] = None,
+        settings: Optional[CosmosDBSettings] = None,
+    ) -> None:
+        super().__init__(
+            lambda settings: settings.get_storage_events_container_name(),
+            PARTITION_KEY,
+            model_type=model_type,
+            db=db,
+            settings=settings,
+            stored_procedures=STORED_PROCEDURES,
+            triggers=TRIGGERS,
+        )
+
+    def get_partition_key(self, model: T) -> str:
+        return model.id

pctasks/core/pctasks/core/cosmos/settings.py

@@ -14,7 +14,9 @@
 DEFAULT_DATABASE_NAME = "pctasks"
 DEFAULT_WORKFLOW_RUNS_CONTAINER = "workflow-runs"
 DEFAULT_WORKFLOWS_CONTAINER = "workflows"
+DEFAULT_ITEMS_CONTAINER = "items"
 DEFAULT_RECORDS_CONTAINER = "records"
+DEFAULT_STORAGE_EVENTS_CONTAINER_NAME = "storage-events"
 
 DEFAULT_SINGLE_PARTITION_KEY_VALUE = "partition_key"
 
@@ -37,7 +39,9 @@ def section_name(cls) -> str:
     database: str = DEFAULT_DATABASE_NAME
     workflows_container_name: str = DEFAULT_WORKFLOWS_CONTAINER
     workflow_runs_container_name: str = DEFAULT_WORKFLOW_RUNS_CONTAINER
+    items_container_name: str = DEFAULT_ITEMS_CONTAINER
     records_container_name: str = DEFAULT_RECORDS_CONTAINER
+    storage_events_container_name: str = DEFAULT_STORAGE_EVENTS_CONTAINER_NAME
 
     max_bulk_put_size: int = 250
 
@@ -54,6 +58,19 @@ def get_workflow_runs_container_name(self) -> str:
             )
         return self.workflow_runs_container_name
 
+    def get_storage_events_container_name(self) -> str:
+        if self.test_container_suffix:
+            return (
+                f"tmp-{self.storage_events_container_name}"
+                f"-{self.test_container_suffix}"
+            )
+        return self.storage_events_container_name
+
+    def get_items_container_name(self) -> str:
+        if self.test_container_suffix:
+            return f"tmp-{self.items_container_name}" f"-{self.test_container_suffix}"
+        return self.items_container_name
+
     def get_records_container_name(self) -> str:
         if self.test_container_suffix:
             return f"tmp-{self.records_container_name}-{self.test_container_suffix}"

pctasks/core/pctasks/core/models/event.py

@@ -5,6 +5,7 @@
 
 from pctasks.core.constants import MICROSOFT_OWNER
 from pctasks.core.models.base import ForeachConfig, PCBaseModel, RunRecordId
+from pctasks.core.models.record import Record
 from pctasks.core.models.utils import tzutc_now
 from pctasks.core.utils import StrEnum
 
@@ -20,6 +21,58 @@ class CloudEvent(PCBaseModel):
     data: Union[Dict[str, Any], BaseModel]
 
 
+class StorageEventType(StrEnum):
+    # Doesn't have an update
+    CREATED = "Microsoft.Storage.BlobCreated"
+    DELETED = "Microsoft.Storage.BlobDeleted"
+
+
+class StorageEventData(PCBaseModel):
+    """
+    A model for Blob Storage Events.
+
+    https://learn.microsoft.com/en-us/azure/event-grid/event-schema-blob-storage?tabs=cloud-event-schema
+    """
+
+    api: str
+    clientRequestId: str
+    requestId: str
+    eTag: str
+    contentType: str
+    contentLength: Optional[int]
+    contentOffset: Optional[int]
+    blobType: str
+    url: str
+    sequencer: str
+    storageDiagnostics: Dict[str, str]
+
+
+class StorageEvent(CloudEvent):
+    type: StorageEventType
+    data: StorageEventData
+
+
+class StorageEventRecord(StorageEvent, Record):
+    """
+    A Cosmos DB record for storage events.
+
+    Cloud Events are valid Cosmos DB records, so we just inherit the
+    additional methods from Record to satisfy mypy.
+    """
+
+    # We need to repeat this type definition to please mypy, which
+    # didn't like the "conflicting" definitions for 'type', despite
+    # them both being strings at the end of the day.
+    type: StorageEventType
+
+    def get_id(self) -> str:
+        return self.id
+
+    @staticmethod
+    def migrate(item: Dict[str, Any]) -> Dict[str, Any]:
+        return item
+
+
 class STACItemEventType(StrEnum):
     CREATED = "Microsoft.PlanetaryComputer.ItemCreated"
     UPDATED = "Microsoft.PlanetaryComputer.ItemUpdated"

pctasks/core/pctasks/core/models/item.py

@@ -0,0 +1,95 @@
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import pystac
+from pydantic import Field, validator
+
+from pctasks.core.models.record import Record
+from pctasks.core.utils import StrEnum
+
+
+class ItemRecordType(StrEnum):
+    STAC_ITEM = "StacItem"
+    ITEM_UPDATED = "ItemUpdated"
+
+
+class ItemRecord(Record):
+    type: ItemRecordType
+    stac_id: str
+
+    @validator("stac_id")
+    def _stac_id_validator(cls, v: str) -> str:
+        # Ensure a single forward slash in the string
+        if v.count("/") != 1:
+            raise ValueError("stac_id must contain a single forward slash")
+        return v
+
+    def get_id(self) -> str:
+        version = getattr(self, "version", "")
+        return f"{self.collection_id}:{self.item_id}:{version}:{self.type}"
+
+    @property
+    def collection_id(self) -> str:
+        return self.stac_id.split("/")[0]
+
+    @property
+    def item_id(self) -> str:
+        return self.stac_id.split("/")[1]
+
+
+class StacItemRecord(ItemRecord):
+    """
+    Record for STAC items.
+
+    These records are used in the items container of the Cosmos DB database.
+
+    Parameters
+    ----------
+    type: ItemRecordType, 'StacItem'
+        This is always 'StacItem'
+    stac_id: str
+        The "STAC identifier" which is the STAC collection ID and Item ID joined
+        by a single ``/``
+    version: str, optional
+        The STAC version identifier.
+    """
+
+    type: ItemRecordType = Field(default=ItemRecordType.STAC_ITEM, const=True)
+    item: Dict[str, Any]
+    deleted: bool = False
+
+    @classmethod
+    def from_item(cls, item: pystac.Item) -> "StacItemRecord":
+        collection_id = item.collection_id
+        item_id = item.id
+        stac_id = f"{collection_id}/{item_id}"
+        return cls(
+            stac_id=stac_id, version=item.properties.get("version"), item=item.to_dict()
+        )
+
+    @property
+    def version(self) -> str:
+        return self.item.get("properties", {}).get("version", "")
+
+
+class ItemUpdatedRecord(ItemRecord):
+    """Record that records an item update.
+
+    Does not specify if the item was created or updated.
+    """
+
+    type: ItemRecordType = Field(default=ItemRecordType.ITEM_UPDATED, const=True)
+
+    run_id: str
+    """The run ID of the workflow that updated this Item version"""
+
+    delete: bool = False
+    """True if the update was to delete this Item version"""
+
+    storage_event_time: Optional[datetime] = None
+    message_inserted_time: Optional[datetime] = None
+    version: Optional[str]
+
+    @validator("version")
+    def _version_validator(cls, v: Optional[str]) -> str:
+        return v or ""

pctasks/core/tests/models/test_item.py

@@ -0,0 +1,57 @@
+import datetime
+import typing
+
+import pystac
+import pytest
+
+import pctasks.core.cosmos.containers.items
+import pctasks.core.models.item
+
+
+@pytest.fixture
+def stac_item() -> typing.Generator[pystac.Item, None, None]:
+    """A simple STAC item for record tests."""
+    item = pystac.Item(
+        "my-item",
+        None,
+        None,
+        datetime.datetime(2000, 1, 1),
+        properties={},
+        collection="my-collection",
+    )
+    yield item
+
+
+def test_from_item(stac_item: pystac.Item):
+    result = pctasks.core.models.item.StacItemRecord.from_item(stac_item)
+    assert result.stac_id == "my-collection/my-item"
+    assert result.collection_id == "my-collection"
+    assert result.item_id == "my-item"
+    assert result.get_id() == "my-collection:my-item::StacItem"
+
+
+def test_id_validator(stac_item: pystac.Item):
+    stac_item.id = "my/item"
+    with pytest.raises(ValueError, match="stac_id must contain"):
+        pctasks.core.models.item.StacItemRecord.from_item(stac_item)
+
+
+def test_get_version(stac_item: pystac.Item):
+    stac_item.properties["version"] = "2"
+    result = pctasks.core.models.item.StacItemRecord.from_item(stac_item)
+    assert result.version == "2"
+
+
+@pytest.mark.parametrize(
+    "version, expected",
+    [
+        (None, "my-collection:my-item::ItemUpdated"),
+        ("2", "my-collection:my-item:2:ItemUpdated"),
+    ],
+)
+def test_item_update_record(version, expected):
+    record = pctasks.core.models.item.ItemUpdatedRecord(
+        stac_id="my-collection/my-item", run_id="test", version=version
+    )
+    result = record.get_id()
+    assert result == expected