Refactor PostgreSQL store to use helper methods for pool management

- Add _ensure_pool_initialized() helper to centralize pool initialization checks
- Add _acquire_connection() async context manager for acquiring connections
- Remove 7 duplicate pool initialization checks across all methods
- Fix linting issues in test file (S105, TRY300)
- Add pyright type ignores for asyncpg's incomplete type stubs

This refactoring improves code maintainability by following DRY principles
and centralizing pool management logic.

Co-authored-by: William Easton <strawgate@users.noreply.github.com>

Author: github-actions[bot]

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

@@ -4,7 +4,8 @@
 This is safe because table names are validated in __init__ to be alphanumeric (plus underscores).
 """
 
-from collections.abc import Sequence
+from collections.abc import AsyncIterator, Sequence
+from contextlib import asynccontextmanager
 from datetime import datetime
 from typing import Any, overload
 
@@ -62,7 +63,7 @@ class PostgreSQLStore(BaseEnumerateCollectionsStore, BaseDestroyCollectionStore,
         ... )
     """
 
-    _pool: asyncpg.Pool[asyncpg.Record] | None
+    _pool: asyncpg.Pool | None  # type: ignore[type-arg]
     _url: str | None
     _host: str
     _port: int
@@ -75,7 +76,7 @@ class PostgreSQLStore(BaseEnumerateCollectionsStore, BaseDestroyCollectionStore,
     def __init__(
         self,
         *,
-        pool: asyncpg.Pool[asyncpg.Record],
+        pool: asyncpg.Pool,  # type: ignore[type-arg]
         table_name: str | None = None,
         default_collection: str | None = None,
     ) -> None:
@@ -130,7 +131,7 @@ def __init__(
     def __init__(
         self,
         *,
-        pool: asyncpg.Pool[asyncpg.Record] | None = None,
+        pool: asyncpg.Pool | None = None,  # type: ignore[type-arg]
         url: str | None = None,
         host: str = DEFAULT_HOST,
         port: int = DEFAULT_PORT,
@@ -158,13 +159,41 @@ def __init__(
 
         super().__init__(default_collection=default_collection)
 
+    def _ensure_pool_initialized(self) -> asyncpg.Pool:  # type: ignore[type-arg]
+        """Ensure the connection pool is initialized.
+
+        Returns:
+            The initialized connection pool.
+
+        Raises:
+            RuntimeError: If the pool is not initialized.
+        """
+        if self._pool is None:
+            msg = "Pool is not initialized. Use async with or call __aenter__() first."
+            raise RuntimeError(msg)
+        return self._pool
+
+    @asynccontextmanager
+    async def _acquire_connection(self) -> AsyncIterator[asyncpg.Connection]:  # type: ignore[type-arg]
+        """Acquire a connection from the pool.
+
+        Yields:
+            A connection from the pool.
+
+        Raises:
+            RuntimeError: If the pool is not initialized.
+        """
+        pool = self._ensure_pool_initialized()
+        async with pool.acquire() as conn:  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            yield conn
+
     @override
     async def __aenter__(self) -> Self:
         if self._pool is None:
             if self._url:
-                self._pool = await asyncpg.create_pool(self._url)
+                self._pool = await asyncpg.create_pool(self._url)  # pyright: ignore[reportUnknownMemberType]
             else:
-                self._pool = await asyncpg.create_pool(
+                self._pool = await asyncpg.create_pool(  # pyright: ignore[reportUnknownMemberType]
                     host=self._host,
                     port=self._port,
                     database=self._database,
@@ -201,10 +230,6 @@ async def _setup_collection(self, *, collection: str) -> None:
         """
         _ = self._sanitize_collection_name(collection=collection)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
         # Create the main table if it doesn't exist
         create_table_sql = f"""  # noqa: S608
         CREATE TABLE IF NOT EXISTS {self._table_name} (
@@ -225,9 +250,9 @@ async def _setup_collection(self, *, collection: str) -> None:
         WHERE expires_at IS NOT NULL
         """
 
-        async with self._pool.acquire() as conn:
-            await conn.execute(create_table_sql)
-            await conn.execute(create_index_sql)
+        async with self._acquire_connection() as conn:
+            await conn.execute(create_table_sql)  # pyright: ignore[reportUnknownMemberType]
+            await conn.execute(create_index_sql)  # pyright: ignore[reportUnknownMemberType]
 
     @override
     async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:
@@ -242,12 +267,8 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
         """
         sanitized_collection = self._sanitize_collection_name(collection=collection)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
-        async with self._pool.acquire() as conn:
-            row = await conn.fetchrow(
+        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",  # noqa: S608
                 sanitized_collection,
                 key,
@@ -258,15 +279,15 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
 
             # Parse the managed entry
             managed_entry = ManagedEntry(
-                value=row["value"],
-                ttl=row["ttl"],
-                created_at=row["created_at"],
-                expires_at=row["expires_at"],
+                value=row["value"],  # pyright: ignore[reportUnknownArgumentType]
+                ttl=row["ttl"],  # pyright: ignore[reportUnknownArgumentType]
+                created_at=row["created_at"],  # pyright: ignore[reportUnknownArgumentType]
+                expires_at=row["expires_at"],  # pyright: ignore[reportUnknownArgumentType]
             )
 
             # Check if expired and delete if so
             if managed_entry.is_expired:
-                await conn.execute(
+                await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                     f"DELETE FROM {self._table_name} WHERE collection = $1 AND key = $2",  # noqa: S608
                     sanitized_collection,
                     key,
@@ -291,13 +312,9 @@ async def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) ->
 
         sanitized_collection = self._sanitize_collection_name(collection=collection)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
-        async with self._pool.acquire() as conn:
+        async with self._acquire_connection() as conn:
             # Use ANY to query for multiple keys
-            rows = await conn.fetch(
+            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[])",  # noqa: S608
                 sanitized_collection,
                 list(keys),
@@ -307,23 +324,23 @@ async def _get_managed_entries(self, *, collection: str, keys: Sequence[str]) ->
             entries_by_key: dict[str, ManagedEntry | None] = dict.fromkeys(keys)
             expired_keys: list[str] = []
 
-            for row in rows:
+            for row in rows:  # pyright: ignore[reportUnknownVariableType]
                 managed_entry = ManagedEntry(
-                    value=row["value"],
-                    ttl=row["ttl"],
-                    created_at=row["created_at"],
-                    expires_at=row["expires_at"],
+                    value=row["value"],  # pyright: ignore[reportUnknownArgumentType]
+                    ttl=row["ttl"],  # pyright: ignore[reportUnknownArgumentType]
+                    created_at=row["created_at"],  # pyright: ignore[reportUnknownArgumentType]
+                    expires_at=row["expires_at"],  # pyright: ignore[reportUnknownArgumentType]
                 )
 
                 if managed_entry.is_expired:
-                    expired_keys.append(row["key"])
+                    expired_keys.append(row["key"])  # pyright: ignore[reportUnknownArgumentType]
                     entries_by_key[row["key"]] = None
                 else:
                     entries_by_key[row["key"]] = managed_entry
 
             # Delete expired entries in batch
             if expired_keys:
-                await conn.execute(
+                await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                     f"DELETE FROM {self._table_name} WHERE collection = $1 AND key = ANY($2::text[])",  # noqa: S608
                     sanitized_collection,
                     expired_keys,
@@ -348,12 +365,8 @@ async def _put_managed_entry(
         """
         sanitized_collection = self._sanitize_collection_name(collection=collection)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
-        async with self._pool.acquire() as conn:
-            await conn.execute(
+        async with self._acquire_connection() as conn:
+            await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                 f"""
                 INSERT INTO {self._table_name} (collection, key, value, ttl, created_at, expires_at)
                 VALUES ($1, $2, $3, $4, $5, $6)
@@ -398,19 +411,15 @@ async def _put_managed_entries(
 
         sanitized_collection = self._sanitize_collection_name(collection=collection)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
         # Prepare data for batch insert
         values = [
             (sanitized_collection, key, entry.value, entry.ttl, entry.created_at, entry.expires_at)
             for key, entry in zip(keys, managed_entries, strict=True)
         ]
 
-        async with self._pool.acquire() as conn:
+        async with self._acquire_connection() as conn:
             # Use executemany for batch insert
-            await conn.executemany(
+            await conn.executemany(  # pyright: ignore[reportUnknownMemberType]
                 f"""
                 INSERT INTO {self._table_name} (collection, key, value, ttl, created_at, expires_at)
                 VALUES ($1, $2, $3, $4, $5, $6)
@@ -437,12 +446,8 @@ async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:
         """
         sanitized_collection = self._sanitize_collection_name(collection=collection)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
-        async with self._pool.acquire() as conn:
-            result = await conn.execute(
+        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",  # noqa: S608
                 sanitized_collection,
                 key,
@@ -466,12 +471,8 @@ async def _delete_managed_entries(self, *, keys: Sequence[str], collection: str)
 
         sanitized_collection = self._sanitize_collection_name(collection=collection)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
-        async with self._pool.acquire() as conn:
-            result = await conn.execute(
+        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[])",  # noqa: S608
                 sanitized_collection,
                 list(keys),
@@ -491,17 +492,13 @@ async def _get_collection_names(self, *, limit: int | None = None) -> list[str]:
         """
         limit = min(limit or DEFAULT_PAGE_SIZE, PAGE_LIMIT)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
-        async with self._pool.acquire() as conn:
-            rows = await conn.fetch(
+        async with self._acquire_connection() as conn:
+            rows = await conn.fetch(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
                 f"SELECT DISTINCT collection FROM {self._table_name} ORDER BY collection LIMIT $1",  # noqa: S608
                 limit,
             )
 
-            return [row["collection"] for row in rows]
+            return [row["collection"] for row in rows]  # pyright: ignore[reportUnknownVariableType]
 
     @override
     async def _delete_collection(self, *, collection: str) -> bool:
@@ -515,12 +512,8 @@ async def _delete_collection(self, *, collection: str) -> bool:
         """
         sanitized_collection = self._sanitize_collection_name(collection=collection)
 
-        if self._pool is None:
-            msg = "Pool is not initialized. Use async with or call __aenter__() first."
-            raise RuntimeError(msg)
-
-        async with self._pool.acquire() as conn:
-            result = await conn.execute(
+        async with self._acquire_connection() as conn:
+            result = await conn.execute(  # pyright: ignore[reportUnknownMemberType]
                 f"DELETE FROM {self._table_name} WHERE collection = $1",  # noqa: S608
                 sanitized_collection,
             )

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

@@ -20,7 +20,7 @@
 POSTGRESQL_HOST = "localhost"
 POSTGRESQL_HOST_PORT = 5432
 POSTGRESQL_USER = "postgres"
-POSTGRESQL_PASSWORD = "test"
+POSTGRESQL_PASSWORD = "test"  # noqa: S105
 POSTGRESQL_TEST_DB = "kv_store_test"
 
 WAIT_FOR_POSTGRESQL_TIMEOUT = 30
@@ -37,17 +37,18 @@ async def ping_postgresql() -> bool:
         return False
 
     try:
-        conn = await asyncpg.connect(
+        conn = await asyncpg.connect(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
             host=POSTGRESQL_HOST,
             port=POSTGRESQL_HOST_PORT,
             user=POSTGRESQL_USER,
             password=POSTGRESQL_PASSWORD,
             database="postgres",
         )
-        await conn.close()
-        return True
+        await conn.close()  # pyright: ignore[reportUnknownMemberType]
     except Exception:
         return False
+    else:
+        return True
 
 
 class PostgreSQLFailedToStartError(Exception):
@@ -96,10 +97,10 @@ async def store(self, setup_postgresql: None) -> PostgreSQLStore:
         # 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]
+                async with store._pool.acquire() as conn:  # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType]
                     # Drop and recreate the kv_store table
                     with contextlib.suppress(Exception):
-                        await conn.execute("DROP TABLE IF EXISTS kv_store")
+                        await conn.execute("DROP TABLE IF EXISTS kv_store")  # pyright: ignore[reportUnknownMemberType]
 
         return store