refactor: stores register cleanup callbacks directly with exit stack

Simplified the BaseContextManagerStore lifecycle management pattern by:
- Removing abstract _register_cleanup_callbacks method
- Initializing AsyncExitStack in __init__ instead of lazily
- Allowing stores to directly register cleanup callbacks via _exit_stack
- Stores now call _exit_stack.push_async_callback() or _exit_stack.callback()
  in their _setup() method to register cleanup operations

This gives stores direct control over exit stack management while keeping
the base class responsible for entering and exiting the stack.

Updated all stores that inherit from BaseContextManagerStore:
- MongoDB, DynamoDB: register client context manager
- Redis, Valkey, Memcached, Elasticsearch: register async close callbacks
- RocksDB, Disk, MultiDisk, DuckDB: register sync close callbacks

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: github-actions[bot]

key-value/key-value-aio/src/key_value/aio/stores/base.py

@@ -434,13 +434,14 @@ class BaseContextManagerStore(BaseStore, ABC):
     the constructor. This ensures the store does not manage the lifecycle of user-provided
     clients (i.e., does not close them).
 
-    Stores that have clients requiring context manager entry should override
-    `_register_cleanup_callbacks(stack)` to register their cleanup with the provided exit stack.
-    This method is only called when the store owns the client (client_provided_by_user=False).
+    The base class provides an AsyncExitStack that stores can use to register cleanup
+    callbacks. Stores should add their cleanup operations to the exit stack as needed.
+    The base class handles entering and exiting the exit stack.
     """
 
     _client_provided_by_user: bool
-    _exit_stack: AsyncExitStack | None
+    _exit_stack: AsyncExitStack
+    _exit_stack_entered: bool
 
     def __init__(self, *, client_provided_by_user: bool = False, **kwargs: Any) -> None:
         """Initialize the context manager store with client ownership configuration.
@@ -452,73 +453,44 @@ def __init__(self, *, client_provided_by_user: bool = False, **kwargs: Any) -> N
             **kwargs: Additional arguments to pass to the base store constructor.
         """
         self._client_provided_by_user = client_provided_by_user
-        self._exit_stack = None
-        super().__init__(**kwargs)
-
-    async def _register_cleanup_callbacks(self, stack: AsyncExitStack) -> None:
-        """Register cleanup callbacks with the exit stack.
-
-        Stores should override this method to register their cleanup callbacks with the exit stack.
-        This method is only called when the store owns the client (client_provided_by_user=False).
-
-        Examples:
-            # For context manager clients:
-            await stack.enter_async_context(self._client)
-
-            # For clients with close() methods:
-            stack.push_async_callback(self._client.aclose)
-
-        Args:
-            stack: The AsyncExitStack to register cleanup callbacks with.
-        """
-
-    async def _ensure_exit_stack(self) -> AsyncExitStack:
-        """Ensure the exit stack exists and register cleanup callbacks if needed.
-
-        Returns:
-            The exit stack instance.
-        """
-        if self._exit_stack is not None:
-            return self._exit_stack
-
         self._exit_stack = AsyncExitStack()
-        await self._exit_stack.__aenter__()
-
-        # Register cleanup callbacks if we own the client
-        if not self._client_provided_by_user:
-            await self._register_cleanup_callbacks(self._exit_stack)
+        self._exit_stack_entered = False
+        super().__init__(**kwargs)
 
-        return self._exit_stack
+    async def _ensure_exit_stack_entered(self) -> None:
+        """Ensure the exit stack has been entered."""
+        if not self._exit_stack_entered:
+            await self._exit_stack.__aenter__()
+            self._exit_stack_entered = True
 
     async def __aenter__(self) -> Self:
-        # Create exit stack and enter client context
-        await self._ensure_exit_stack()
+        # Enter the exit stack
+        await self._ensure_exit_stack_entered()
         await self.setup()
         return self
 
     async def __aexit__(
         self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
     ) -> None:
         # Close the exit stack, which handles all cleanup
-        if self._exit_stack is not None:
-            await self._exit_stack.aclose()
-            self._exit_stack = None
+        if self._exit_stack_entered:
+            await self._exit_stack.__aexit__(exc_type, exc_value, traceback)
+            self._exit_stack_entered = False
 
     async def close(self) -> None:
-        # Close the exit stack if it exists
-        if self._exit_stack is not None:
+        # Close the exit stack if it has been entered
+        if self._exit_stack_entered:
             await self._exit_stack.aclose()
-            self._exit_stack = None
+            self._exit_stack_entered = False
 
     async def setup(self) -> None:
         """Initialize the store if not already initialized.
 
-        This override ensures that if a client needs context manager entry, the exit stack
-        is created and the client is entered before the store's _setup() method is called.
-        This allows stores to work correctly even when not used with `async with`.
+        This override ensures the exit stack is entered before the store's _setup()
+        method is called, allowing stores to register cleanup callbacks during setup.
         """
-        # Ensure exit stack exists and client context is entered (if needed)
-        await self._ensure_exit_stack()
+        # Ensure exit stack is entered
+        await self._ensure_exit_stack_entered()
         # Call parent setup
         await super().setup()
 

key-value/key-value-aio/src/key_value/aio/stores/disk/multi_store.py

@@ -106,6 +106,11 @@ def default_disk_cache_factory(collection: str) -> Cache:
             stable_api=True,
         )
 
+    @override
+    async def _setup(self) -> None:
+        """Register cache cleanup."""
+        self._exit_stack.callback(self._sync_close)
+
     @override
     async def _setup_collection(self, *, collection: str) -> None:
         self._cache[collection] = self._disk_cache_factory(collection)
@@ -148,8 +153,5 @@ def _sync_close(self) -> None:
         for cache in self._cache.values():
             cache.close()
 
-    async def _close(self) -> None:
-        self._sync_close()
-
     def __del__(self) -> None:
         self._sync_close()

key-value/key-value-aio/src/key_value/aio/stores/disk/store.py

@@ -86,6 +86,12 @@ def __init__(
             stable_api=True,
         )
 
+    @override
+    async def _setup(self) -> None:
+        """Register cache cleanup if we own the cache."""
+        if not self._client_provided_by_user:
+            self._exit_stack.callback(self._cache.close)
+
     @override
     async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:
         combo_key: str = compound_key(collection=collection, key=key)
@@ -126,9 +132,6 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
 
         return self._cache.delete(key=combo_key, retry=True)
 
-    async def _close(self) -> None:
-        self._cache.close()
-
     def __del__(self) -> None:
         if not getattr(self, "_client_provided_by_user", False) and hasattr(self, "_cache"):
             self._cache.close()

key-value/key-value-aio/src/key_value/aio/stores/duckdb/store.py

@@ -263,6 +263,16 @@ async def _setup(self) -> None:
         - Metadata queries without JSON deserialization
         - Native JSON column support for rich querying capabilities
         """
+        # Register connection cleanup if we own the connection
+        if not self._client_provided_by_user:
+
+            def close_connection() -> None:
+                if not self._is_closed:
+                    self._connection.close()
+                    self._is_closed = True
+
+            self._exit_stack.callback(close_connection)
+
         # Create the main table for storing key-value entries
         self._connection.execute(self._get_create_table_sql())
 
@@ -362,12 +372,6 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
         deleted_rows = result.fetchall()
         return len(deleted_rows) > 0
 
-    async def _close(self) -> None:
-        """Close the DuckDB connection."""
-        if not self._is_closed:
-            self._connection.close()
-            self._is_closed = True
-
     def __del__(self) -> None:
         """Clean up the DuckDB connection on deletion."""
         try:

key-value/key-value-aio/src/key_value/aio/stores/dynamodb/store.py

@@ -1,4 +1,3 @@
-from contextlib import AsyncExitStack
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, overload
 
@@ -123,12 +122,6 @@ def __init__(
             client_provided_by_user=client_provided,
         )
 
-    @override
-    async def _register_cleanup_callbacks(self, stack: AsyncExitStack) -> None:
-        """Register DynamoDB client cleanup with the exit stack."""
-        if hasattr(self, "_raw_client"):
-            self._client = await stack.enter_async_context(self._raw_client)
-
     @property
     def _connected_client(self) -> DynamoDBClient:
         if not self._client:
@@ -139,6 +132,9 @@ def _connected_client(self) -> DynamoDBClient:
     @override
     async def _setup(self) -> None:
         """Setup the DynamoDB client and ensure table exists."""
+        # Register client cleanup if we own the client
+        if not self._client_provided_by_user and hasattr(self, "_raw_client"):
+            self._client = await self._exit_stack.enter_async_context(self._raw_client)
         try:
             await self._connected_client.describe_table(TableName=self._table_name)  # pyright: ignore[reportUnknownMemberType]
         except self._connected_client.exceptions.ResourceNotFoundException:  # pyright: ignore[reportUnknownMemberType]

key-value/key-value-aio/src/key_value/aio/stores/elasticsearch/store.py

@@ -259,6 +259,10 @@ def __init__(
 
     @override
     async def _setup(self) -> None:
+        # Register client cleanup if we own the client
+        if not self._client_provided_by_user:
+            self._exit_stack.push_async_callback(self._client.close)
+
         cluster_info = await self._client.options(ignore_status=404).info()
 
         self._is_serverless = cluster_info.get("version", {}).get("build_flavor") == "serverless"
@@ -551,6 +555,3 @@ async def _cull(self) -> None:
                 },
             },
         )
-
-    async def _close(self) -> None:
-        await self._client.close()

key-value/key-value-aio/src/key_value/aio/stores/memcached/store.py

@@ -81,6 +81,12 @@ def __init__(
             client_provided_by_user=client_provided,
         )
 
+    @override
+    async def _setup(self) -> None:
+        """Register client cleanup if we own the client."""
+        if not self._client_provided_by_user:
+            self._exit_stack.push_async_callback(self._client.close)
+
     @override
     async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:
         combo_key: str = self._sanitize_key(compound_key(collection=collection, key=key))
@@ -151,6 +157,3 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
     async def _delete_store(self) -> bool:
         _ = await self._client.flush_all()
         return True
-
-    async def _close(self) -> None:
-        await self._client.close()

key-value/key-value-aio/src/key_value/aio/stores/mongodb/store.py

@@ -1,5 +1,4 @@
 from collections.abc import Sequence
-from contextlib import AsyncExitStack
 from datetime import datetime, timezone
 from typing import Any, overload
 
@@ -194,9 +193,10 @@ def __init__(
         )
 
     @override
-    async def _register_cleanup_callbacks(self, stack: AsyncExitStack) -> None:
-        """Register MongoDB client cleanup with the exit stack."""
-        await stack.enter_async_context(self._client)
+    async def _setup(self) -> None:
+        """Register client cleanup if we own the client."""
+        if not self._client_provided_by_user:
+            await self._exit_stack.enter_async_context(self._client)
 
     @override
     async def _setup_collection(self, *, collection: str) -> None:

key-value/key-value-aio/src/key_value/aio/stores/redis/store.py

@@ -218,9 +218,12 @@ async def _get_collection_keys(self, *, collection: str, limit: int | None = Non
 
         return get_keys_from_compound_keys(compound_keys=keys, collection=collection)
 
+    @override
+    async def _setup(self) -> None:
+        """Register client cleanup if we own the client."""
+        if not self._client_provided_by_user:
+            self._exit_stack.push_async_callback(self._client.aclose)
+
     @override
     async def _delete_store(self) -> bool:
         return await self._client.flushdb()  # pyright: ignore[reportUnknownMemberType, reportAny]
-
-    async def _close(self) -> None:
-        await self._client.aclose()

key-value/key-value-aio/src/key_value/aio/stores/rocksdb/store.py

@@ -85,8 +85,12 @@ def __init__(
             client_provided_by_user=client_provided,
         )
 
-    async def _close(self) -> None:
-        self._close_and_flush()
+    @override
+    async def _setup(self) -> None:
+        """Register database cleanup if we own the database."""
+        if not self._client_provided_by_user:
+            # Register a callback to close and flush the database
+            self._exit_stack.callback(self._close_and_flush)
 
     def _close_and_flush(self) -> None:
         if not self._is_closed: