cosmodol 0.0.2__py3-none-any.whl

cosmodol/__init__.py

@@ -0,0 +1,100 @@
+"""Access Azure Cosmos DB (NoSQL/Core API) through a Mapping interface.
+
+``cosmodol`` exposes Azure Cosmos DB as ``dol``-style ``Mapping`` /
+``MutableMapping`` interfaces, layered over the official ``azure-cosmos`` SDK.
+
+Quick start::
+
+    from cosmodol import cosmos_store
+
+    store = cosmos_store(
+        connection_string="AccountEndpoint=https://localhost:8081/;AccountKey=...",
+        database="mydb",
+        container="mycontainer",
+        partition_key_value="tenant-X",
+    )
+
+    store["k1"] = {"name": "Alice", "age": 30}
+    store["k1"]              # → {"id": "k1", "<pk>": "tenant-X", "name": "Alice", "age": 30}
+    "k1" in store            # → True
+    del store["k1"]
+
+See ``misc/docs/architecture.md`` for the layered design.
+
+If your Cosmos account was provisioned with the **MongoDB API**, use ``pymongo`` +
+``mongodol`` directly — this package only targets the NoSQL/Core API.
+"""
+
+from cosmodol.connection import CosmosConnection, resolve_credential
+from cosmodol.base import (
+    ResponseHeaders,
+    batch,
+    point_contains,
+    point_delete,
+    point_get,
+    point_replace,
+    point_upsert,
+    query,
+)
+from cosmodol.errors import (
+    ContainerNotEmptyError,
+    ContainerNotFoundError,
+    CosmosThrottleError,
+    DatabaseNotEmptyError,
+    DatabaseNotFoundError,
+    ItemAlreadyExistsError,
+    ItemNotFoundError,
+    KeyMismatchError,
+    translate_cosmos_errors,
+    validate_cosmos_id,
+)
+from cosmodol.stores import (
+    CosmosItems,
+    CosmosPartitionedItems,
+    SYSTEM_FIELDS,
+)
+from cosmodol.trees import (
+    CosmosAccount,
+    CosmosDatabase,
+)
+from cosmodol.recipes import (
+    cosmos_store,
+    strip_system_fields,
+)
+
+
+__all__ = [
+    # connection
+    "CosmosConnection",
+    "resolve_credential",
+    # base / metal layer
+    "ResponseHeaders",
+    "batch",
+    "point_contains",
+    "point_delete",
+    "point_get",
+    "point_replace",
+    "point_upsert",
+    "query",
+    # errors
+    "ContainerNotEmptyError",
+    "ContainerNotFoundError",
+    "CosmosThrottleError",
+    "DatabaseNotEmptyError",
+    "DatabaseNotFoundError",
+    "ItemAlreadyExistsError",
+    "ItemNotFoundError",
+    "KeyMismatchError",
+    "translate_cosmos_errors",
+    "validate_cosmos_id",
+    # stores
+    "CosmosItems",
+    "CosmosPartitionedItems",
+    "SYSTEM_FIELDS",
+    # trees
+    "CosmosAccount",
+    "CosmosDatabase",
+    # recipes
+    "cosmos_store",
+    "strip_system_fields",
+]

cosmodol/base.py

@@ -0,0 +1,175 @@
+"""Close-to-metal CRUD facade for Cosmos DB.
+
+Functional layer over ``ContainerProxy``. Each function takes the container plus
+normalized args and returns ``(value, ResponseHeaders)`` so callers can observe RU
+charges and ETags. The Mapping-shaped stores in ``cosmodol.stores`` route all their
+work through these.
+
+See ``misc/docs/architecture.md`` Layer A.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Iterator, NamedTuple, Optional
+
+from azure.cosmos import ContainerProxy
+
+from cosmodol.errors import translate_cosmos_errors
+
+
+class ResponseHeaders(NamedTuple):
+    """Subset of Cosmos response headers we surface for observability.
+
+    Attributes:
+        request_charge: RU consumed by the operation. ``None`` when the emulator (which
+            does not populate this header) was the backend.
+        etag: ETag of the (created / read / replaced) item, when applicable.
+    """
+
+    request_charge: Optional[float]
+    etag: Optional[str]
+
+
+def _headers(container: ContainerProxy) -> ResponseHeaders:
+    """Extract our observability tuple from the last response on the container proxy."""
+    h = container.client_connection.last_response_headers or {}
+    rc = h.get("x-ms-request-charge")
+    return ResponseHeaders(
+        request_charge=float(rc) if rc is not None else None,
+        etag=h.get("etag"),
+    )
+
+
+@translate_cosmos_errors(key_arg="id")
+def point_get(
+    container: ContainerProxy,
+    id: str,
+    partition_key: Any,
+) -> tuple[dict, ResponseHeaders]:
+    """Point read of one item. ~1 RU/KB. Raises ``ItemNotFoundError`` on missing."""
+    item = container.read_item(item=id, partition_key=partition_key)
+    return dict(item), _headers(container)
+
+
+@translate_cosmos_errors(key_arg="id")
+def point_upsert(
+    container: ContainerProxy,
+    body: dict,
+    *,
+    etag: Optional[str] = None,
+) -> tuple[dict, ResponseHeaders]:
+    """Insert-or-replace an item. ``id`` and the partition-key value are extracted from
+    ``body``.
+
+    Note: ``etag`` parameter accepted for symmetry but Cosmos does not honor it on
+    ``upsert_item``; use ``point_replace`` for ETag-conditional writes.
+    """
+    _ = etag  # accepted for symmetry; see docstring
+    item = container.upsert_item(body=body)
+    return dict(item), _headers(container)
+
+
+@translate_cosmos_errors(key_arg="id")
+def point_replace(
+    container: ContainerProxy,
+    id: str,
+    body: dict,
+    partition_key: Any,
+    *,
+    etag: Optional[str] = None,
+) -> tuple[dict, ResponseHeaders]:
+    """Full replace of an item. With ``etag``, performs an If-Match conditional write."""
+    kwargs = {}
+    if etag is not None:
+        kwargs["if_match_etag"] = etag
+    item = container.replace_item(item=id, body=body, **kwargs)
+    return dict(item), _headers(container)
+
+
+@translate_cosmos_errors(key_arg="id")
+def point_delete(
+    container: ContainerProxy,
+    id: str,
+    partition_key: Any,
+    *,
+    etag: Optional[str] = None,
+) -> ResponseHeaders:
+    """Point delete. Raises ``ItemNotFoundError`` on missing."""
+    kwargs = {}
+    if etag is not None:
+        kwargs["if_match_etag"] = etag
+    container.delete_item(item=id, partition_key=partition_key, **kwargs)
+    return _headers(container)
+
+
+def point_contains(
+    container: ContainerProxy,
+    id: str,
+    partition_key: Any,
+) -> tuple[bool, ResponseHeaders]:
+    """Existence check via point read. ``True/False``. Never raises ``KeyError``.
+
+    Cheaper than a ``SELECT VALUE COUNT(1)`` query — ~1 RU vs ≥ 2.3 RU.
+    """
+    from azure.cosmos.exceptions import CosmosResourceNotFoundError
+
+    try:
+        container.read_item(item=id, partition_key=partition_key)
+        return True, _headers(container)
+    except CosmosResourceNotFoundError:
+        return False, _headers(container)
+
+
+def query(
+    container: ContainerProxy,
+    sql: str,
+    *,
+    parameters: Optional[list[dict]] = None,
+    partition_key: Any = None,
+    cross_partition: bool = False,
+    max_item_count: Optional[int] = None,
+) -> Iterator[dict]:
+    """Run a SQL query. Yields dicts.
+
+    Either pass ``partition_key=`` (cheap, single-partition) or
+    ``cross_partition=True`` (RU scales with data). One of the two is required by Cosmos.
+    """
+    kwargs: dict = {"query": sql}
+    if parameters is not None:
+        kwargs["parameters"] = parameters
+    if partition_key is not None:
+        kwargs["partition_key"] = partition_key
+    elif cross_partition:
+        kwargs["enable_cross_partition_query"] = True
+    if max_item_count is not None:
+        kwargs["max_item_count"] = max_item_count
+    for item in container.query_items(**kwargs):
+        # ``SELECT VALUE ...`` queries yield raw scalars (str, int, float, bool, list).
+        # Other queries yield CosmosDict (a dict subclass). Pass through either shape
+        # so callers can distinguish; convert to plain dict in the dict-shaped case.
+        if isinstance(item, dict):
+            yield dict(item)
+        else:
+            yield item
+
+
+def batch(
+    container: ContainerProxy,
+    operations: list[tuple],
+    partition_key: Any,
+) -> list[dict]:
+    """Transactional batch within one logical partition.
+
+    Args:
+        operations: List of ``(op_name, args_tuple, kwargs_dict)`` triples. Cosmos op
+            names: ``"create"``, ``"upsert"``, ``"replace"``, ``"patch"``, ``"read"``,
+            ``"delete"``. ≤ 100 ops, ≤ 1.2 MB total.
+        partition_key: All ops must share this partition-key value.
+
+    Returns:
+        List of per-op result dicts as returned by the SDK.
+    """
+    results = container.execute_item_batch(
+        batch_operations=operations, partition_key=partition_key
+    )
+    return [dict(r) for r in results]

cosmodol/connection.py

@@ -0,0 +1,211 @@
+"""Connection-layer for cosmodol.
+
+Owns the expensive resource (``CosmosClient``) and the credential cascade.
+See ``misc/docs/architecture.md`` Layer 0 and ``misc/docs/design_decisions.md`` §14.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import Any, Optional, Union
+
+from azure.cosmos import CosmosClient
+
+
+CredentialLike = Union[str, dict, Any, None]
+
+
+# Env vars consulted by ``resolve_credential``.
+_ENV_CONNECTION_STRING = "AZURE_COSMOS_CONNECTION_STRING"
+_ENV_ENDPOINT = "AZURE_COSMOS_ENDPOINT"
+_ENV_KEY = "AZURE_COSMOS_KEY"
+
+
+_CONN_STR_ENDPOINT_RE = re.compile(r"AccountEndpoint=([^;]+)", re.IGNORECASE)
+_CONN_STR_KEY_RE = re.compile(r"AccountKey=([^;]+)", re.IGNORECASE)
+
+
+def _looks_like_connection_string(s: str) -> bool:
+    return "AccountEndpoint=" in s and "AccountKey=" in s
+
+
+def _parse_connection_string(cs: str) -> tuple[str, str]:
+    """Parse a Cosmos connection string into ``(endpoint, key)``."""
+    ep = _CONN_STR_ENDPOINT_RE.search(cs)
+    k = _CONN_STR_KEY_RE.search(cs)
+    if not (ep and k):
+        raise ValueError(
+            "Cosmos connection string must contain AccountEndpoint=... and AccountKey=..."
+        )
+    return ep.group(1), k.group(1)
+
+
+def resolve_credential(
+    *,
+    credential: CredentialLike = None,
+    connection_string: Optional[str] = None,
+    endpoint: Optional[str] = None,
+    key: Optional[str] = None,
+) -> dict:
+    """Resolve a credential into a normalized form that can build a ``CosmosClient``.
+
+    Cascade (first hit wins):
+
+    1. Explicit ``credential=`` (with ``endpoint=`` for the URL)
+    2. Explicit ``connection_string=``
+    3. Explicit ``endpoint=`` + ``key=`` (or just ``endpoint=`` + AAD)
+    4. Env var ``AZURE_COSMOS_CONNECTION_STRING``
+    5. Env vars ``AZURE_COSMOS_ENDPOINT`` + ``AZURE_COSMOS_KEY``
+    6. Env var ``AZURE_COSMOS_ENDPOINT`` alone + ``DefaultAzureCredential``
+
+    Returns:
+        A dict with: ``{"url": "...", "credential": <obj>}``
+
+    Raises:
+        ValueError: if no source resolves.
+    """
+    # 1. Explicit credential
+    if credential is not None:
+        if isinstance(credential, str) and _looks_like_connection_string(credential):
+            ep, k = _parse_connection_string(credential)
+            return {"url": ep, "credential": k}
+        ep = endpoint or os.environ.get(_ENV_ENDPOINT)
+        if ep is None:
+            raise ValueError(
+                "credential=... was provided but no endpoint. Pass endpoint=... or set "
+                f"the {_ENV_ENDPOINT} env var."
+            )
+        return {"url": ep, "credential": credential}
+
+    # 2. Explicit connection string
+    if connection_string is not None:
+        ep, k = _parse_connection_string(connection_string)
+        return {"url": ep, "credential": k}
+
+    # 3. Explicit endpoint + key (or endpoint alone)
+    if endpoint is not None and key is not None:
+        return {"url": endpoint, "credential": key}
+    if endpoint is not None:
+        return {"url": endpoint, "credential": _default_aad_credential()}
+
+    # 4. Env: connection string
+    env_cs = os.environ.get(_ENV_CONNECTION_STRING)
+    if env_cs:
+        ep, k = _parse_connection_string(env_cs)
+        return {"url": ep, "credential": k}
+
+    # 5. Env: endpoint + key
+    env_ep = os.environ.get(_ENV_ENDPOINT)
+    env_k = os.environ.get(_ENV_KEY)
+    if env_ep and env_k:
+        return {"url": env_ep, "credential": env_k}
+
+    # 6. Env: endpoint alone + AAD
+    if env_ep:
+        return {"url": env_ep, "credential": _default_aad_credential()}
+
+    raise ValueError(
+        "Could not resolve Cosmos DB credentials. Provide one of:\n"
+        "  - credential=<obj or connection string>\n"
+        "  - connection_string=<str>\n"
+        "  - endpoint=<url> + key=<key>\n"
+        "  - env var AZURE_COSMOS_CONNECTION_STRING\n"
+        "  - env vars AZURE_COSMOS_ENDPOINT + AZURE_COSMOS_KEY\n"
+        "  - env var AZURE_COSMOS_ENDPOINT alone (uses DefaultAzureCredential)\n"
+    )
+
+
+def _default_aad_credential():
+    """Lazy import to keep ``azure-identity`` an optional dependency."""
+    try:
+        from azure.identity import DefaultAzureCredential
+    except ImportError as e:
+        raise ImportError(
+            "azure-identity is required for AAD credentials. Install with: "
+            "`pip install azure-identity`"
+        ) from e
+    return DefaultAzureCredential()
+
+
+@dataclass
+class CosmosConnection:
+    """Holds a resolved credential and a lazy ``CosmosClient``.
+
+    This is the dependency-injection seam for the package. Tests construct one pointing
+    at the emulator without touching any store class.
+
+    Args:
+        credential: explicit credential object or account-key string.
+        connection_string: full ``AccountEndpoint=...;AccountKey=...`` string.
+        endpoint: Cosmos endpoint URL (with or without credential).
+        key: account master key.
+        consistency_level: optional override; defaults to inheriting the account default.
+        client_kwargs: extra kwargs forwarded to ``CosmosClient``.
+    """
+
+    credential: CredentialLike = None
+    connection_string: Optional[str] = None
+    endpoint: Optional[str] = None
+    key: Optional[str] = None
+    consistency_level: Optional[str] = None
+    client_kwargs: dict = field(default_factory=dict)
+
+    @cached_property
+    def _resolved(self) -> dict:
+        return resolve_credential(
+            credential=self.credential,
+            connection_string=self.connection_string,
+            endpoint=self.endpoint,
+            key=self.key,
+        )
+
+    @cached_property
+    def client(self) -> CosmosClient:
+        """The lazily-constructed ``CosmosClient``. Cached for the connection's lifetime."""
+        r = self._resolved
+        kw = dict(self.client_kwargs)
+        if self.consistency_level is not None:
+            kw.setdefault("consistency_level", self.consistency_level)
+        return CosmosClient(url=r["url"], credential=r["credential"], **kw)
+
+    def database(self, name: str):
+        return self.client.get_database_client(name)
+
+    def container(self, database: str, container: str):
+        return self.client.get_database_client(database).get_container_client(container)
+
+    @classmethod
+    def from_anything(cls, source) -> "CosmosConnection":
+        """Convenience: build a ``CosmosConnection`` from a thing-or-spec.
+
+        Accepts:
+            - ``CosmosConnection`` (returned as-is)
+            - ``CosmosClient`` (wrapped without further resolution)
+            - ``str`` (connection string)
+            - ``dict`` (passed as kwargs)
+            - ``None`` (defer to env / AAD)
+        """
+        if isinstance(source, cls):
+            return source
+        if isinstance(source, CosmosClient):
+            inst = cls.__new__(cls)
+            inst.credential = None
+            inst.connection_string = None
+            inst.endpoint = None
+            inst.key = None
+            inst.consistency_level = None
+            inst.client_kwargs = {}
+            inst.__dict__["client"] = source
+            return inst
+        if isinstance(source, str):
+            return cls(connection_string=source)
+        if isinstance(source, dict):
+            return cls(**source)
+        if source is None:
+            return cls()
+        raise TypeError(
+            f"Cannot build CosmosConnection from {type(source).__name__}: {source!r}"
+        )

cosmodol/errors.py

@@ -0,0 +1,143 @@
+"""Custom exceptions and error-translation decorator for cosmodol.
+
+All Azure Cosmos SDK exception handling for the Mapping-shaped methods on close-to-metal
+stores funnels through ``translate_cosmos_errors`` so the auth/throttle vs not-found
+distinction stays auditable from one place. See ``misc/docs/design_decisions.md`` §11.
+"""
+
+from __future__ import annotations
+
+import inspect
+from functools import wraps
+from typing import Callable
+
+from azure.cosmos.exceptions import (
+    CosmosResourceExistsError,
+    CosmosResourceNotFoundError,
+)
+
+
+def _extract_key(func, args, kwargs, key_arg):
+    """Resolve the user-facing key from *args/**kwargs.
+
+    Supports either a positional int index OR a name (resolved via the function's
+    signature so it works whether the caller passed by position or by keyword).
+    """
+    if isinstance(key_arg, int):
+        return args[key_arg] if key_arg < len(args) else None
+    # By name: prefer kwargs; fall back to positional lookup via signature.
+    if key_arg in kwargs:
+        return kwargs[key_arg]
+    try:
+        sig = inspect.signature(func)
+        params = list(sig.parameters)
+        idx = params.index(key_arg)
+        return args[idx] if idx < len(args) else None
+    except (ValueError, TypeError):
+        return None
+
+
+class ItemNotFoundError(KeyError):
+    """Raised when an item does not exist for a (partition_key, id)."""
+
+
+class ItemAlreadyExistsError(KeyError):
+    """Raised on strict-create attempts when the item already exists."""
+
+
+class ContainerNotFoundError(KeyError):
+    """Raised when a Cosmos container does not exist."""
+
+
+class DatabaseNotFoundError(KeyError):
+    """Raised when a Cosmos database does not exist."""
+
+
+class ContainerNotEmptyError(RuntimeError):
+    """Raised on ``del db_store[name]`` when the container has items. See
+    ``misc/docs/design_decisions.md`` §8.
+    """
+
+
+class DatabaseNotEmptyError(RuntimeError):
+    """Raised on ``del account_store[name]`` when the database has containers. See
+    ``misc/docs/design_decisions.md`` §8.
+    """
+
+
+class KeyMismatchError(ValueError):
+    """Raised when a written body's ``id`` (or partition-key value) disagrees with the
+    inferred value (from the dict key). See ``misc/docs/design_decisions.md`` §9.
+    """
+
+
+class CosmosThrottleError(RuntimeError):
+    """Wraps ``CosmosHttpResponseError`` with HTTP 429 (RU exhaustion)."""
+
+
+# Cosmos id-string validation.
+_ID_FORBIDDEN = set("/\\?#")
+_ID_MAX_LEN = 255
+
+
+def validate_cosmos_id(k) -> None:
+    """Raise ``ValueError`` if ``k`` is not a valid Cosmos item ``id``.
+
+    Cosmos may accept invalid ids on write but the item then becomes unreachable from
+    the SDK. We fail loudly at write time. See ``misc/docs/cosmos_db_reference.md``
+    §"id rules".
+    """
+    if not isinstance(k, str):
+        raise ValueError(f"Cosmos item id must be a str, got {type(k).__name__}: {k!r}")
+    if not k:
+        raise ValueError("Cosmos item id must be a non-empty string.")
+    if len(k) > _ID_MAX_LEN:
+        raise ValueError(f"Cosmos item id length {len(k)} exceeds {_ID_MAX_LEN}: {k!r}")
+    bad = sorted(set(k) & _ID_FORBIDDEN)
+    if bad:
+        raise ValueError(
+            f"Cosmos item id contains forbidden characters {bad}: {k!r}. "
+            "Use URL-safe base64 if you need to encode arbitrary strings."
+        )
+
+
+def translate_cosmos_errors(
+    *,
+    key_arg: int | str = 0,
+    not_found_cls: type[KeyError] = ItemNotFoundError,
+    exists_cls: type[KeyError] = ItemAlreadyExistsError,
+) -> Callable:
+    """Decorator: translate Cosmos SDK exceptions into ``KeyError`` subclasses.
+
+    Auth errors and any other Cosmos errors propagate untouched. HTTP 429 throttling is
+    wrapped in ``CosmosThrottleError`` so callers can distinguish it. See
+    ``misc/docs/design_decisions.md`` §11.
+
+    Args:
+        key_arg: Position (int) or name (str) of the key argument in the wrapped
+            method's signature. Used to populate ``KeyError(key)``. Default 0; for
+            methods the user-facing key is typically at index 1 (``self`` at 0).
+        not_found_cls: Exception class to raise on ``CosmosResourceNotFoundError``.
+        exists_cls: Exception class to raise on ``CosmosResourceExistsError``.
+    """
+    from azure.cosmos.exceptions import CosmosHttpResponseError
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except CosmosResourceNotFoundError as e:
+                k = _extract_key(func, args, kwargs, key_arg)
+                raise not_found_cls(k) from e
+            except CosmosResourceExistsError as e:
+                k = _extract_key(func, args, kwargs, key_arg)
+                raise exists_cls(k) from e
+            except CosmosHttpResponseError as e:
+                if getattr(e, "status_code", None) == 429:
+                    raise CosmosThrottleError(str(e)) from e
+                raise
+
+        return wrapper
+
+    return decorator