fix: address CodeRabbit PR feedback on lifecycle management

- MongoDB: Remove duplicate client cleanup from __aexit__, update _close() to call __aexit__(None, None, None) for symmetry with __aenter__
- DynamoDB: Move client cleanup before super().__aexit__() to prevent double-cleanup, set _client = None after closing
- RocksDB: Add try-except in __del__ for best-effort cleanup with appropriate noqa comments
- DuckDB: Use getattr for _is_closed check in __del__ for consistency
- Disk: Set _client_provided_by_user early for exception safety
- Memcached: Change truthiness check to explicit 'is not None' for consistency
- Valkey: Fix _close() to close the correct client (_connected_client) and set it to None for idempotency

Co-authored-by: William Easton <[email protected]>

Author: github-actions[bot]

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

@@ -66,6 +66,7 @@ def __init__(
             raise ValueError(msg)
 
         client_provided = disk_cache is not None
+        self._client_provided_by_user = client_provided
 
         if disk_cache:
             self._cache = disk_cache

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

@@ -372,7 +372,11 @@ async def _close(self) -> None:
     def __del__(self) -> None:
         """Clean up the DuckDB connection on deletion."""
         try:
-            if not self._is_closed and not getattr(self, "_client_provided_by_user", False) and hasattr(self, "_connection"):
+            if (
+                not getattr(self, "_is_closed", False)
+                and not getattr(self, "_client_provided_by_user", False)
+                and hasattr(self, "_connection")
+            ):
                 self._connection.close()
                 self._is_closed = True
         except Exception:  # noqa: S110

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

@@ -135,10 +135,12 @@ async def __aenter__(self) -> Self:
     async def __aexit__(
         self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
     ) -> None:
-        await super().__aexit__(exc_type, exc_value, traceback)
-        # Only exit the client's context manager if the store created it
+        # Exit the client's context manager before calling super().__aexit__() to avoid double-cleanup
         if not self._client_provided_by_user and self._client:
-            await self._client.__aexit__(exc_type, exc_value, traceback)
+            client = self._client
+            self._client = None
+            await client.__aexit__(exc_type, exc_value, traceback)
+        await super().__aexit__(exc_type, exc_value, traceback)
 
     @property
     def _connected_client(self) -> DynamoDBClient:
@@ -266,5 +268,4 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
     @override
     async def _close(self) -> None:
         """Close the DynamoDB client."""
-        if self._client:
-            await self._client.__aexit__(None, None, None)  # pyright: ignore[reportUnknownMemberType]
+        # Client cleanup is handled in __aexit__

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

@@ -70,7 +70,7 @@ def __init__(
         """
         client_provided = client is not None
 
-        if client:
+        if client is not None:
             self._client = client
         else:
             self._client = Client(host=host, port=port)

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

@@ -203,9 +203,6 @@ async def __aenter__(self) -> Self:
     @override
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:  # pyright: ignore[reportAny]
         await super().__aexit__(exc_type, exc_val, exc_tb)
-        # Only exit the client's context manager if the store created it
-        if not self._client_provided_by_user:
-            await self._client.__aexit__(exc_type, exc_val, exc_tb)
 
     @override
     async def _setup_collection(self, *, collection: str) -> None:
@@ -345,4 +342,4 @@ async def _delete_collection(self, *, collection: str) -> bool:
 
     @override
     async def _close(self) -> None:
-        await self._client.close()
+        await self._client.__aexit__(None, None, None)

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

@@ -194,4 +194,7 @@ async def _delete_managed_entries(self, *, keys: Sequence[str], collection: str)
 
     def __del__(self) -> None:
         if not getattr(self, "_client_provided_by_user", False):
-            self._close_and_flush()
+            try:  # noqa: SIM105
+                self._close_and_flush()
+            except (AttributeError, Exception):  # noqa: S110
+                pass  # Best-effort cleanup during finalization

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

@@ -168,4 +168,5 @@ async def _delete_managed_entries(self, *, keys: Sequence[str], collection: str)
     @override
     async def _close(self) -> None:
         if self._connected_client is not None:
-            await self._client.close()
+            await self._connected_client.close()
+            self._connected_client = None

key-value/key-value-sync/src/key_value/sync/code_gen/stores/disk/store.py

@@ -69,6 +69,7 @@ def __init__(
             raise ValueError(msg)
 
         client_provided = disk_cache is not None
+        self._client_provided_by_user = client_provided
 
         if disk_cache:
             self._cache = disk_cache

key-value/key-value-sync/src/key_value/sync/code_gen/stores/duckdb/store.py

@@ -317,7 +317,11 @@ def _close(self) -> None:
     def __del__(self) -> None:
         """Clean up the DuckDB connection on deletion."""
         try:
-            if not self._is_closed and (not getattr(self, "_client_provided_by_user", False)) and hasattr(self, "_connection"):
+            if (
+                not getattr(self, "_is_closed", False)
+                and (not getattr(self, "_client_provided_by_user", False))
+                and hasattr(self, "_connection")
+            ):
                 self._connection.close()
                 self._is_closed = True
         except Exception:  # noqa: S110

key-value/key-value-sync/src/key_value/sync/code_gen/stores/mongodb/store.py

@@ -201,9 +201,6 @@ def __enter__(self) -> Self:
     @override
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:  # pyright: ignore[reportAny]
         super().__exit__(exc_type, exc_val, exc_tb)
-        # Only exit the client's context manager if the store created it
-        if not self._client_provided_by_user:
-            self._client.__exit__(exc_type, exc_val, exc_tb)
 
     @override
     def _setup_collection(self, *, collection: str) -> None:
@@ -320,4 +317,4 @@ def _delete_collection(self, *, collection: str) -> bool:
 
     @override
     def _close(self) -> None:
-        self._client.close()
+        self._client.__exit__(None, None, None)