refactor: stores register cleanup callbacks directly with exit stack

Previously, stores returned clients via `_get_client_for_context()` and the base
class managed them. Now stores directly register their cleanup callbacks with the
exit stack via `_register_cleanup_callbacks(stack)`.

This gives stores more control over their cleanup registration and simplifies
the base class by removing an abstraction layer.

Also updated codegen to map AsyncExitStack -> ExitStack and
enter_async_context -> enter_context for sync code generation.

🤖 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,14 +434,13 @@ 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 can override `_get_client_for_context()`
-    to return the client that needs to be entered. The base class will automatically enter the client's
-    context using an exit stack, which handles cleanup automatically.
+    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).
     """
 
     _client_provided_by_user: bool
     _exit_stack: AsyncExitStack | None
-    _entered_client_result: Any
 
     def __init__(self, *, client_provided_by_user: bool = False, **kwargs: Any) -> None:
         """Initialize the context manager store with client ownership configuration.
@@ -454,32 +453,27 @@ def __init__(self, *, client_provided_by_user: bool = False, **kwargs: Any) -> N
         """
         self._client_provided_by_user = client_provided_by_user
         self._exit_stack = None
-        self._entered_client_result = None
         super().__init__(**kwargs)
 
-    def _get_client_for_context(self) -> Any | None:
-        """Return the client that needs context manager entry, or None if not applicable.
+    async def _register_cleanup_callbacks(self, stack: AsyncExitStack) -> None:
+        """Register cleanup callbacks with the exit stack.
 
-        Stores with clients that require context manager entry (e.g., MongoDB's AsyncMongoClient,
-        DynamoDB's aioboto3 client) should override this method to return the raw client object
-        that needs to be entered.
+        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).
 
-        Returns:
-            The client object that has __aenter__/__aexit__ methods, or None if no context
-            management is needed.
-        """
-        return None
+        Examples:
+            # For context manager clients:
+            await stack.enter_async_context(self._client)
 
-    def _get_entered_client(self) -> Any | None:
-        """Get the result of entering the client's context manager.
+            # For clients with close() methods:
+            stack.push_async_callback(self._client.aclose)
 
-        Returns:
-            The entered client (typically the client itself), or None if context wasn't entered.
+        Args:
+            stack: The AsyncExitStack to register cleanup callbacks with.
         """
-        return self._entered_client_result
 
     async def _ensure_exit_stack(self) -> AsyncExitStack:
-        """Ensure the exit stack exists and enter client context if needed.
+        """Ensure the exit stack exists and register cleanup callbacks if needed.
 
         Returns:
             The exit stack instance.
@@ -490,12 +484,9 @@ async def _ensure_exit_stack(self) -> AsyncExitStack:
         self._exit_stack = AsyncExitStack()
         await self._exit_stack.__aenter__()
 
-        # Enter client context if we own it
+        # Register cleanup callbacks if we own the client
         if not self._client_provided_by_user:
-            client = self._get_client_for_context()
-            if client is not None and hasattr(client, "__aenter__"):
-                result = await self._exit_stack.enter_async_context(client)
-                self._entered_client_result = result
+            await self._register_cleanup_callbacks(self._exit_stack)
 
         return self._exit_stack
 

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

@@ -148,7 +148,6 @@ def _sync_close(self) -> None:
         for cache in self._cache.values():
             cache.close()
 
-    @override
     async def _close(self) -> None:
         self._sync_close()
 

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

@@ -126,7 +126,6 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
 
         return self._cache.delete(key=combo_key, retry=True)
 
-    @override
     async def _close(self) -> None:
         self._cache.close()
 

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

@@ -362,7 +362,6 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
         deleted_rows = result.fetchall()
         return len(deleted_rows) > 0
 
-    @override
     async def _close(self) -> None:
         """Close the DuckDB connection."""
         if not self._is_closed:

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

@@ -1,3 +1,4 @@
+from contextlib import AsyncExitStack
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, overload
 
@@ -123,9 +124,10 @@ def __init__(
         )
 
     @override
-    def _get_client_for_context(self) -> Any | None:
-        """Return the raw client (context manager) for context management."""
-        return self._raw_client if hasattr(self, "_raw_client") else None
+    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:
@@ -137,13 +139,6 @@ def _connected_client(self) -> DynamoDBClient:
     @override
     async def _setup(self) -> None:
         """Setup the DynamoDB client and ensure table exists."""
-
-        # Get the entered client from the base class context management
-        if not self._client:
-            result = self._get_entered_client()
-            if result is not None:
-                self._client = result
-
         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

@@ -552,6 +552,5 @@ async def _cull(self) -> None:
             },
         )
 
-    @override
     async def _close(self) -> None:
         await self._client.close()

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

@@ -152,6 +152,5 @@ async def _delete_store(self) -> bool:
         _ = await self._client.flush_all()
         return True
 
-    @override
     async def _close(self) -> None:
         await self._client.close()

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

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+from contextlib import AsyncExitStack
 from datetime import datetime, timezone
 from typing import Any, overload
 
@@ -193,9 +194,9 @@ def __init__(
         )
 
     @override
-    def _get_client_for_context(self) -> Any | None:
-        """Return the MongoDB client for context management."""
-        return self._client
+    async def _register_cleanup_callbacks(self, stack: AsyncExitStack) -> None:
+        """Register MongoDB client cleanup with the exit stack."""
+        await 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

@@ -222,6 +222,5 @@ async def _get_collection_keys(self, *, collection: str, limit: int | None = Non
     async def _delete_store(self) -> bool:
         return await self._client.flushdb()  # pyright: ignore[reportUnknownMemberType, reportAny]
 
-    @override
     async def _close(self) -> None:
         await self._client.aclose()

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

@@ -85,7 +85,6 @@ def __init__(
             client_provided_by_user=client_provided,
         )
 
-    @override
     async def _close(self) -> None:
         self._close_and_flush()