perf(storage): use asyncpg connection pool in PostgreSQL kvstore (#5734)

## Summary

- Replace per-operation `asyncpg.connect()`/`close()` with
`asyncpg.create_pool()` to avoid TCP handshake and auth overhead on
every get/set/delete call
- Add retry-on-connection-error that catches
`ConnectionDoesNotExistError`/`OSError`, expires stale pool connections,
and retries once
- Add `pool_size` and `max_overflow` config fields to
`PostgresKVStoreConfig`, matching `PostgresSqlStoreConfig` naming
conventions
- Add integration test that restarts postgres mid-session to verify
recovery
- Add kvstore-related path triggers to the SqlStore CI workflow

## Test plan

- [x] Unit tests pass: `uv run pytest
tests/unit/utils/kvstore/test_postgres_kvstore.py -x --tb=short` (18
passed, 3 xfailed)
- [x] Integration test `test_kvstore_recovers_after_postgres_restart`
passes in CI
- [x] CI green on fork

Signed-off-by: Derek Higgins <derekh@redhat.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: derekhiggins

.github/workflows/integration-sql-store-tests.yml

@@ -13,7 +13,9 @@ on:
       - 'release-[0-9]+.[0-9]+.x'
     paths:
       - 'src/ogx/providers/utils/sqlstore/**'
+      - 'src/ogx/core/storage/kvstore/postgres/**'
       - 'tests/integration/sqlstore/**'
+      - 'tests/integration/providers/utils/kvstore/**'
       - 'uv.lock'
       - 'pyproject.toml'
       - 'requirements.txt'
@@ -70,6 +72,17 @@ jobs:
         run: |
           uv run pytest -sv tests/integration/providers/utils/sqlstore/
 
+      - name: Run KVStore Connection Resilience Test
+        env:
+          ENABLE_POSTGRES_TESTS: "true"
+          POSTGRES_HOST: localhost
+          POSTGRES_PORT: 5432
+          POSTGRES_DB: ogx
+          POSTGRES_USER: ogx
+          POSTGRES_PASSWORD: ogx
+        run: |
+          uv run pytest -sv tests/integration/providers/utils/kvstore/
+
       - name: Upload test logs
         if: ${{ always() }}
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1

src/ogx/core/storage/datatypes.py

@@ -92,6 +92,8 @@ class PostgresKVStoreConfig(CommonConfig):
     ssl_mode: str | None = None
     ca_cert_path: str | None = None
     table_name: str = "ogx_kvstore"
+    pool_size: int = Field(default=5, ge=1, description="Number of persistent connections in the pool")
+    max_overflow: int = Field(default=10, ge=0, description="Max additional connections beyond pool_size")
 
     @classmethod
     def sample_run_config(cls, table_name: str = "ogx_kvstore", **kwargs: object) -> dict[str, str]:

src/ogx/core/storage/kvstore/postgres/postgres.py

@@ -4,7 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
+from collections.abc import Callable, Coroutine
 from datetime import datetime
+from typing import TypeVar
 
 import asyncpg  # type: ignore[import-untyped]
 
@@ -15,12 +18,16 @@
 
 log = get_logger(name=__name__, category="providers::utils")
 
+T = TypeVar("T")
+
 
 class PostgresKVStoreImpl(KVStore):
     """PostgreSQL-backed key-value store implementation."""
 
     def __init__(self, config: PostgresKVStoreConfig):
         self.config = config
+        self._pool: asyncpg.Pool | None = None
+        self._loop: asyncio.AbstractEventLoop | None = None
         self._table_created = False
 
     async def initialize(self) -> None:
@@ -35,22 +42,34 @@ def _build_ssl(self) -> object:
             return self.config.ssl_mode
         return None
 
-    async def _connect(self) -> asyncpg.Connection:
-        try:
-            conn = await asyncpg.connect(
-                host=self.config.host,
-                port=int(self.config.port),
-                database=self.config.db,
-                user=self.config.user,
-                password=self.config.password,
-                ssl=self._build_ssl(),
-            )
-        except Exception as e:
-            log.exception("Could not connect to PostgreSQL database server")
-            raise RuntimeError("Could not connect to PostgreSQL database server") from e
+    async def _acquire(self) -> asyncpg.Pool:
+        loop = asyncio.get_running_loop()
+        if self._pool is not None and self._loop is not loop:
+            # Pool was created in a different event loop (e.g., during init in a
+            # temporary asyncio.run() loop). Discard it -- the old connections are
+            # already dead since that loop is closed.
+            self._pool = None
+            self._table_created = False
 
-        if not self._table_created:
+        if self._pool is None:
             try:
+                self._pool = await asyncpg.create_pool(
+                    host=self.config.host,
+                    port=int(self.config.port),
+                    database=self.config.db,
+                    user=self.config.user,
+                    password=self.config.password,
+                    ssl=self._build_ssl(),
+                    min_size=self.config.pool_size,
+                    max_size=self.config.pool_size + self.config.max_overflow,
+                )
+                self._loop = loop
+            except Exception as e:
+                log.exception("Could not connect to PostgreSQL database server")
+                raise RuntimeError("Could not connect to PostgreSQL database server") from e
+
+        if not self._table_created:
+            async with self._pool.acquire() as conn:
                 await conn.execute(
                     f"""
                     CREATE TABLE IF NOT EXISTS {self.config.table_name} (
@@ -61,10 +80,25 @@ async def _connect(self) -> asyncpg.Connection:
                     """
                 )
                 self._table_created = True
-            except Exception:
-                await conn.close()
-                raise
-        return conn
+
+        return self._pool
+
+    async def _execute_with_retry(self, fn: Callable[[asyncpg.Connection], Coroutine[None, None, T]]) -> T:
+        """Execute fn with a pooled connection, retrying once on connection error."""
+        pool = await self._acquire()
+        try:
+            async with pool.acquire() as conn:
+                return await fn(conn)
+        except (
+            asyncpg.exceptions.ConnectionDoesNotExistError,
+            asyncpg.exceptions.InterfaceError,
+            OSError,
+            RuntimeError,
+        ):
+            log.warning("PostgreSQL connection lost, expiring pool connections")
+            await pool.expire_connections()
+            async with pool.acquire() as conn:
+                return await fn(conn)
 
     def _namespaced_key(self, key: str) -> str:
         if not self.config.namespace:
@@ -78,8 +112,8 @@ def _strip_namespace(self, key: str) -> str:
 
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         key = self._namespaced_key(key)
-        conn = await self._connect()
-        try:
+
+        async def _do(conn: asyncpg.Connection) -> None:
             await conn.execute(
                 f"""
                 INSERT INTO {self.config.table_name} (key, value, expiration)
@@ -91,13 +125,13 @@ async def set(self, key: str, value: str, expiration: datetime | None = None) ->
                 value,
                 expiration,
             )
-        finally:
-            await conn.close()
+
+        await self._execute_with_retry(_do)
 
     async def get(self, key: str) -> str | None:
         key = self._namespaced_key(key)
-        conn = await self._connect()
-        try:
+
+        async def _do(conn: asyncpg.Connection) -> str | None:
             row = await conn.fetchrow(
                 f"""
                 SELECT value FROM {self.config.table_name}
@@ -107,26 +141,25 @@ async def get(self, key: str) -> str | None:
                 key,
             )
             return row["value"] if row else None
-        finally:
-            await conn.close()
+
+        return await self._execute_with_retry(_do)
 
     async def delete(self, key: str) -> None:
         key = self._namespaced_key(key)
-        conn = await self._connect()
-        try:
+
+        async def _do(conn: asyncpg.Connection) -> None:
             await conn.execute(
                 f"DELETE FROM {self.config.table_name} WHERE key = $1",
                 key,
             )
-        finally:
-            await conn.close()
+
+        await self._execute_with_retry(_do)
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
 
-        conn = await self._connect()
-        try:
+        async def _do(conn: asyncpg.Connection) -> list[str]:
             rows = await conn.fetch(
                 f"""
                 SELECT value FROM {self.config.table_name}
@@ -138,15 +171,14 @@ async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
                 end_key,
             )
             return [row["value"] for row in rows]
-        finally:
-            await conn.close()
+
+        return await self._execute_with_retry(_do)
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
 
-        conn = await self._connect()
-        try:
+        async def _do(conn: asyncpg.Connection) -> list[str]:
             rows = await conn.fetch(
                 f"""
                 SELECT key FROM {self.config.table_name}
@@ -158,8 +190,10 @@ async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
                 end_key,
             )
             return [self._strip_namespace(row["key"]) for row in rows]
-        finally:
-            await conn.close()
+
+        return await self._execute_with_retry(_do)
 
     async def shutdown(self) -> None:
-        pass
+        if self._pool:
+            await self._pool.close()
+            self._pool = None

tests/integration/providers/utils/kvstore/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

tests/integration/providers/utils/kvstore/test_postgres_connection_resilience.py

@@ -0,0 +1,99 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import subprocess
+import time
+
+import pytest
+
+from ogx.core.storage.kvstore.config import PostgresKVStoreConfig
+from ogx.core.storage.kvstore.postgres.postgres import PostgresKVStoreImpl
+
+POSTGRES_ENABLED = os.environ.get("ENABLE_POSTGRES_TESTS", "").lower() == "true"
+
+pytestmark = pytest.mark.skipif(
+    not POSTGRES_ENABLED,
+    reason="PostgreSQL tests disabled (set ENABLE_POSTGRES_TESTS=true)",
+)
+
+
+def get_postgres_config():
+    return PostgresKVStoreConfig(
+        host=os.environ.get("POSTGRES_HOST", "localhost"),
+        port=int(os.environ.get("POSTGRES_PORT", "5432")),
+        db=os.environ.get("POSTGRES_DB", "ogx"),
+        user=os.environ.get("POSTGRES_USER", "ogx"),
+        password=os.environ.get("POSTGRES_PASSWORD", "ogx"),
+        table_name="test_kvstore_resilience",
+    )
+
+
+def _find_postgres_container():
+    """Find the running postgres container ID."""
+    result = subprocess.run(
+        ["docker", "ps", "-q", "--filter", "ancestor=postgres:15"],
+        capture_output=True,
+        text=True,
+    )
+    container_id = result.stdout.strip()
+    if not container_id:
+        # Fallback: match any postgres container
+        result = subprocess.run(
+            ["docker", "ps", "-q", "--filter", "name=postgres"],
+            capture_output=True,
+            text=True,
+        )
+        container_id = result.stdout.strip()
+    if not container_id:
+        pytest.skip("No postgres container found to restart")
+    # Take first if multiple
+    return container_id.split("\n")[0]
+
+
+def _restart_postgres(container_id: str, timeout: int = 30):
+    """Restart the postgres container and wait for it to accept connections."""
+    subprocess.run(["docker", "restart", container_id], check=True)
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        result = subprocess.run(
+            [
+                "docker",
+                "exec",
+                container_id,
+                "pg_isready",
+                "-U",
+                "ogx",
+            ],
+            capture_output=True,
+        )
+        if result.returncode == 0:
+            return
+        time.sleep(1)
+    raise TimeoutError("PostgreSQL did not become ready after restart")
+
+
+async def test_kvstore_recovers_after_postgres_restart():
+    config = get_postgres_config()
+    store = PostgresKVStoreImpl(config)
+    await store.initialize()
+
+    # Baseline: kvstore works
+    await store.set("test_key", "before_restart")
+    value = await store.get("test_key")
+    assert value == "before_restart"
+
+    # Restart postgres
+    container_id = _find_postgres_container()
+    _restart_postgres(container_id)
+
+    # After restart, kvstore should still work.
+    # Without reconnect logic this will raise InterfaceError.
+    await store.set("test_key", "after_restart")
+    value = await store.get("test_key")
+    assert value == "after_restart"
+
+    await store.shutdown()