refactor: simplify client ownership by initializing flag in base class

- Initialize _client_provided_by_user = False in BaseContextManagerStore.__init__()
- Remove redundant = False assignments from all store implementations
- Add ownership checks in _close() methods to prevent closing user-provided clients
- Add ownership checks in __del__() methods where applicable (RocksDB, Disk)
- Fixes test failures where stores were closing user-provided clients

This simplifies the ownership pattern: stores only set the flag to True
when users provide a client, relying on the base class default of False
for internally-created clients.

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

Author: github-actions[bot]

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

@@ -434,6 +434,11 @@ class BaseContextManagerStore(BaseStore, ABC):
 
     _client_provided_by_user: bool
 
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the context manager store with default client ownership."""
+        self._client_provided_by_user = False
+        super().__init__(**kwargs)
+
     async def __aenter__(self) -> Self:
         await self.setup()
         return self

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

@@ -77,7 +77,6 @@ def __init__(
                 self._cache = Cache(directory=directory, size_limit=max_size)
             else:
                 self._cache = Cache(directory=directory, eviction_policy="none")
-            self._client_provided_by_user = False
 
         self._stable_api = True
 
@@ -125,7 +124,9 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
 
     @override
     async def _close(self) -> None:
-        self._cache.close()
+        if not self._client_provided_by_user:
+            self._cache.close()
 
     def __del__(self) -> None:
-        self._cache.close()
+        if not self._client_provided_by_user:
+            self._cache.close()

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

@@ -154,7 +154,6 @@ def __init__(
                 self._connection = duckdb.connect(":memory:")
             else:
                 self._connection = duckdb.connect(database=database_path)
-            self._client_provided_by_user = False
 
         self._is_closed = False
         self._adapter = DuckDBSerializationAdapter()

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

@@ -116,7 +116,6 @@ def __init__(
             self._raw_client = session.client(service_name="dynamodb", endpoint_url=endpoint_url)  # pyright: ignore[reportUnknownMemberType]
 
             self._client = None
-            self._client_provided_by_user = False
 
         super().__init__(default_collection=default_collection)
 
@@ -263,5 +262,5 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
     @override
     async def _close(self) -> None:
         """Close the DynamoDB client."""
-        if self._client:
+        if self._client and not self._client_provided_by_user:
             await self._client.__aexit__(None, None, None)  # pyright: ignore[reportUnknownMemberType]

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

@@ -236,7 +236,6 @@ def __init__(
             self._client = AsyncElasticsearch(
                 hosts=[url], api_key=api_key, http_compress=True, request_timeout=10, retry_on_timeout=True, max_retries=3
             )
-            self._client_provided_by_user = False
         else:
             msg = "Either elasticsearch_client or url must be provided"
             raise ValueError(msg)
@@ -553,4 +552,5 @@ async def _cull(self) -> None:
 
     @override
     async def _close(self) -> None:
-        await self._client.close()
+        if not self._client_provided_by_user:
+            await self._client.close()

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

@@ -73,7 +73,6 @@ def __init__(
             self._client_provided_by_user = True
         else:
             self._client = Client(host=host, port=port)
-            self._client_provided_by_user = False
 
         super().__init__(
             default_collection=default_collection,
@@ -153,4 +152,5 @@ async def _delete_store(self) -> bool:
 
     @override
     async def _close(self) -> None:
-        await self._client.close()
+        if not self._client_provided_by_user:
+            await self._client.close()

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

@@ -174,11 +174,9 @@ def __init__(
             self._client_provided_by_user = True
         elif url:
             self._client = AsyncMongoClient(url)
-            self._client_provided_by_user = False
         else:
             # Defaults to localhost
             self._client = AsyncMongoClient()
-            self._client_provided_by_user = False
 
         db_name = db_name or DEFAULT_DB
         coll_name = coll_name or DEFAULT_COLLECTION
@@ -345,4 +343,5 @@ async def _delete_collection(self, *, collection: str) -> bool:
 
     @override
     async def _close(self) -> None:
-        await self._client.close()
+        if not self._client_provided_by_user:
+            await self._client.close()

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

@@ -76,7 +76,6 @@ def __init__(
                 password=parsed_url.password or password,
                 decode_responses=True,
             )
-            self._client_provided_by_user = False
         else:
             self._client = Redis(
                 host=host,
@@ -85,7 +84,6 @@ def __init__(
                 password=password,
                 decode_responses=True,
             )
-            self._client_provided_by_user = False
 
         self._stable_api = True
         self._adapter = BasicSerializationAdapter(date_format="isoformat", value_format="dict")
@@ -222,4 +220,5 @@ async def _delete_store(self) -> bool:
 
     @override
     async def _close(self) -> None:
-        await self._client.aclose()
+        if not self._client_provided_by_user:
+            await self._client.aclose()

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

@@ -76,15 +76,15 @@ def __init__(
             opts.create_if_missing(True)
 
             self._db = Rdict(str(path), options=opts)
-            self._client_provided_by_user = False
 
         self._is_closed = False
 
         super().__init__(default_collection=default_collection)
 
     @override
     async def _close(self) -> None:
-        self._close_and_flush()
+        if not self._client_provided_by_user:
+            self._close_and_flush()
 
     def _close_and_flush(self) -> None:
         if not self._is_closed:
@@ -190,4 +190,5 @@ async def _delete_managed_entries(self, *, keys: Sequence[str], collection: str)
         return deleted_count
 
     def __del__(self) -> None:
-        self._close_and_flush()
+        if not self._client_provided_by_user:
+            self._close_and_flush()

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

@@ -74,7 +74,6 @@ def __init__(
             credentials: ServerCredentials | None = ServerCredentials(password=password, username=username) if password else None
             self._client_config = GlideClientConfiguration(addresses=addresses, database_id=db, credentials=credentials)
             self._connected_client = None
-            self._client_provided_by_user = False
 
         self._stable_api = True
 
@@ -165,6 +164,5 @@ async def _delete_managed_entries(self, *, keys: Sequence[str], collection: str)
 
     @override
     async def _close(self) -> None:
-        if self._connected_client is None:
-            return
-        await self._client.close()
+        if self._connected_client is not None and not self._client_provided_by_user:
+            await self._client.close()