Add routing wrapper

[strawgate]

A routing wrapper is a helper wrapper that allows routing requests to different stores using a routing function to return the store

We'll ship some simple ones that do like collection to store via a map

[claude]

Claude finished @strawgate's task —— View job

---

Issue #114 Analysis: Add routing wrapper

---

Recommendation

Implement a RoutingWrapper that routes operations to different backing stores based on a routing function. This is a natural extension to the existing wrapper pattern and addresses a clear use case: directing different collections (or keys) to different storage backends.

The feature should include:

1. RoutingWrapper - A base wrapper that accepts a routing function to determine which store handles each request
2. CollectionRoutingWrapper - A pre-built implementation that routes based on collection name using a simple map (matches the "collection to store via a map" requirement from the issue)

This follows established patterns in the codebase (like FallbackWrapper) and provides both flexibility and convenience.

---

Findings

Architecture Context

The py-key-value project has a well-established wrapper architecture:

• Base Pattern: All wrappers extend BaseWrapper which implements the AsyncKeyValue protocol (base.py:9-52)
• Protocol-based: Wrappers implement AsyncKeyValueProtocol with methods: get, put, delete, ttl, and their _many variants
• Existing multi-store pattern: FallbackWrapper already demonstrates routing between multiple stores based on exception handling (fallback/wrapper.py:10-119)
• Collection-based storage: The protocol uses collection as an optional parameter to organize keys into logical namespaces

Key Observations:

1. No existing wrapper routes based on collection or key patterns
2. The FallbackWrapper shows it's viable to route operations to different stores
3. The SingleCollectionWrapper shows collection name manipulation patterns
4. The MultiDiskStore creates separate store instances per collection, but this is at the store level, not wrapper level

Use Cases

A routing wrapper would enable several valuable patterns:

1. Multi-region/Multi-backend Storage: Route different data types to optimal backends

   # Route sessions to Redis, user data to DynamoDB
   router = CollectionRoutingWrapper(
       collection_map={
           "sessions": RedisStore(url="redis://localhost"),
           "users": DynamoDBStore(table_name="users"),
       },
       default_store=MemoryStore()
   )

2. Cost Optimization: Route hot data to expensive fast storage, cold data to cheaper storage

3. Compliance: Route PII data to encrypted stores, other data to standard stores

4. Development/Testing: Route certain collections to mock stores while others use real backends

Existing Patterns Analysis

Similar implementations in the codebase:

Pattern	File	Relevance
FallbackWrapper	fallback/wrapper.py	Routes to secondary store on exception - demonstrates multi-store routing
SingleCollectionWrapper	single_collection/wrapper.py	Shows collection name manipulation patterns
MultiDiskStore	disk/multi_store.py	Creates per-collection stores but at store level
BaseWrapper	base.py	Standard passthrough pattern all wrappers follow

Key implementation insights:

• Wrappers use composition, not inheritance (hold reference to AsyncKeyValue)
• All 8 protocol methods must be implemented (get, put, delete, ttl, and _many variants)
• Type signatures must match protocol exactly
• Test structure follows BaseStoreTests pattern

---

Detailed Action Plan

1. Create RoutingWrapper base class

File: key-value/key-value-aio/src/key_value/aio/wrappers/routing/wrapper.py

from collections.abc import Callable, Sequence
from typing import Any, SupportsFloat
from typing_extensions import override
from key_value.aio.protocols.key_value import AsyncKeyValue
from key_value.aio.wrappers.base import BaseWrapper

RoutingFunction = Callable[[str | None], AsyncKeyValue]

class RoutingWrapper(BaseWrapper):
    """Routes operations to different stores based on a routing function.
    
    The routing function receives the collection name and returns the appropriate store.
    """
    
    def __init__(
        self,
        routing_function: RoutingFunction,
        default_store: AsyncKeyValue | None = None,
    ) -> None:
        """Initialize the routing wrapper.
        
        Args:
            routing_function: Function that takes a collection name and returns the store to use
            default_store: Optional fallback store if routing_function returns None
        """
        self.routing_function = routing_function
        self.default_store = default_store
    
    def _get_store(self, collection: str | None) -> AsyncKeyValue:
        """Get the appropriate store for the given collection."""
        store = self.routing_function(collection)
        if store is None and self.default_store is not None:
            return self.default_store
        if store is None:
            msg = f"No store found for collection: {collection}"
            raise ValueError(msg)
        return store
    
    @override
    async def get(self, key: str, *, collection: str | None = None) -> dict[str, Any] | None:
        store = self._get_store(collection)
        return await store.get(key=key, collection=collection)
    
    # ... implement all other protocol methods similarly (get_many, put, put_many, 
    # delete, delete_many, ttl, ttl_many)

2. Create CollectionRoutingWrapper convenience class

File: key-value/key-value-aio/src/key_value/aio/wrappers/routing/collection_routing.py

from key_value.aio.protocols.key_value import AsyncKeyValue
from key_value.aio.wrappers.routing.wrapper import RoutingWrapper

class CollectionRoutingWrapper(RoutingWrapper):
    """Routes operations based on collection name using a simple map.
    
    Example:
        router = CollectionRoutingWrapper(
            collection_map={
                "sessions": redis_store,
                "users": dynamo_store,
            },
            default_store=memory_store
        )
    """
    
    def __init__(
        self,
        collection_map: dict[str, AsyncKeyValue],
        default_store: AsyncKeyValue | None = None,
    ) -> None:
        """Initialize collection-based routing.
        
        Args:
            collection_map: Mapping from collection name to store
            default_store: Store to use for unmapped collections
        """
        self.collection_map = collection_map
        
        def route_by_collection(collection: str | None) -> AsyncKeyValue | None:
            if collection is None:
                return default_store
            return self.collection_map.get(collection, default_store)
        
        super().__init__(
            routing_function=route_by_collection,
            default_store=default_store
        )

3. Create comprehensive tests

File: key-value/key-value-aio/tests/stores/wrappers/test_routing.py

Test cases should cover:

• Basic routing to different stores based on collection
• Fallback to default_store for unmapped collections
• Error when no default and collection unmapped
• All CRUD operations (get, put, delete, ttl)
• Bulk operations (get_many, put_many, delete_many, ttl_many)
• Custom routing functions beyond collection mapping

4. Update exports and documentation

• Add to key-value/key-value-aio/src/key_value/aio/wrappers/__init__.py
• Create key-value/key-value-aio/src/key_value/aio/wrappers/routing/__init__.py
• Add usage examples to README.md under the wrappers section
• Document the pattern in docstrings with clear examples

5. Generate sync version

After async implementation is complete:

uv run python scripts/build_sync_library.py

This will auto-generate the sync version in key-value-sync.

6. Run full test suite

uv run pytest key-value/key-value-aio/tests -v
uv run pytest key-value/key-value-sync/tests -v
make lint
make typecheck

---

Related Items

Related Issues and Pull Requests

Repository	Issue or PR	Relevance
strawgate/py-key-value	#49: Add store middleware	Discusses evolution of wrapper architecture - routing wrapper should consider future middleware design
strawgate/py-key-value	#50: Add Retry, Timeout, Fallback... Wrappers	Merged PR showing wrapper implementation patterns to follow
strawgate/py-key-value	#75: Add EncryptionWrapper	Example of recent wrapper implementation

Related Files

Repository	File	Relevance	Sections
strawgate/py-key-value	wrappers/base.py	Base class for all wrappers - must extend this	9-52
strawgate/py-key-value	wrappers/fallback/wrapper.py	Closest pattern - routes between two stores	10-119
strawgate/py-key-value	protocols/key_value.py	Protocol definition that must be implemented	5-181
strawgate/py-key-value	wrappers/single_collection/wrapper.py	Shows collection name handling patterns	12-85
strawgate/py-key-value	tests/stores/wrappers/test_fallback.py	Test pattern for multi-store wrappers	23-77

---

[strawgate]

@claude let's add a routing wrapper

[claude]

Claude finished @strawgate's task —— View job • claude/issue-114-20251026-0414 • Create PR ➔

---

Adding Routing Wrapper

I've successfully implemented the routing wrapper feature for py-key-value!

Summary

Created two new wrapper classes that enable flexible multi-backend storage routing:

1. RoutingWrapper - Base wrapper that accepts a custom routing function to determine which store handles each request
2. CollectionRoutingWrapper - Convenience wrapper that routes based on collection names using a simple dictionary map

Implementation Details

Files Created:

• key-value/key-value-aio/src/key_value/aio/wrappers/routing/wrapper.py - Base routing wrapper
• key-value/key-value-aio/src/key_value/aio/wrappers/routing/collection_routing.py - Collection-based routing
• key-value/key-value-aio/src/key_value/aio/wrappers/routing/__init__.py - Package exports
• key-value/key-value-aio/tests/stores/wrappers/test_routing.py - Comprehensive test suite

Key Features:

• Support for all AsyncKeyValue protocol methods (get, put, delete, ttl, and their _many variants)
• Flexible custom routing functions
• Optional default store fallback
• Proper error handling when no store is found

Testing:

• ✅ 11/11 tests passed
• ✅ Python linting passed
• ✅ Type checking passed for routing wrapper code

Use Cases Enabled:

# Route different data types to optimal backends
router = CollectionRoutingWrapper(
    collection_map={
        "sessions": redis_store,
        "users": dynamo_store,
        "cache": memory_store,
    },
    default_store=disk_store
)

# Custom routing logic
def route_by_pattern(collection: str | None) -> AsyncKeyValue | None:
    if collection and collection.startswith("temp_"):
        return memory_store
    return persistent_store

router = RoutingWrapper(routing_function=route_by_pattern)

Next Steps

The sync version will be generated automatically by the CI pipeline (requires Python 3.10). The implementation is ready for review!

--- |