refactor: optimize PostgreSQL store to use one-time table setup

- Move table/index creation from _setup_collection to _setup (called once)
- Remove collection sanitization since collection names are column values
- Remove PostgreSQLV1CollectionSanitizationStrategy class and exports
- Update tests to verify collection names work without restrictions
- Update docstrings to clarify collections are stored as values

This simplifies the implementation since all collections share a single
table, eliminating unnecessary per-collection setup overhead.

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

Author: github-actions[bot]

key-value/key-value-aio/src/key_value/aio/stores/postgresql/__init__.py

@@ -1,9 +1,9 @@
 """PostgreSQL store for py-key-value-aio."""
 
 try:
-    from key_value.aio.stores.postgresql.store import PostgreSQLStore, PostgreSQLV1CollectionSanitizationStrategy
+    from key_value.aio.stores.postgresql.store import PostgreSQLStore
 except ImportError as e:
     msg = 'PostgreSQLStore requires the "postgresql" extra. Install via: pip install "py-key-value-aio[postgresql]"'
     raise ImportError(msg) from e
 
-__all__ = ["PostgreSQLStore", "PostgreSQLV1CollectionSanitizationStrategy"]
+__all__ = ["PostgreSQLStore"]

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

@@ -12,8 +12,6 @@
 from typing import Any, overload
 
 from key_value.shared.utils.managed_entry import ManagedEntry
-from key_value.shared.utils.sanitization import HybridSanitizationStrategy, SanitizationStrategy
-from key_value.shared.utils.sanitize import ALPHANUMERIC_CHARACTERS
 from typing_extensions import Self, override
 
 from key_value.aio.stores.base import BaseContextManagerStore, BaseDestroyCollectionStore, BaseEnumerateCollectionsStore, BaseStore
@@ -34,31 +32,15 @@
 PAGE_LIMIT = 10000
 
 # PostgreSQL table name length limit is 63 characters
-# Use 200 for consistency with MongoDB
-MAX_COLLECTION_LENGTH = 200
 POSTGRES_MAX_IDENTIFIER_LEN = 63
-COLLECTION_ALLOWED_CHARACTERS = ALPHANUMERIC_CHARACTERS + "_"
-
-
-class PostgreSQLV1CollectionSanitizationStrategy(HybridSanitizationStrategy):
-    def __init__(self) -> None:
-        super().__init__(
-            replacement_character="_",
-            max_length=MAX_COLLECTION_LENGTH,
-            allowed_characters=COLLECTION_ALLOWED_CHARACTERS,
-        )
 
 
 class PostgreSQLStore(BaseEnumerateCollectionsStore, BaseDestroyCollectionStore, BaseContextManagerStore, BaseStore):
     """PostgreSQL-based key-value store using asyncpg.
 
-    This store uses a single table with columns for collection, key, value (JSONB), and metadata.
-    Collections are stored as a column value rather than separate tables.
-
-    By default, collections are not sanitized. This means that there are character and length restrictions on
-    collection names that may cause errors when trying to get and put entries.
-
-    To avoid issues, you may want to consider leveraging the `PostgreSQLV1CollectionSanitizationStrategy` strategy.
+    This store uses a single shared table with columns for collection, key, value (JSONB), and metadata.
+    Collections are stored as values in the collection column, not as separate tables or SQL identifiers,
+    so there are no character restrictions on collection names.
 
     Example:
         Basic usage with default connection:
@@ -99,15 +81,13 @@ def __init__(
         pool: asyncpg.Pool,  # type: ignore[type-arg]
         table_name: str | None = None,
         default_collection: str | None = None,
-        collection_sanitization_strategy: SanitizationStrategy | None = None,
     ) -> None:
         """Initialize the PostgreSQL store with an existing connection pool.
 
         Args:
             pool: An existing asyncpg connection pool to use.
             table_name: The name of the table to use for storage (default: kv_store).
             default_collection: The default collection to use if no collection is provided.
-            collection_sanitization_strategy: The sanitization strategy to use for collections.
         """
 
     @overload
@@ -117,15 +97,13 @@ def __init__(
         url: str,
         table_name: str | None = None,
         default_collection: str | None = None,
-        collection_sanitization_strategy: SanitizationStrategy | None = None,
     ) -> None:
         """Initialize the PostgreSQL store with a connection URL.
 
         Args:
             url: PostgreSQL connection URL (e.g., postgresql://user:pass@localhost/dbname).
             table_name: The name of the table to use for storage (default: kv_store).
             default_collection: The default collection to use if no collection is provided.
-            collection_sanitization_strategy: The sanitization strategy to use for collections.
         """
 
     @overload
@@ -139,7 +117,6 @@ def __init__(
         password: str | None = None,
         table_name: str | None = None,
         default_collection: str | None = None,
-        collection_sanitization_strategy: SanitizationStrategy | None = None,
     ) -> None:
         """Initialize the PostgreSQL store with connection parameters.
 
@@ -151,7 +128,6 @@ def __init__(
             password: Database password (default: None).
             table_name: The name of the table to use for storage (default: kv_store).
             default_collection: The default collection to use if no collection is provided.
-            collection_sanitization_strategy: The sanitization strategy to use for collections.
         """
 
     def __init__(
@@ -166,7 +142,6 @@ def __init__(
         password: str | None = None,
         table_name: str | None = None,
         default_collection: str | None = None,
-        collection_sanitization_strategy: SanitizationStrategy | None = None,
     ) -> None:
         """Initialize the PostgreSQL store."""
         self._pool = pool
@@ -178,7 +153,7 @@ def __init__(
         self._user = user
         self._password = password
 
-        # Validate and sanitize table name to prevent SQL injection and invalid identifiers
+        # Validate table name to prevent SQL injection and invalid identifiers
         table_name = table_name or DEFAULT_TABLE
         if not table_name.replace("_", "").isalnum():
             msg = f"Table name must be alphanumeric (with underscores): {table_name}"
@@ -192,10 +167,7 @@ def __init__(
             raise ValueError(msg)
         self._table_name = table_name
 
-        super().__init__(
-            default_collection=default_collection,
-            collection_sanitization_strategy=collection_sanitization_strategy,
-        )
+        super().__init__(default_collection=default_collection)
 
     def _ensure_pool_initialized(self) -> asyncpg.Pool:  # type: ignore[type-arg]
         """Ensure the connection pool is initialized.
@@ -250,14 +222,12 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:  #
             await self._pool.close()
 
     @override
-    async def _setup_collection(self, *, collection: str) -> None:
+    async def _setup(self) -> None:
         """Set up the database table and indexes if they don't exist.
 
-        Args:
-            collection: The collection name (used for validation, but all collections share the same table).
+        This is called once when the store is first used. Since all collections share the same table,
+        we only need to set up the schema once.
         """
-        _ = self._sanitize_collection(collection=collection)
-
         # Create the main table if it doesn't exist
         table_sql = (
             f"CREATE TABLE IF NOT EXISTS {self._table_name} ("
@@ -295,12 +265,10 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
         Returns:
             The managed entry if found and not expired, None otherwise.
         """
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         async with self._acquire_connection() as conn:
             row = await conn.fetchrow(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
                 f"SELECT value, ttl, created_at, expires_at FROM {self._table_name} WHERE collection = $1 AND key = $2",
-                sanitized_collection,
+                collection,
                 key,
             )
 
@@ -318,7 +286,7 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
             if managed_entry.is_expired:
                 await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                     f"DELETE FROM {self._table_name} WHERE collection = $1 AND key = $2",
-                    sanitized_collection,
+                    collection,
                     key,
                 )
                 return None
@@ -339,13 +307,11 @@ async def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) ->
         if not keys:
             return []
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         async with self._acquire_connection() as conn:
             # Use ANY to query for multiple keys
             rows = await conn.fetch(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
                 f"SELECT key, value, ttl, created_at, expires_at FROM {self._table_name} WHERE collection = $1 AND key = ANY($2::text[])",
-                sanitized_collection,
+                collection,
                 list(keys),
             )
 
@@ -370,7 +336,7 @@ async def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) ->
             if expired_keys:
                 await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                     f"DELETE FROM {self._table_name} WHERE collection = $1 AND key = ANY($2::text[])",
-                    sanitized_collection,
+                    collection,
                     expired_keys,
                 )
 
@@ -391,7 +357,6 @@ async def _put_managed_entry(
             collection: The collection to store in.
             managed_entry: The managed entry to store.
         """
-        sanitized_collection = self._sanitize_collection(collection=collection)
 
         async with self._acquire_connection() as conn:
             upsert_sql = (
@@ -403,7 +368,7 @@ async def _put_managed_entry(
             )
             await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                 upsert_sql,
-                sanitized_collection,
+                collection,
                 key,
                 managed_entry.value,
                 managed_entry.ttl,
@@ -435,12 +400,8 @@ async def _put_managed_entries(
         if not keys:
             return
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         # Prepare data for batch insert using method-level ttl/created_at/expires_at
-        values = [
-            (sanitized_collection, key, entry.value, ttl, created_at, expires_at) for key, entry in zip(keys, managed_entries, strict=True)
-        ]
+        values = [(collection, key, entry.value, ttl, created_at, expires_at) for key, entry in zip(keys, managed_entries, strict=True)]
 
         async with self._acquire_connection() as conn:
             # Use executemany for batch insert
@@ -467,12 +428,10 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
         Returns:
             True if the entry was deleted, False if it didn't exist.
         """
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         async with self._acquire_connection() as conn:
             result = await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                 f"DELETE FROM {self._table_name} WHERE collection = $1 AND key = $2",
-                sanitized_collection,
+                collection,
                 key,
             )
             # PostgreSQL execute returns a string like "DELETE N" where N is the number of rows deleted
@@ -492,12 +451,10 @@ async def _delete_managed_entries(self, *, keys: Sequence[str], collection: str)
         if not keys:
             return 0
 
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         async with self._acquire_connection() as conn:
             result = await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                 f"DELETE FROM {self._table_name} WHERE collection = $1 AND key = ANY($2::text[])",
-                sanitized_collection,
+                collection,
                 list(keys),
             )
             # PostgreSQL execute returns a string like "DELETE N" where N is the number of rows deleted
@@ -535,12 +492,10 @@ async def _delete_collection(self, *, collection: str) -> bool:
         Returns:
             True if any entries were deleted, False otherwise.
         """
-        sanitized_collection = self._sanitize_collection(collection=collection)
-
         async with self._acquire_connection() as conn:
             result = await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                 f"DELETE FROM {self._table_name} WHERE collection = $1",
-                sanitized_collection,
+                collection,
             )
             # Return True if any rows were deleted
             return result.split()[-1] != "0"

key-value/key-value-aio/tests/stores/postgresql/test_postgresql.py

@@ -7,7 +7,7 @@
 from typing_extensions import override
 
 from key_value.aio.stores.base import BaseStore
-from key_value.aio.stores.postgresql import PostgreSQLStore, PostgreSQLV1CollectionSanitizationStrategy
+from key_value.aio.stores.postgresql import PostgreSQLStore
 from tests.conftest import docker_container, should_skip_docker_tests
 from tests.stores.base import BaseStoreTests, ContextManagerStoreTestMixin
 
@@ -104,65 +104,25 @@ async def store(self, setup_postgresql: None) -> PostgreSQLStore:
 
         return store
 
-    @pytest.fixture
-    async def postgresql_store(self, store: PostgreSQLStore) -> PostgreSQLStore:
-        """Provide the PostgreSQL store fixture."""
-        return store
-
-    @pytest.fixture
-    async def sanitizing_store(self, setup_postgresql: None) -> PostgreSQLStore:
-        """Create a PostgreSQL store with collection sanitization enabled."""
-        store = PostgreSQLStore(
-            host=POSTGRESQL_HOST,
-            port=POSTGRESQL_HOST_PORT,
-            database=POSTGRESQL_TEST_DB,
-            user=POSTGRESQL_USER,
-            password=POSTGRESQL_PASSWORD,
-            table_name="kv_store_sanitizing",
-            collection_sanitization_strategy=PostgreSQLV1CollectionSanitizationStrategy(),
-        )
-
-        # Clean up the database before each test
-        async with store:
-            if store._pool is not None:  # pyright: ignore[reportPrivateUsage]
-                async with store._pool.acquire() as conn:  # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType]
-                    # Drop and recreate the kv_store_sanitizing table
-                    with contextlib.suppress(Exception):
-                        await conn.execute("DROP TABLE IF EXISTS kv_store_sanitizing")  # pyright: ignore[reportUnknownMemberType]
-
-        return store
-
     @pytest.mark.skip(reason="Distributed Caches are unbounded")
     @override
     async def test_not_unbounded(self, store: BaseStore): ...
 
     @override
-    async def test_long_collection_name(self, store: PostgreSQLStore, sanitizing_store: PostgreSQLStore):  # pyright: ignore[reportIncompatibleMethodOverride]
-        """Test that long collection names fail without sanitization but work with it."""
-        with pytest.raises(Exception):  # noqa: B017, PT011
-            await store.put(collection="test_collection" * 100, key="test_key", value={"test": "test"})
-
-        await sanitizing_store.put(collection="test_collection" * 100, key="test_key", value={"test": "test"})
-        assert await sanitizing_store.get(collection="test_collection" * 100, key="test_key") == {"test": "test"}
+    async def test_long_collection_name(self, store: PostgreSQLStore):  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Test that long collection names work since they're just column values."""
+        # Long collection names should work fine since they're stored as column values, not SQL identifiers
+        long_collection = "test_collection" * 100
+        await store.put(collection=long_collection, key="test_key", value={"test": "test"})
+        assert await store.get(collection=long_collection, key="test_key") == {"test": "test"}
 
     @override
-    async def test_special_characters_in_collection_name(self, store: PostgreSQLStore, sanitizing_store: PostgreSQLStore):  # pyright: ignore[reportIncompatibleMethodOverride]
-        """Test that special characters in collection names fail without sanitization but work with it."""
-        # Without sanitization, special characters should work (PostgreSQL allows them in column values)
-        # but may cause issues with certain characters
-        await store.put(collection="test_collection", key="test_key", value={"test": "test"})
-        assert await store.get(collection="test_collection", key="test_key") == {"test": "test"}
-
-        # With sanitization, special characters should work
-        await sanitizing_store.put(collection="test_collection!@#$%^&*()", key="test_key", value={"test": "test"})
-        assert await sanitizing_store.get(collection="test_collection!@#$%^&*()", key="test_key") == {"test": "test"}
-
-    async def test_postgresql_collection_name_sanitization(self, sanitizing_store: PostgreSQLStore):
-        """Test that the V1 sanitization strategy produces expected collection names."""
-        await sanitizing_store.put(collection="test_collection!@#$%^&*()", key="test_key", value={"test": "test"})
-        assert await sanitizing_store.get(collection="test_collection!@#$%^&*()", key="test_key") == {"test": "test"}
-
-        collections = await sanitizing_store.collections()
-        # The sanitized collection name should only contain alphanumeric characters and underscores
-        assert len(collections) == 1
-        assert all(c.isalnum() or c in "_-" for c in collections[0])
+    async def test_special_characters_in_collection_name(self, store: PostgreSQLStore):  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Test that special characters in collection names work since they're just column values."""
+        # Special characters should work fine since collection names are stored as column values
+        await store.put(collection="test_collection!@#$%^&*()", key="test_key", value={"test": "test"})
+        assert await store.get(collection="test_collection!@#$%^&*()", key="test_key") == {"test": "test"}
+
+        # Verify the collection name is stored as-is
+        collections = await store.collections()
+        assert "test_collection!@#$%^&*()" in collections