diff --git a/src/key_value/aio/stores/postgresql/store.py b/src/key_value/aio/stores/postgresql/store.py index 97da2839..a58c4bf8 100644 --- a/src/key_value/aio/stores/postgresql/store.py +++ b/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. diff --git a/tests/stores/postgresql/test_postgresql.py b/tests/stores/postgresql/test_postgresql.py index 9048244a..5f8a2e0e 100644 --- a/tests/stores/postgresql/test_postgresql.py +++ b/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"]