feat: support for key enumeration in the postgres store

src/key_value/aio/stores/postgresql/store.py

@@ -13,7 +13,13 @@
 from typing_extensions import override
 
 from key_value.aio._utils.managed_entry import ManagedEntry, dump_to_json, load_from_json
-from key_value.aio.stores.base import BaseContextManagerStore, BaseDestroyCollectionStore, BaseEnumerateCollectionsStore, BaseStore
+from key_value.aio.stores.base import (
+    BaseContextManagerStore,
+    BaseDestroyCollectionStore,
+    BaseEnumerateCollectionsStore,
+    BaseEnumerateKeysStore,
+    BaseStore,
+)
 
 try:
     import asyncpg
@@ -140,7 +146,9 @@ def _postgresql_row_get_datetime(row: asyncpg.Record, key: str) -> datetime | No
     return row[key]
 
 
-class PostgreSQLStore(BaseEnumerateCollectionsStore, BaseDestroyCollectionStore, BaseContextManagerStore, BaseStore):
+class PostgreSQLStore(
+    BaseEnumerateCollectionsStore, BaseEnumerateKeysStore, BaseDestroyCollectionStore, BaseContextManagerStore, BaseStore
+):
     """PostgreSQL-based key-value store using asyncpg.
 
     This store uses a single shared table with columns for collection, key, value (JSONB), and metadata.
@@ -493,6 +501,32 @@ async def _get_collection_names(self, *, limit: int | None = None) -> list[str]:
 
         return [_postgresql_row_get_str(row, "collection") for row in rows]
 
+    @override
+    async def _get_collection_keys(self, *, collection: str, limit: int | None = None) -> list[str]:
+        """List all keys in a collection.
+
+        Args:
+            collection: The collection to list keys from.
+            limit: Maximum number of keys to return.
+
+        Returns:
+            A list of keys in the collection.
+        """
+        if limit is None or limit <= 0:
+            limit = DEFAULT_PAGE_SIZE
+        limit = min(limit, PAGE_LIMIT)
+
+        pool = self._initialized_pool
+
+        rows = await _postgresql_fetch(
+            pool,
+            f"SELECT key FROM {self._table_name} WHERE collection = $1 ORDER BY key LIMIT $2",
+            collection,
+            limit,
+        )
+
+        return [_postgresql_row_get_str(row, "key") for row in rows]
+
     @override
     async def _delete_collection(self, *, collection: str) -> bool:
         """Delete all entries in a collection.

tests/stores/postgresql/test_postgresql.py

@@ -122,3 +122,42 @@ async def store(self, setup_postgresql: None, postgresql_host: str, postgresql_p
     @pytest.mark.skip(reason="Distributed Caches are unbounded")
     @override
     async def test_not_unbounded(self, store: BaseStore): ...
+
+    async def test_keys_empty(self, store: PostgreSQLStore):
+        """keys() on an empty collection returns an empty list."""
+        assert await store.keys(collection="test_collection") == []
+
+    async def test_keys_after_put(self, store: PostgreSQLStore):
+        """keys() returns keys that were put into the collection."""
+        await store.put(collection="test_collection", key="alpha", value={"v": 1})
+        await store.put(collection="test_collection", key="beta", value={"v": 2})
+        await store.put(collection="test_collection", key="gamma", value={"v": 3})
+
+        # _get_collection_keys sorts by key via ORDER BY.
+        assert await store.keys(collection="test_collection") == ["alpha", "beta", "gamma"]
+
+    async def test_keys_are_collection_scoped(self, store: PostgreSQLStore):
+        """keys() only returns keys from the requested collection."""
+        await store.put(collection="collection_a", key="a_key", value={"v": 1})
+        await store.put(collection="collection_b", key="b_key", value={"v": 2})
+
+        assert await store.keys(collection="collection_a") == ["a_key"]
+        assert await store.keys(collection="collection_b") == ["b_key"]
+
+    async def test_keys_respects_limit(self, store: PostgreSQLStore):
+        """keys() honors the limit parameter."""
+        for i in range(5):
+            await store.put(collection="test_collection", key=f"key_{i}", value={"v": i})
+
+        limited = await store.keys(collection="test_collection", limit=2)
+        assert len(limited) == 2
+        # With ORDER BY key, the first two are key_0 and key_1.
+        assert limited == ["key_0", "key_1"]
+
+    async def test_keys_after_delete(self, store: PostgreSQLStore):
+        """keys() reflects deletes."""
+        await store.put(collection="test_collection", key="keep", value={"v": 1})
+        await store.put(collection="test_collection", key="drop", value={"v": 2})
+        await store.delete(collection="test_collection", key="drop")
+
+        assert await store.keys(collection="test_collection") == ["keep"]