Add azure cosmosdb store

[strawgate]

[claude]

Claude finished @strawgate's task —— View job

---

Issue #44 Analysis Complete ✓

---

Recommendation

Implement an AzureCosmosDBStore class following the established pattern from MongoDB and DynamoDB stores. Azure Cosmos DB is a strong fit for py-key-value due to its native TTL support, async Python SDK, and document-oriented NoSQL API that aligns well with the existing store architecture.

The implementation should:

1. Use the azure-cosmos package (latest: 4.14.0) with async client from azure.cosmos.aio
2. Store values as documents in a Cosmos DB container with partition key on collection and item id as key
3. Leverage Cosmos DB's native TTL support via the ttl property on documents
4. Follow the pattern established by MongoDB and DynamoDB stores for consistency
5. Include comprehensive test coverage using Azure Cosmos DB Emulator via Docker

---

Findings

Azure Cosmos DB Capabilities

Native TTL Support:

• Cosmos DB provides built-in TTL at both container and item level
• TTL is set in seconds via the ttl property on documents
• Items with TTL automatically expire without manual cleanup
• Container-level defaultTtl can be set to -1 (no expiration) or a positive number
• Per-item TTL overrides container-level TTL

Async Python SDK:

• Package: azure-cosmos (v4.14.0, released Oct 2025)
• Async client: from azure.cosmos.aio import CosmosClient
• Supports all standard CRUD operations: create_item, read_item, delete_item, upsert_item
• Context manager support with async with
• Requires Python ≥3.9

Storage Model:

• Document-oriented NoSQL database (similar to MongoDB)
• Supports partition keys for horizontal scaling
• Can store JSON documents natively
• Supports atomic operations on single items

Store Architecture Pattern Analysis

Based on examination of existing stores (MongoDB at key-value-aio/src/key_value/aio/stores/mongodb/store.py:1-210 and DynamoDB at key-value-aio/src/key_value/aio/stores/dynamodb/store.py:1-236):

Common Pattern:

1. Inherit from BaseStore, BaseContextManagerStore
2. Optional mixins: BaseEnumerateCollectionsStore, BaseDestroyCollectionStore
3. Implement core methods: _get_managed_entry, _put_managed_entry, _delete_managed_entry
4. Support two initialization modes: client injection or connection parameters
5. Implement __aenter__/__aexit__ for async context manager
6. Use _setup() for one-time initialization (ensure database/container exists)
7. Optional: _setup_collection() for per-collection initialization
8. Store values as ManagedEntry.to_json() or native documents (trend toward native)

Key Differences Between Stores:

• MongoDB: Uses collections within a database, stores as JSON strings currently (being migrated to native BSON in PR feat: Switch MongoDB and Elasticsearch to native document storage #104)
• DynamoDB: Uses composite primary key (collection + key), stores as JSON strings in a single table
• Elasticsearch: Uses indices, being migrated to native flattened field storage (PR feat: Switch MongoDB and Elasticsearch to native document storage #104)

Recent Architecture Trends

Issue #88 and PR #104 indicate the project is moving toward:

• Storing values natively when the backend supports it (rather than JSON strings)
• Using native TTL when available (rather than managed TTL)
• This would apply to Cosmos DB since it supports both

---

Detailed Action Plan

Implementation Steps

A junior developer could implement this by following these steps:

1. Create Store Implementation

File: key-value/key-value-aio/src/key_value/aio/stores/cosmosdb/store.py

from typing import Any, overload
from key_value.shared.utils.managed_entry import ManagedEntry
from typing_extensions import Self, override
from key_value.aio.stores.base import BaseContextManagerStore, BaseStore

try:
    from azure.cosmos.aio import CosmosClient
    from azure.cosmos.aio import ContainerProxy
    from azure.cosmos import PartitionKey
except ImportError as e:
    msg = "CosmosDBStore requires py-key-value-aio[cosmosdb]"
    raise ImportError(msg) from e

Key implementation details:

• Use partition key: collection (like DynamoDB pattern)
• Use item id: key
• Store ManagedEntry as JSON string in value field initially (can migrate to native later)
• Map ManagedEntry.ttl to document ttl property
• Include created_at, expires_at for compatibility
• Implement overloaded __init__ for both client injection and URL-based init
• Use upsert_item for _put_managed_entry (atomic create/update)
• Implement _setup() to create database and container if they don't exist
• Set container with partition key on /collection and default TTL of -1

2. Add Package Dependency

File: key-value/key-value-aio/pyproject.toml

Add to [project.optional-dependencies]:

cosmosdb = ["azure-cosmos>=4.14.0"]

Add to [dependency-groups] dev dependencies:

"py-key-value-aio[cosmosdb]",

3. Create init.py

File: key-value/key-value-aio/src/key_value/aio/stores/cosmosdb/__init__.py

from key_value.aio.stores.cosmosdb.store import CosmosDBStore

__all__ = ["CosmosDBStore"]

4. Create Test Suite

File: key-value/key-value-aio/tests/stores/cosmosdb/test_cosmosdb.py

Follow pattern from test_mongodb.py and test_dynamodb.py:

• Use Docker container for Cosmos DB Emulator: mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator
• Implement ping_cosmosdb() health check function
• Create session-scoped fixture setup_cosmosdb with Docker container
• Inherit from BaseStoreTests and ContextManagerStoreTestMixin
• Override test_not_unbounded and skip (distributed caches are unbounded)
• Clean up test database in store fixture before tests
• Add Cosmos-specific tests (e.g., partition key handling)

Note: Azure Cosmos DB Emulator Docker setup:

with docker_container(
    "cosmosdb-test",
    "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator",
    {"8081": 8081, "10251": 10251, "10252": 10252, "10253": 10253, "10254": 10254}
):

Connection string for emulator:

"AccountEndpoint=https://localhost:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="

5. Update Documentation

File: README.md

Add Cosmos DB to the list of supported backends:

- **Multiple backends**: DynamoDB, Elasticsearch, Memcached, MongoDB, Redis,
  RocksDB, Valkey, Azure Cosmos DB, and In-memory, Disk, etc

Add usage example in appropriate section.

6. Generate Sync Version

Run the sync library generator:

uv run python scripts/build_sync_library.py

This will auto-generate the synchronous version at key-value/key-value-sync/src/key_value/sync/code_gen/stores/cosmosdb/

7. Testing & Validation

# Install dependencies
uv sync --all-extras --all-packages

# Run tests
uv run pytest key-value/key-value-aio/tests/stores/cosmosdb -v

# Run linting
make lint

# Run type checking
make typecheck

# Generate sync library
uv run python scripts/build_sync_library.py

Challenges & Considerations

1. Cosmos DB Emulator Limitations:

   • Linux emulator may have specific SSL/certificate requirements
   • May need to disable SSL verification in tests: connection_verify=False
   • Emulator startup can be slow (30-60 seconds)

2. Partition Key Strategy:

   • Using collection as partition key provides good isolation
   • May need to consider partition key limits (20 GB per partition)
   • Alternative: hash-based partition key for better distribution at scale

3. Cost Considerations:

   • Document in README that Cosmos DB is a premium service
   • Mention Request Units (RU) consumption patterns
   • Suggest using emulator for development

4. TTL Behavior:

   • Cosmos DB TTL cleanup is eventual (within minutes to hours)
   • Unlike Redis which expires immediately
   • May need to document this difference

5. Future Native Storage:

   • Per PR feat: Switch MongoDB and Elasticsearch to native document storage #104 trend, consider storing values as native objects
   • Would allow querying within stored values
   • Can implement initially as JSON strings, migrate later

---

Related Items

Related Issues and Pull Requests

Repository	Issue/PR	Relevance
strawgate/py-key-value	Issue #48: Add postgresql data store	Similar request to add a new database backend store
strawgate/py-key-value	Issue #90: Review mongodb for inclusion as stable store	MongoDB review addressing native document storage strategy (applicable to Cosmos DB)
strawgate/py-key-value	Issue #88: Store values directly for stores that natively support ttl	Discusses native TTL support strategy (Cosmos DB has native TTL)
strawgate/py-key-value	PR #104: Switch MongoDB and Elasticsearch to native document storage	Shows current trend toward native storage; pattern to follow for Cosmos DB
strawgate/py-key-value	PR #54: Add DynamoDB Store	Successfully merged DynamoDB implementation to reference
strawgate/py-key-value	PR #58: Add Aerospike store implementation	Example of adding distributed NoSQL store with native TTL

---

Related Files

Repository	File	Relevance	Sections
strawgate/py-key-value	mongodb/store.py	Best reference implementation - document database with similar data model to Cosmos DB	43-210
strawgate/py-key-value	dynamodb/store.py	Reference for cloud NoSQL service pattern, partition key strategy	29-236
strawgate/py-key-value	stores/base.py	Base class defining store interface and required methods	28-100
strawgate/py-key-value	tests/stores/mongodb/test_mongodb.py	Test pattern for document databases with Docker container setup	38-74
strawgate/py-key-value	tests/stores/dynamodb/test_dynamodb.py	Test pattern for cloud NoSQL services with local emulator	44-80
strawgate/py-key-value	pyproject.toml	Where to add cosmosdb optional dependency	34-52

---

Related Webpages

Name	URL	Relevance
Azure Cosmos DB Python SDK Documentation	https://learn.microsoft.com/en-us/python/api/overview/azure/cosmos-readme	Official documentation for the azure-cosmos Python package, including async client usage
Azure Cosmos DB TTL Documentation	https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/time-to-live	Comprehensive guide on how TTL works in Cosmos DB at container and item level
azure-cosmos PyPI Package	https://pypi.org/project/azure-cosmos/	Package information, latest version (4.14.0), and requirements
Azure Cosmos DB Async Python SDK Blog	https://devblogs.microsoft.com/cosmosdb/async-python-sdk-preview/	Detailed explanation of async IO support in the Python SDK
Azure SDK for Python - Cosmos Samples	https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples	Official code samples showing create_item, read_item, delete_item patterns

---