feat: Add Azure Cosmos DB store implementation #297
Conversation
Add CosmosDBStore for Azure Cosmos DB NoSQL API support: - Implements AsyncKeyValue protocol with full CRUD operations - Uses partition key for collection isolation - Supports TTL via Cosmos DB's native TTL feature - Includes collection sanitization strategy - Supports both client injection and URL/credential configuration - Adds optional dependency: azure-cosmos>=4.7.0 Closes #292 Co-authored-by: Bill Easton <[email protected]>
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the
✨ Finishing touches🧪 Generate unit tests (beta)
Important Action Needed: IP Allowlist UpdateIf your organization protects your Git platform with IP whitelisting, please add the new CodeRabbit IP address to your allowlist:
Failure to add the new IP will result in interrupted reviews. Comment |
Add
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
pyproject.toml (1)
54-65:⚠️ Potential issue | 🟡 MinorUpdate azure-cosmos dependency to explicitly support Python 3.13 or verify compatibility.
azure-cosmosversion 4.7.0 has explicit classifiers for Python 3.10–3.12 but lacks Python 3.13 support. If the project targets Python 3.13, either upgradeazure-cosmosto a version that officially supports 3.13 or document that 3.13 compatibility is unverified.
🤖 Fix all issues with AI agents
In `@src/key_value/aio/stores/cosmosdb/store.py`:
- Around line 198-291: The constructor (__init__) currently allows conflicting
inputs and mixes truthiness checks; change validation to explicitly enforce
mutual exclusivity between client and url/credential by using "is not None"
checks: raise ValueError if both client is not None and (url is not None or
credential is not None), and also raise ValueError if client is None and (url is
None or credential is None); compute client_provided = (client is not None) and
only call _create_cosmosdb_client when client is None, keeping references to
_client, _create_cosmosdb_client, and client_provided consistent.
- Around line 408-465: The _put_managed_entry method fails when
ManagedEntry.expires_at is a tz-naive datetime; before doing TTL math, normalize
managed_entry.expires_at to an aware UTC datetime (e.g., if expires_at has no
tzinfo, attach UTC) so subtraction against now = datetime.now(tz=timezone.utc)
won't raise; update logic around "if managed_entry.expires_at is not None" to
convert the value to an aware UTC datetime, then store the ISO string and
compute ttl_seconds, keeping use of _make_item_id, _sanitize_collection,
_serialization_adapter.dump_json and calling _upsert_item unchanged.
In `@tests/stores/cosmosdb/test_cosmosdb.py`:
- Around line 32-46: The cleanup query in clean_cosmosdb_container performs a
cross-partition scan but does not enable cross-partition queries; update the
call to store._container.query_items (in function clean_cosmosdb_container) to
pass enable_cross_partition_query=True so the SDK will allow fan-out queries
across partitions and avoid runtime errors when iterating items.
🧹 Nitpick comments (1)
tests/stores/cosmosdb/test_cosmosdb.py (1)
82-125: Close stores in fixtures to avoid leaked client sessions.The fixtures currently return the store without guaranteed teardown. A yield fixture with
close()infinallykeeps resources tidy.♻️ Suggested teardown pattern
`@override` `@pytest.fixture` async def store(self) -> CosmosDBStore: @@ - await clean_cosmosdb_container(store=store) - - return store + await clean_cosmosdb_container(store=store) + try: + yield store + finally: + await store.close() @@ `@pytest.fixture` async def sanitizing_store(self) -> CosmosDBStore: @@ - await clean_cosmosdb_container(store=store) - - return store + await clean_cosmosdb_container(store=store) + try: + yield store + finally: + await store.close()
| @overload | ||
| def __init__( | ||
| self, | ||
| *, | ||
| client: CosmosClient, | ||
| database_name: str | None = None, | ||
| container_name: str | None = None, | ||
| default_collection: str | None = None, | ||
| collection_sanitization_strategy: SanitizationStrategy | None = None, | ||
| auto_create: bool = True, | ||
| ) -> None: | ||
| """Initialize the Cosmos DB store. | ||
|
|
||
| Args: | ||
| client: The CosmosClient to use. | ||
| database_name: The name of the Cosmos DB database. | ||
| container_name: The name of the Cosmos DB container. | ||
| default_collection: The default collection to use if no collection is provided. | ||
| collection_sanitization_strategy: The sanitization strategy to use for collections. | ||
| auto_create: Whether to automatically create database/container if they don't exist. Defaults to True. | ||
| """ | ||
|
|
||
| @overload | ||
| def __init__( | ||
| self, | ||
| *, | ||
| url: str, | ||
| credential: str, | ||
| database_name: str | None = None, | ||
| container_name: str | None = None, | ||
| default_collection: str | None = None, | ||
| collection_sanitization_strategy: SanitizationStrategy | None = None, | ||
| auto_create: bool = True, | ||
| ) -> None: | ||
| """Initialize the Cosmos DB store. | ||
|
|
||
| Args: | ||
| url: The Cosmos DB account URL. | ||
| credential: The Cosmos DB account key or credential. | ||
| database_name: The name of the Cosmos DB database. | ||
| container_name: The name of the Cosmos DB container. | ||
| default_collection: The default collection to use if no collection is provided. | ||
| collection_sanitization_strategy: The sanitization strategy to use for collections. | ||
| auto_create: Whether to automatically create database/container if they don't exist. Defaults to True. | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| *, | ||
| client: CosmosClient | None = None, | ||
| url: str | None = None, | ||
| credential: str | None = None, | ||
| database_name: str | None = None, | ||
| container_name: str | None = None, | ||
| default_collection: str | None = None, | ||
| collection_sanitization_strategy: SanitizationStrategy | None = None, | ||
| auto_create: bool = True, | ||
| ) -> None: | ||
| """Initialize the Cosmos DB store. | ||
|
|
||
| Args: | ||
| client: The CosmosClient to use (mutually exclusive with url/credential). If provided, the store | ||
| will not manage the client's lifecycle (will not close it). The caller is responsible for | ||
| managing the client's lifecycle. | ||
| url: The Cosmos DB account URL (mutually exclusive with client). | ||
| credential: The Cosmos DB account key or credential (mutually exclusive with client). | ||
| database_name: The name of the Cosmos DB database. | ||
| container_name: The name of the Cosmos DB container. | ||
| default_collection: The default collection to use if no collection is provided. | ||
| collection_sanitization_strategy: The sanitization strategy to use for collections. | ||
| auto_create: Whether to automatically create database/container if they don't exist. Defaults to True. | ||
| When False, raises ValueError if the database or container doesn't exist. | ||
| """ | ||
| client_provided = client is not None | ||
|
|
||
| if client: | ||
| self._client = client | ||
| else: | ||
| if not url or not credential: | ||
| msg = "Either 'client' or both 'url' and 'credential' must be provided" | ||
| raise ValueError(msg) | ||
| self._client = _create_cosmosdb_client(url=url, credential=credential) | ||
|
|
||
| self._database_name = database_name or DEFAULT_DATABASE | ||
| self._container_name = container_name or DEFAULT_CONTAINER | ||
| self._database = None | ||
| self._container = None | ||
| self._auto_create = auto_create | ||
|
|
||
| super().__init__( | ||
| default_collection=default_collection, | ||
| collection_sanitization_strategy=collection_sanitization_strategy, | ||
| client_provided_by_user=client_provided, | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's read the actual file to see the current implementation
cat -n src/key_value/aio/stores/cosmosdb/store.py | sed -n '198,291p'Repository: strawgate/py-key-value
Length of output: 4850
🏁 Script executed:
# Check other store implementations for constructor patterns
fd -e py -path "*/stores/*" -name "store.py" | head -20Repository: strawgate/py-key-value
Length of output: 234
🏁 Script executed:
# Search for how client_provided_by_user is used
rg -n "client_provided_by_user" --type py -A 3 -B 3Repository: strawgate/py-key-value
Length of output: 20817
🏁 Script executed:
# Check for truthiness checks on client in the cosmosdb store
rg -n "if client" --type py src/key_value/aio/stores/cosmosdb/Repository: strawgate/py-key-value
Length of output: 131
🏁 Script executed:
# Check imports and CosmosClient type
head -50 src/key_value/aio/stores/cosmosdb/store.py | grep -E "^import|^from|CosmosClient"Repository: strawgate/py-key-value
Length of output: 948
🏁 Script executed:
# Search for CosmosClient definition or any __bool__ method
rg -n "CosmosClient" src/key_value/aio/stores/cosmosdb/store.py | head -20Repository: strawgate/py-key-value
Length of output: 784
🏁 Script executed:
# Check if other stores with similar patterns (firestore, disk) handle mutual exclusivity
rg -n -A 10 "def __init__" src/key_value/aio/stores/firestore/store.py | head -40Repository: strawgate/py-key-value
Length of output: 1177
🏁 Script executed:
# Check disk.py for comparison
rg -n -A 10 "def __init__" src/key_value/aio/stores/disk/store.py | head -50Repository: strawgate/py-key-value
Length of output: 1563
🏁 Script executed:
# Check firestore validation logic
sed -n '58,110p' src/key_value/aio/stores/firestore/store.pyRepository: strawgate/py-key-value
Length of output: 2350
🏁 Script executed:
# Check disk validation logic
sed -n '137,190p' src/key_value/aio/stores/disk/store.pyRepository: strawgate/py-key-value
Length of output: 2248
🏁 Script executed:
# Check if CosmosClient can be falsy - look at azure SDK documentation or imports
rg -n "if client" src/key_value/aio/stores/cosmosdb/store.py -A 5 -B 2Repository: strawgate/py-key-value
Length of output: 378
🏁 Script executed:
# Let me also check if there are any tests that might reveal whether both params are ever provided
rg -n "CosmosDBStore" --type py -A 5 | grep -E "client=|url=|credential=" | head -30Repository: strawgate/py-key-value
Length of output: 333
Enforce mutual exclusivity and add explicit validation for conflicting parameters.
The docstring documents mutual exclusivity between client and url/credential, but the implementation silently ignores conflicting parameters. Additionally, line 271 uses explicit is not None checks while line 273 uses a truthiness check, creating an inconsistent pattern.
Compare with disk.py, which properly validates mutual exclusivity:
if disk_cache is not None and directory is not None:
msg = "Provide only one of disk_cache or directory"
raise ValueError(msg)Adopt this approach for consistency and to catch configuration errors:
Suggested validation tightening
- client_provided = client is not None
+ if client is not None and (url or credential):
+ msg = "Provide either 'client' or 'url'/'credential', not both"
+ raise ValueError(msg)
- if client:
+ client_provided = client is not None
+
+ if client is not None:
self._client = client
else:
if not url or not credential:
msg = "Either 'client' or both 'url' and 'credential' must be provided"
raise ValueError(msg)
self._client = _create_cosmosdb_client(url=url, credential=credential)🤖 Prompt for AI Agents
In `@src/key_value/aio/stores/cosmosdb/store.py` around lines 198 - 291, The
constructor (__init__) currently allows conflicting inputs and mixes truthiness
checks; change validation to explicitly enforce mutual exclusivity between
client and url/credential by using "is not None" checks: raise ValueError if
both client is not None and (url is not None or credential is not None), and
also raise ValueError if client is None and (url is None or credential is None);
compute client_provided = (client is not None) and only call
_create_cosmosdb_client when client is None, keeping references to _client,
_create_cosmosdb_client, and client_provided consistent.
| @override | ||
| async def _put_managed_entry( | ||
| self, | ||
| *, | ||
| key: str, | ||
| collection: str, | ||
| managed_entry: ManagedEntry, | ||
| ) -> None: | ||
| if self._container is None: | ||
| msg = "Container not initialized" | ||
| raise ValueError(msg) | ||
|
|
||
| sanitized_collection = self._sanitize_collection(collection=collection) | ||
| item_id = self._make_item_id(key=key, collection=sanitized_collection) | ||
| json_value = self._serialization_adapter.dump_json(entry=managed_entry, key=key, collection=collection) | ||
|
|
||
| item: dict[str, Any] = { | ||
| "id": item_id, | ||
| "collection": sanitized_collection, | ||
| "key": key, | ||
| "value": json_value, | ||
| } | ||
|
|
||
| # Add TTL if present | ||
| if managed_entry.expires_at is not None: | ||
| # Store expires_at as ISO string for our own tracking | ||
| item["expires_at"] = managed_entry.expires_at.isoformat() | ||
|
|
||
| # Calculate TTL in seconds from now for Cosmos DB's TTL feature | ||
| now = datetime.now(tz=timezone.utc) | ||
| ttl_seconds = int((managed_entry.expires_at - now).total_seconds()) | ||
| if ttl_seconds > 0: | ||
| item["ttl"] = ttl_seconds | ||
|
|
||
| await _upsert_item(self._container, item) | ||
|
|
||
| @override | ||
| async def _put_managed_entries( | ||
| self, | ||
| *, | ||
| collection: str, | ||
| keys: Sequence[str], | ||
| managed_entries: Sequence[ManagedEntry], | ||
| ttl: float | None, | ||
| created_at: datetime, | ||
| expires_at: datetime | None, | ||
| ) -> None: | ||
| if not keys: | ||
| return | ||
|
|
||
| # Cosmos DB doesn't have native batch upsert in the Python SDK, | ||
| # so we iterate through entries | ||
| for key, managed_entry in zip(keys, managed_entries, strict=True): | ||
| await self._put_managed_entry( | ||
| key=key, | ||
| collection=collection, | ||
| managed_entry=managed_entry, | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# Find ManagedEntry class definition
fd --type f --name "*.py" | xargs grep -l "class ManagedEntry"Repository: strawgate/py-key-value
Length of output: 318
🏁 Script executed:
# Also search for imports of ManagedEntry in the file being reviewed
head -50 src/key_value/aio/stores/cosmosdb/store.py | grep -E "^import|^from"Repository: strawgate/py-key-value
Length of output: 524
🏁 Script executed:
# Check the full imports section of the store.py file
rg -n "^import|^from" src/key_value/aio/stores/cosmosdb/store.py | head -30Repository: strawgate/py-key-value
Length of output: 541
🏁 Script executed:
fd --type f managed_entry.pyRepository: strawgate/py-key-value
Length of output: 137
🏁 Script executed:
# Get file size first
wc -l key_value/shared/managed_entry.pyRepository: strawgate/py-key-value
Length of output: 130
🏁 Script executed:
wc -l src/key_value/shared/managed_entry.pyRepository: strawgate/py-key-value
Length of output: 107
🏁 Script executed:
cat -n src/key_value/shared/managed_entry.pyRepository: strawgate/py-key-value
Length of output: 4900
🏁 Script executed:
fd --type f time_to_live.pyRepository: strawgate/py-key-value
Length of output: 135
🏁 Script executed:
wc -l src/key_value/shared/time_to_live.pyRepository: strawgate/py-key-value
Length of output: 106
🏁 Script executed:
cat -n src/key_value/shared/time_to_live.pyRepository: strawgate/py-key-value
Length of output: 3881
🏁 Script executed:
# Check if there's any validation of timezone info in ManagedEntry or related code
rg -n "tzinfo" src/key_value/Repository: strawgate/py-key-value
Length of output: 721
🏁 Script executed:
# Check how ManagedEntry is created/deserialized in stores
rg -n "ManagedEntry" src/key_value/aio/stores/cosmosdb/store.py | head -20Repository: strawgate/py-key-value
Length of output: 688
🏁 Script executed:
# Check the serialization adapter to see how ManagedEntry is deserialized
rg -n "load_json\|load_entry" src/key_value/aio/stores/cosmosdb/store.py -A 5Repository: strawgate/py-key-value
Length of output: 48
🏁 Script executed:
# Find serialization adapter implementations
fd --type f "serialization" | grep -E "\.py$"Repository: strawgate/py-key-value
Length of output: 137
🏁 Script executed:
wc -l src/key_value/shared/serialization.pyRepository: strawgate/py-key-value
Length of output: 107
🏁 Script executed:
cat -n src/key_value/shared/serialization.pyRepository: strawgate/py-key-value
Length of output: 9233
Normalize expires_at before TTL math to avoid naive/aware mismatches.
The ManagedEntry.expires_at field lacks a timezone constraint and can be deserialized as tz-naive (via datetime.fromisoformat() on strings without timezone info). Subtracting it from the tz-aware now() raises TypeError. This pattern is already implemented in the MongoDB and DuckDB stores; normalizing to UTC here ensures consistency and resilience.
🕒 Suggested normalization
- if managed_entry.expires_at is not None:
- # Store expires_at as ISO string for our own tracking
- item["expires_at"] = managed_entry.expires_at.isoformat()
-
- # Calculate TTL in seconds from now for Cosmos DB's TTL feature
- now = datetime.now(tz=timezone.utc)
- ttl_seconds = int((managed_entry.expires_at - now).total_seconds())
+ if managed_entry.expires_at is not None:
+ expires_at = managed_entry.expires_at
+ if expires_at.tzinfo is None:
+ expires_at = expires_at.replace(tzinfo=timezone.utc)
+
+ # Store expires_at as ISO string for our own tracking
+ item["expires_at"] = expires_at.isoformat()
+
+ # Calculate TTL in seconds from now for Cosmos DB's TTL feature
+ now = datetime.now(tz=timezone.utc)
+ ttl_seconds = int((expires_at - now).total_seconds())🤖 Prompt for AI Agents
In `@src/key_value/aio/stores/cosmosdb/store.py` around lines 408 - 465, The
_put_managed_entry method fails when ManagedEntry.expires_at is a tz-naive
datetime; before doing TTL math, normalize managed_entry.expires_at to an aware
UTC datetime (e.g., if expires_at has no tzinfo, attach UTC) so subtraction
against now = datetime.now(tz=timezone.utc) won't raise; update logic around "if
managed_entry.expires_at is not None" to convert the value to an aware UTC
datetime, then store the ISO string and compute ttl_seconds, keeping use of
_make_item_id, _sanitize_collection, _serialization_adapter.dump_json and
calling _upsert_item unchanged.
| async def clean_cosmosdb_container(store: CosmosDBStore) -> None: | ||
| """Clean all items from the test container.""" | ||
| if store._container is None: | ||
| return | ||
|
|
||
| # Query all items and delete them | ||
| query = "SELECT c.id, c.collection FROM c" | ||
| items: list[dict[str, str]] = [] | ||
| async for item in store._container.query_items(query=query): | ||
| items.append(item) # noqa: PERF401 - async comprehensions not supported here | ||
|
|
||
| for item in items: | ||
| with contextlib.suppress(Exception): | ||
| await store._container.delete_item(item=item["id"], partition_key=item["collection"]) | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
Azure Cosmos SDK Python v4 ContainerProxy query_items enable_cross_partition_query documentation
💡 Result:
In azure-cosmos (Python SDK v4), ContainerProxy.query_items(...) supports the keyword argument enable_cross_partition_query to allow a query to fan out across partitions.
- Allowed values:
None(default) orTrue(there isn’t a documentedFalsebehavior—useNone/ omit it instead). [1] - Usage example (from Microsoft docs):
(Docs show it as a boolean option in examples.) [2]
items = container.query_items( query=sql_stmt, parameters=params, enable_cross_partition_query=False, # commonly shown in samples )
References:
[1] PyPI package docs for azure-cosmos (section “How to use enable_cross_partition_query”) (pypi.org)
[2] Microsoft Learn Python quickstart for Azure Cosmos DB (examples using enable_cross_partition_query) (learn.microsoft.com)
Citations:
- 1: https://pypi.org/project/azure-cosmos/?utm_source=openai
- 2: https://learn.microsoft.com/en-us/azure/cosmos-db/quickstart-python?utm_source=openai
🌐 Web query:
Azure Cosmos DB Python SDK v4 query_items cross-partition query fail without enable_cross_partition_query required
💡 Result:
In azure-cosmos (Python SDK) v4, a query_items() call that isn’t scoped to a single logical partition (i.e., you didn’t pass partition_key=... and your SQL doesn’t resolve to one partition key value) is treated as a cross-partition query, and cross-partition fan-out is disabled unless you explicitly enable it.
What to do
-
Best (recommended): target one partition
items = container.query_items( query="SELECT * FROM c WHERE c.ticker = `@t`", parameters=[{"name":"@t","value":"MSFT"}], partition_key="MSFT", )
If you pass
partition_key, you typically don’t needenable_cross_partition_query. (SDK docs describepartition_keyas targeting the request;Noneimplies cross-partition.) [2] -
If you truly need cross-partition results
items = container.query_items( query="SELECT * FROM c WHERE c.location = `@loc`", parameters=[{"name":"@loc","value":"Seattle"}], enable_cross_partition_query=True, )
The SDK documents
enable_cross_partition_queryas the switch that allows the SDK to send multiple requests when the query isn’t scoped to a single partition key value. [1][2]
Microsoft’s query guidance also explains that queries without a partition-key filter must fan out across partitions. [3]
If you’re on async + hitting a TypeError with enable_cross_partition_query
If you see something like ClientSession._request() got an unexpected keyword argument 'enable_cross_partition_query', that’s a known bug reported for azure-cosmos 4.4.0 async; upgrading to a newer azure-cosmos version is the practical fix. [4]
Sources: PyPI azure-cosmos docs (latest metadata + parameter behavior) [1], Azure SDK docs for query_items/partition_key/enable_cross_partition_query [2], Microsoft Learn on cross-partition queries [3], GitHub issue for async failure in 4.4.0 [4].
Enable cross-partition query when scanning all items in the container.
The cleanup query SELECT c.id, c.collection FROM c is a cross-partition query without a partition key filter. Azure Cosmos SDK v4 disables cross-partition fan-out by default; this query requires enable_cross_partition_query=True to avoid runtime errors.
Suggested fix
- async for item in store._container.query_items(query=query):
+ async for item in store._container.query_items(query=query, enable_cross_partition_query=True):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async def clean_cosmosdb_container(store: CosmosDBStore) -> None: | |
| """Clean all items from the test container.""" | |
| if store._container is None: | |
| return | |
| # Query all items and delete them | |
| query = "SELECT c.id, c.collection FROM c" | |
| items: list[dict[str, str]] = [] | |
| async for item in store._container.query_items(query=query): | |
| items.append(item) # noqa: PERF401 - async comprehensions not supported here | |
| for item in items: | |
| with contextlib.suppress(Exception): | |
| await store._container.delete_item(item=item["id"], partition_key=item["collection"]) | |
| async def clean_cosmosdb_container(store: CosmosDBStore) -> None: | |
| """Clean all items from the test container.""" | |
| if store._container is None: | |
| return | |
| # Query all items and delete them | |
| query = "SELECT c.id, c.collection FROM c" | |
| items: list[dict[str, str]] = [] | |
| async for item in store._container.query_items(query=query, enable_cross_partition_query=True): | |
| items.append(item) # noqa: PERF401 - async comprehensions not supported here | |
| for item in items: | |
| with contextlib.suppress(Exception): | |
| await store._container.delete_item(item=item["id"], partition_key=item["collection"]) |
🤖 Prompt for AI Agents
In `@tests/stores/cosmosdb/test_cosmosdb.py` around lines 32 - 46, The cleanup
query in clean_cosmosdb_container performs a cross-partition scan but does not
enable cross-partition queries; update the call to store._container.query_items
(in function clean_cosmosdb_container) to pass enable_cross_partition_query=True
so the SDK will allow fan-out queries across partitions and avoid runtime errors
when iterating items.
1 participant
Summary
Add CosmosDBStore for Azure Cosmos DB NoSQL API support.
Changes
Closes #292
Generated with Claude Code
Summary by CodeRabbit
New Features
Chores
Tests