[Persistence]: Backend implementation

backend_py/primary/primary/config.py

@@ -35,5 +35,5 @@
 REDIS_CACHE_URL = "redis://redis-cache:6379"
 
 COSMOS_DB_PROD_CONNECTION_STRING = os.environ.get("WEBVIZ_DB_CONNECTION_STRING", None)
-COSMOS_DB_EMULATOR_URI = "https://host.docker.internal:8081/"
-COSMOS_DB_EMULATOR_KEY = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==;"
+COSMOS_DB_EMULATOR_URI = "https://cosmos-db-emulator:8081"
+COSMOS_DB_EMULATOR_KEY = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="

backend_py/primary/primary/main.py

@@ -59,10 +59,7 @@
 logging.getLogger("primary.routers.grid3d").setLevel(logging.DEBUG)
 logging.getLogger("primary.routers.dev").setLevel(logging.DEBUG)
 logging.getLogger("primary.routers.surface").setLevel(logging.DEBUG)
-logging.getLogger("primary.persistence.cosmosdb").setLevel(logging.DEBUG)
-logging.getLogger("primary.persistence.session_store").setLevel(logging.DEBUG)
-logging.getLogger("primary.persistence.snapshot_store").setLevel(logging.DEBUG)
-logging.getLogger("primary.persistence.tasks").setLevel(logging.DEBUG)
+logging.getLogger("primary.persistence").setLevel(logging.DEBUG)
 # logging.getLogger("primary.auth").setLevel(logging.DEBUG)
 # logging.getLogger("uvicorn.error").setLevel(logging.DEBUG)
 # logging.getLogger("uvicorn.access").setLevel(logging.DEBUG)

backend_py/primary/primary/persistence/_utils.py

@@ -1,18 +1,38 @@
 import hashlib
-from typing import Any, cast
+from typing import Any, TypeVar, cast
+
+from azure.core.async_paging import AsyncPageIterator, AsyncItemPaged
+
+
+T = TypeVar("T")
 
 
 # Utility function to hash a JSON string using SHA-256
 # This function mimics the behavior of TextEncoder in JavaScript, which encodes strings to
 # UTF-8 before hashing. The output is a hexadecimal string representation of the hash.
 #
 # It is important that this function returns the same hash as the JavaScript version
-def hash_sha256(json_string: str) -> str:
-    data = json_string.encode("utf-8")  # Matches TextEncoder behavior
+def hash_session_content_string(string: str) -> str:
+    data = string.encode("utf-8")  # Matches TextEncoder behavior
     hash_bytes = hashlib.sha256(data).digest()
     hash_hex = "".join(f"{b:02x}" for b in hash_bytes)
     return hash_hex
 
 
 def cast_query_params(params: list[dict[str, Any]]) -> list[dict[str, object]]:
     return cast(list[dict[str, object]], params)
+
+
+def query_by_page(query_iterable: AsyncItemPaged[T], page_token: str | None) -> AsyncPageIterator[T]:
+    """
+    Cosmosdb's `by_page` returns a more narrow subtype than anticipated. This makes
+    extra's like `.continuation_token` not show up in returned value's type.
+
+    This util function correctly casts the return value to the expected type
+    """
+    pager = query_iterable.by_page(page_token)
+
+    if not isinstance(pager, AsyncPageIterator):
+        raise TypeError("Expected AsyncPageIterator from query_items_by_page_token_async")
+
+    return cast(AsyncPageIterator[T], pager)

backend_py/primary/primary/persistence/cosmosdb/cosmos_container.py

@@ -5,9 +5,12 @@
 from azure.cosmos import exceptions
 from pydantic import BaseModel, ValidationError
 
+from primary.persistence._utils import query_by_page
+
 from .cosmos_database import CosmosDatabase
 from .exceptions import (
     DatabaseAccessError,
+    DatabaseAccessIntegrityError,
     DatabaseAccessNotFoundError,
     DatabaseAccessConflictError,
     DatabaseAccessPreconditionFailedError,
@@ -21,16 +24,16 @@
 
 T = TypeVar("T", bound=BaseModel)
 
-"""
-CosmosContainer provides access to a specific container in a Cosmos DB database.
-It allows for querying, inserting, updating, and deleting items in the container.
-It uses a Pydantic model for item validation and serialization.
 
-It is designed to be used with asynchronous context management, ensuring proper resource cleanup.
-"""
+class CosmosContainer(Generic[T]):
+    """
+    CosmosContainer provides access to a specific container in a Cosmos DB database.
+    It allows for querying, inserting, updating, and deleting items in the container.
+    It uses a Pydantic model for item validation and serialization.
 
+    It is designed to be used with asynchronous context management, ensuring proper resource cleanup.
+    """
 
-class CosmosContainer(Generic[T]):
     def __init__(
         self,
         database_name: str,
@@ -46,19 +49,21 @@ def __init__(
         self._validation_model: Type[T] = validation_model
 
     @classmethod
-    def create(cls, database_name: str, container_name: str, validation_model: Type[T]) -> "CosmosContainer[T]":
+    def create_instance(
+        cls, database_name: str, container_name: str, validation_model: Type[T]
+    ) -> "CosmosContainer[T]":
         """Create a CosmosContainer instance."""
-        database = CosmosDatabase.create(database_name)
+        database = CosmosDatabase.create_instance(database_name)
         container = database.get_container(container_name)
         logger.debug("[CosmosContainer] Created for container '%s' in database '%s'", container_name, database_name)
         return cls(database_name, container_name, database, container, validation_model)
 
-    async def __aenter__(self) -> "CosmosContainer[T]":  # pylint: disable=C9001
+    async def __aenter__(self) -> "CosmosContainer[T]":
         return self
 
     async def __aexit__(
         self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
-    ) -> None:  # pylint: disable=C9001
+    ) -> None:
         await self.close_async()
 
     def _make_exception(self, op: str, exc: exceptions.CosmosHttpResponseError) -> DatabaseAccessError:
@@ -124,6 +129,24 @@ async def query_items_async(self, query: str, parameters: Optional[List[Dict[str
         except exceptions.CosmosHttpResponseError as error:
             raise self._make_exception("query_items_async", error)
 
+    async def query_items_by_page_token_async(
+        self,
+        query: str,
+        page_token: str | None,
+        parameters: Optional[List[Dict[str, object]]] = None,
+        page_size: Optional[int] = None,
+    ) -> tuple[list[T], str | None]:
+        query_iterable = self._container.query_items(query=query, parameters=parameters, max_item_count=page_size)
+
+        pager = query_by_page(query_iterable, page_token)
+        page = await anext(pager)
+
+        token = pager.continuation_token
+
+        items = [self._validation_model.model_validate(item) async for item in page]
+
+        return (items, token)
+
     async def get_item_async(self, item_id: str, partition_key: str) -> T:
         try:
             item = await self._container.read_item(item=item_id, partition_key=partition_key)
@@ -152,14 +175,14 @@ async def delete_item_async(self, item_id: str, partition_key: str) -> None:
         except exceptions.CosmosHttpResponseError as error:
             raise self._make_exception("delete_item_async", error) from error
 
-    async def update_item_async(self, item_id: str, partition_key: str, updated_item: T) -> None:
+    async def update_item_async(self, item_id: str, updated_item: T) -> None:
         try:
             validated = self._validation_model.model_validate(updated_item).model_dump(by_alias=True, mode="json")
 
             if validated.get("id") and validated["id"] != item_id:
-                raise ValueError(f"id mismatch: payload id {validated['id']} != path id {item_id}")
+                raise DatabaseAccessIntegrityError(f"id mismatch: payload id {validated['id']} != path id {item_id}")
 
-            await self._container.replace_item(item=item_id, body=validated, partition_key=partition_key)
+            await self._container.replace_item(item=item_id, body=validated)
 
             logger.debug("[CosmosContainer] Updated item '%s' in '%s'", item_id, self._container_name)
         except ValidationError as validation_error:

backend_py/primary/primary/persistence/cosmosdb/cosmos_database.py

@@ -8,13 +8,20 @@
 
 
 class CosmosDatabase:
+    """
+    CosmosDatabase provides access to a Cosmos DB database.
+    It allows for getting container proxies within the database.
+
+    It is designed to be used with asynchronous context management, ensuring proper resource cleanup.
+    """
+
     def __init__(self, database_name: str, client: CosmosClient):
         self._database_name = database_name
         self._client = client
         self._database = self._client.get_database_client(database_name)
 
     @classmethod
-    def create(cls, database_name: str) -> "CosmosDatabase":
+    def create_instance(cls, database_name: str) -> "CosmosDatabase":
         if COSMOS_DB_PROD_CONNECTION_STRING:
             client = CosmosClient.from_connection_string(COSMOS_DB_PROD_CONNECTION_STRING)
         elif COSMOS_DB_EMULATOR_URI and COSMOS_DB_EMULATOR_KEY:
@@ -26,12 +33,12 @@ def create(cls, database_name: str) -> "CosmosDatabase":
         self = cls(database_name, client)
         return self
 
-    async def __aenter__(self) -> "CosmosDatabase":  # pylint: disable=C9001
+    async def __aenter__(self) -> "CosmosDatabase":
         return self
 
     async def __aexit__(
         self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
-    ) -> None:  # pylint: disable=C9001
+    ) -> None:
         await self.close_async()
 
     def _make_exception(self, message: str) -> ServiceRequestError:

backend_py/primary/primary/persistence/cosmosdb/error_converter.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Dict, NoReturn, Optional, Type
+from typing import Dict, NoReturn, Type
 
 from primary.services.service_exceptions import Service, ServiceRequestError
 
@@ -27,15 +27,12 @@
 
 def convert_data_access_error_to_service_error(
     err: DatabaseAccessError,
-    *,
-    context: Optional[str] = None,
-    messages: Optional[Dict[Type[DatabaseAccessError], str]] = None,
 ) -> ServiceRequestError:
     """
     Convert a DatabaseAccess* error to a ServiceRequestError (without raising).
     You can customize messages per exception type via the 'messages' dict.
     """
-    msgs = {**_DEFAULT_MESSAGES, **(messages or {})}
+    msgs = {**_DEFAULT_MESSAGES}
 
     # Find the most specific message for the concrete type
     msg = None
@@ -55,22 +52,17 @@ def convert_data_access_error_to_service_error(
     if getattr(err, "activity_id", None):
         details.append(f"activity_id={err.activity_id}")
 
-    prefix = f"{context}: " if context else ""
     suffix = f" ({', '.join(details)})" if details else ""
-    message = f"{prefix}{msg}{suffix}"
+    message = f"{msg}{suffix}"
 
-    # Chain the original exception for traceback preservation
     return ServiceRequestError(message, Service.DATABASE)
 
 
 def raise_service_error_from_database_access(
     err: DatabaseAccessError,
-    *,
-    context: Optional[str] = None,
-    messages: Optional[Dict[Type[DatabaseAccessError], str]] = None,
 ) -> NoReturn:
     """
     Convert and raise immediately, chaining the original error.
     """
-    service_err = convert_data_access_error_to_service_error(err, context=context, messages=messages)
+    service_err = convert_data_access_error_to_service_error(err)
     raise service_err from err

backend_py/primary/primary/persistence/cosmosdb/exceptions.py

@@ -35,3 +35,7 @@ class DatabaseAccessThrottledError(DatabaseAccessError):
 
 class DatabaseAccessTransportError(DatabaseAccessError):
     """Other transport / HTTP errors."""
+
+
+class DatabaseAccessIntegrityError(DatabaseAccessError):
+    """Data integrity error."""