refactor: pass ttl and timestamps as parameters to _put_managed_entries

Instead of extracting TTL and timestamps from managed entries, these values
are now passed as explicit parameters to _put_managed_entries. This avoids
the need to inspect managed entries to extract metadata.

Changes:
- Base store: Added ttl, created_at, expires_at parameters to signature
- Base store put_many: Pre-calculates timestamps once for all entries
- Redis store: Uses ttl parameter instead of extracting from first entry
- MongoDB store: Uses timestamp parameters instead of extracting from first entry
- Elasticsearch store: Uses timestamp parameters instead of extracting from first entry
- RocksDB store: Updated signature (stores full JSON, doesn't use params)

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

Author: github-actions[bot]

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

@@ -6,6 +6,7 @@
 from asyncio.locks import Lock
 from collections import defaultdict
 from collections.abc import Mapping, Sequence
+from datetime import datetime, timedelta
 from types import MappingProxyType, TracebackType
 from typing import Any, SupportsFloat
 
@@ -206,9 +207,26 @@ async def _put_managed_entry(self, *, collection: str, key: str, managed_entry:
         """Store a managed entry by key in the specified collection."""
         ...
 
-    async def _put_managed_entries(self, *, collection: str, keys: Sequence[str], managed_entries: Sequence[ManagedEntry]) -> None:
-        """Store multiple managed entries by key in the specified collection."""
+    async def _put_managed_entries(
+        self,
+        *,
+        collection: str,
+        keys: Sequence[str],
+        managed_entries: Sequence[ManagedEntry],
+        ttl: float | None,  # noqa: ARG002
+        created_at: datetime,  # noqa: ARG002
+        expires_at: datetime | None,  # noqa: ARG002
+    ) -> None:
+        """Store multiple managed entries by key in the specified collection.
 
+        Args:
+            collection: The collection to store entries in
+            keys: The keys for the entries
+            managed_entries: The managed entries to store
+            ttl: The TTL in seconds (None for no expiration)
+            created_at: The creation timestamp for all entries
+            expires_at: The expiration timestamp for all entries (None if no TTL)
+        """
         for key, managed_entry in zip(keys, managed_entries, strict=True):
             await self._put_managed_entry(
                 collection=collection,
@@ -261,9 +279,20 @@ async def put_many(
 
         keys, values, ttl_for_entries = self._prepare_put_many(keys=keys, values=values, ttl=ttl)
 
-        managed_entries: list[ManagedEntry] = [ManagedEntry(value=value, ttl=ttl_for_entries, created_at=now()) for value in values]
+        # Pre-calculate timestamps once for all entries
+        created_at = now()
+        expires_at = created_at + timedelta(seconds=ttl_for_entries) if ttl_for_entries is not None else None
+
+        managed_entries: list[ManagedEntry] = [ManagedEntry(value=value, ttl=ttl_for_entries, created_at=created_at) for value in values]
 
-        await self._put_managed_entries(collection=collection, keys=keys, managed_entries=managed_entries)
+        await self._put_managed_entries(
+            collection=collection,
+            keys=keys,
+            managed_entries=managed_entries,
+            ttl=ttl_for_entries,
+            created_at=created_at,
+            expires_at=expires_at,
+        )
 
     @abstractmethod
     async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:

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

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, overload
 
 from key_value.shared.utils.compound import compound_key
@@ -35,8 +36,6 @@
     raise ImportError(msg) from e
 
 if TYPE_CHECKING:
-    from datetime import datetime
-
     from elastic_transport import ObjectApiResponse
 
 DEFAULT_INDEX_PREFIX = "kv_store"
@@ -267,15 +266,22 @@ async def _put_managed_entry(
         )
 
     @override
-    async def _put_managed_entries(self, *, collection: str, keys: Sequence[str], managed_entries: Sequence[ManagedEntry]) -> None:
+    async def _put_managed_entries(
+        self,
+        *,
+        collection: str,
+        keys: Sequence[str],
+        managed_entries: Sequence[ManagedEntry],
+        ttl: float | None,
+        created_at: datetime,
+        expires_at: datetime | None,
+    ) -> None:
         if not keys:
             return
 
-        # All entries in a batch have the same timestamps (from BaseStore.put_many)
-        # Extract timestamps once from the first entry
-        first_entry = managed_entries[0]
-        created_at_iso: str | None = first_entry.created_at.isoformat() if first_entry.created_at else None
-        expires_at_iso: str | None = first_entry.expires_at.isoformat() if first_entry.expires_at else None
+        # Convert timestamps to ISO format for Elasticsearch
+        created_at_iso: str = created_at.isoformat()
+        expires_at_iso: str | None = expires_at.isoformat() if expires_at else None
 
         # Use bulk API for efficient batch indexing
         operations: list[dict[str, Any]] = []
@@ -286,10 +292,9 @@ async def _put_managed_entries(self, *, collection: str, keys: Sequence[str], ma
                 "collection": collection,
                 "key": key,
                 "value": managed_entry.to_json(include_metadata=False),
+                "created_at": created_at_iso,
             }
 
-            if created_at_iso:
-                document["created_at"] = created_at_iso
             if expires_at_iso:
                 document["expires_at"] = expires_at_iso
 

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

@@ -203,17 +203,24 @@ async def _put_managed_entry(
         )
 
     @override
-    async def _put_managed_entries(self, *, collection: str, keys: Sequence[str], managed_entries: Sequence[ManagedEntry]) -> None:
+    async def _put_managed_entries(
+        self,
+        *,
+        collection: str,
+        keys: Sequence[str],
+        managed_entries: Sequence[ManagedEntry],
+        ttl: float | None,
+        created_at: datetime,
+        expires_at: datetime | None,
+    ) -> None:
         if not keys:
             return
 
         collection = self._sanitize_collection_name(collection=collection)
 
-        # All entries in a batch have the same timestamps (from BaseStore.put_many)
-        # Extract timestamps once from the first entry
-        first_entry = managed_entries[0]
-        created_at_iso: str | None = first_entry.created_at.isoformat() if first_entry.created_at else None
-        expires_at_iso: str | None = first_entry.expires_at.isoformat() if first_entry.expires_at else None
+        # Convert timestamps to ISO format for MongoDB
+        created_at_iso: str = created_at.isoformat()
+        expires_at_iso: str | None = expires_at.isoformat() if expires_at else None
         updated_at_iso: str = now().isoformat()
 
         # Use bulk_write for efficient batch operations

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

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+from datetime import datetime
 from typing import Any, overload
 from urllib.parse import urlparse
 
@@ -134,17 +135,24 @@ async def _put_managed_entry(
             _ = await self._client.set(name=combo_key, value=json_value)  # pyright: ignore[reportAny]
 
     @override
-    async def _put_managed_entries(self, *, collection: str, keys: Sequence[str], managed_entries: Sequence[ManagedEntry]) -> None:
+    async def _put_managed_entries(
+        self,
+        *,
+        collection: str,
+        keys: Sequence[str],
+        managed_entries: Sequence[ManagedEntry],
+        ttl: float | None,
+        created_at: datetime,
+        expires_at: datetime | None,
+    ) -> None:
         if not keys:
             return
 
-        # All entries in a batch have the same TTL (from BaseStore.put_many)
-        # Extract TTL once from the first entry
-        first_entry = managed_entries[0]
+        # Convert TTL to integer seconds for Redis
         ttl_seconds: int | None = None
-        if first_entry.ttl is not None:
+        if ttl is not None:
             # Redis does not support <= 0 TTLs
-            ttl_seconds = max(int(first_entry.ttl), 1)
+            ttl_seconds = max(int(ttl), 1)
 
         # Use pipeline for bulk operations
         pipeline = self._client.pipeline()

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

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+from datetime import datetime
 from pathlib import Path
 from typing import Any, overload
 
@@ -128,7 +129,16 @@ async def _put_managed_entry(
         self._db[combo_key] = json_value.encode("utf-8")
 
     @override
-    async def _put_managed_entries(self, *, collection: str, keys: Sequence[str], managed_entries: Sequence[ManagedEntry]) -> None:
+    async def _put_managed_entries(
+        self,
+        *,
+        collection: str,
+        keys: Sequence[str],
+        managed_entries: Sequence[ManagedEntry],
+        ttl: float | None,
+        created_at: datetime,
+        expires_at: datetime | None,
+    ) -> None:
         self._fail_on_closed_store()
 
         if not keys: