fix(db): keep the DSN's application_name across pool reuse (blank pg_stat_activity behind pgbouncer)

hindsight-api-slim/hindsight_api/engine/db/postgresql.py

@@ -8,9 +8,10 @@
 """
 
 import logging
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import asynccontextmanager
 from typing import Any
+from urllib.parse import parse_qs, urlparse
 
 import asyncpg  # noqa: F401
 
@@ -64,6 +65,54 @@ async def copy_records_to_table(
         await self._conn.copy_records_to_table(table_name, records=records, columns=columns, timeout=timeout)
 
 
+def application_name_from_dsn(dsn: str) -> str | None:
+    """Extract the ``application_name`` query parameter from a PostgreSQL DSN.
+
+    asyncpg already forwards this to the server in the startup packet (it
+    passes unrecognized DSN query parameters through as ``server_settings``),
+    so a direct connection is labelled correctly in ``pg_stat_activity``.
+    The value is extracted here so it can be re-applied per acquire — see
+    ``_application_name_setup``.
+    """
+    try:
+        values = parse_qs(urlparse(dsn).query).get("application_name")
+    except ValueError:
+        return None
+    if not values:
+        return None
+    # libpq semantics: the last occurrence of a repeated parameter wins.
+    return values[-1] or None
+
+
+def _application_name_setup(app_name: str, init_callback: Any | None) -> Callable[[Any], Awaitable[None]]:
+    """Wrap ``init_callback`` so every acquire re-asserts ``application_name``.
+
+    asyncpg runs ``RESET ALL`` when a connection is released back to the pool.
+    Connected straight to PostgreSQL that is harmless: ``RESET ALL`` restores
+    the value from the startup packet, which carried the DSN's name.
+
+    Behind a connection pooler (pgbouncer) it is not. The server connection's
+    startup packet is the *pooler's*, with no application_name; pgbouncer
+    applies the client's value with a ``SET`` when it links client to server.
+    ``RESET ALL`` therefore resets it to empty, and pgbouncer — which already
+    believes the value is applied — does not re-issue it. Only the first
+    acquire on each server connection is attributed; every later one reports
+    an empty application_name, which is exactly the sort of gap that shows up
+    in production but never under psql.
+
+    Re-asserting it on every acquire fixes both topologies. ``set_config``
+    rather than ``SET`` because the name is operator-supplied and ``SET`` does
+    not accept bind parameters.
+    """
+
+    async def _setup(conn: Any) -> None:
+        await conn.execute("SELECT set_config('application_name', $1, false)", app_name)
+        if init_callback is not None:
+            await init_callback(conn)
+
+    return _setup
+
+
 class PostgreSQLBackend(DatabaseBackend):
     """DatabaseBackend implementation wrapping an asyncpg connection pool."""
 
@@ -101,6 +150,11 @@ async def initialize(
         # the wait it names: a pool-exhaustion stall never surfaced as an error,
         # it just hung (#3002). 0 restores the unbounded behaviour.
         self._acquire_timeout_s = acquire_timeout if acquire_timeout > 0 else None
+        # The DSN's application_name survives RESET ALL only on a direct
+        # connection; behind pgbouncer it has to be re-asserted per acquire
+        # (see _application_name_setup).
+        app_name = application_name_from_dsn(dsn)
+        pool_setup = _application_name_setup(app_name, init_callback) if app_name else init_callback
         self._pool = await asyncpg.create_pool(
             dsn,
             min_size=min_size,
@@ -109,11 +163,11 @@ async def initialize(
             statement_cache_size=statement_cache_size,
             timeout=acquire_timeout,
             # init runs once per new connection; setup runs on every acquire,
-            # after asyncpg's release-time RESET ALL. Passing init_callback as
+            # after asyncpg's release-time RESET ALL. Passing the callback as
             # both keeps the per-connection session GUCs (hnsw.ef_search, etc.)
             # applied after a connection is reused, not just on first creation.
-            init=init_callback,
-            setup=init_callback,
+            init=pool_setup,
+            setup=pool_setup,
         )
         logger.info(
             f"PostgreSQL pool created (min={min_size}, max={max_size}, "

hindsight-api-slim/tests/test_dsn_application_name.py

@@ -0,0 +1,137 @@
+"""The DSN's ``application_name`` must survive connection reuse, not just the
+first acquire.
+
+asyncpg forwards the DSN's ``application_name`` in the startup packet, so a
+direct connection is labelled correctly. But asyncpg also runs ``RESET ALL``
+when a connection is released to the pool, and behind pgbouncer the server
+connection's startup packet is the *pooler's* — pgbouncer applies the client's
+name with a ``SET`` at link time, so ``RESET ALL`` clears it and pgbouncer does
+not re-issue it. Only the first acquire on each server connection ends up
+attributed in ``pg_stat_activity``; the rest report an empty name. The backend
+therefore re-asserts the name on every acquire via the pool's setup hook.
+
+Deterministic (no DB): asyncpg.create_pool is monkeypatched to capture kwargs.
+"""
+
+import pytest
+
+from hindsight_api.engine.db import postgresql as pg_mod
+from hindsight_api.engine.db.postgresql import PostgreSQLBackend, application_name_from_dsn
+
+
+class TestApplicationNameFromDsn:
+    def test_present(self):
+        assert application_name_from_dsn("postgresql://u:p@h:5432/db?application_name=worker-0") == "worker-0"
+
+    def test_absent(self):
+        assert application_name_from_dsn("postgresql://u:p@h:5432/db") is None
+
+    def test_other_params_only(self):
+        assert application_name_from_dsn("postgresql://u:p@h:5432/db?sslmode=disable") is None
+
+    def test_alongside_other_params(self):
+        dsn = "postgresql://u:p@h:5432/db?application_name=api-1&sslmode=disable"
+        assert application_name_from_dsn(dsn) == "api-1"
+
+    def test_last_occurrence_wins(self):
+        # libpq semantics for repeated URL parameters.
+        dsn = "postgresql://u:p@h/db?application_name=a&application_name=b"
+        assert application_name_from_dsn(dsn) == "b"
+
+    def test_empty_value_is_none(self):
+        assert application_name_from_dsn("postgresql://u:p@h/db?application_name=") is None
+
+    def test_unparseable_dsn_is_none(self):
+        assert application_name_from_dsn("not a dsn at all ::::") is None
+
+
+class _FakePool:
+    def get_size(self):
+        return 0
+
+    def get_idle_size(self):
+        return 0
+
+
+class _RecordingConnection:
+    """Captures the statements the pool's setup hook issues on acquire."""
+
+    def __init__(self) -> None:
+        self.statements: list[tuple[str, tuple]] = []
+
+    async def execute(self, query: str, *args) -> None:
+        self.statements.append((query, args))
+
+
+@pytest.fixture
+def captured_pool_kwargs(monkeypatch):
+    captured: dict = {}
+
+    async def fake_create_pool(dsn, **kwargs):
+        captured["dsn"] = dsn
+        captured.update(kwargs)
+        return _FakePool()
+
+    monkeypatch.setattr(pg_mod.asyncpg, "create_pool", fake_create_pool)
+    return captured
+
+
+@pytest.mark.asyncio
+async def test_setup_hook_reasserts_application_name_on_every_acquire(captured_pool_kwargs):
+    backend = PostgreSQLBackend()
+    await backend.initialize("postgresql://u:p@h:5432/db?application_name=worker-3")
+
+    # asyncpg runs `setup` on every acquire (after the release-time RESET ALL),
+    # which is the only hook that survives pgbouncer clearing the value.
+    setup = captured_pool_kwargs["setup"]
+    assert setup is captured_pool_kwargs["init"]
+
+    conn = _RecordingConnection()
+    await setup(conn)
+    assert conn.statements == [("SELECT set_config('application_name', $1, false)", ("worker-3",))]
+
+    # A second acquire re-applies it rather than assuming it stuck.
+    await setup(conn)
+    assert len(conn.statements) == 2
+
+
+@pytest.mark.asyncio
+async def test_setup_hook_still_runs_the_callers_init_callback(captured_pool_kwargs):
+    seen: list[object] = []
+
+    async def init_callback(conn):
+        seen.append(conn)
+
+    backend = PostgreSQLBackend()
+    await backend.initialize(
+        "postgresql://u:p@h:5432/db?application_name=worker-3",
+        init_callback=init_callback,
+    )
+
+    conn = _RecordingConnection()
+    await captured_pool_kwargs["setup"](conn)
+    # The session GUCs (hnsw.ef_search, statement_timeout, ...) must still be
+    # applied — wrapping the callback must not displace it.
+    assert seen == [conn]
+
+
+@pytest.mark.asyncio
+async def test_without_application_name_the_callback_is_passed_through(captured_pool_kwargs):
+    async def init_callback(conn):
+        pass
+
+    backend = PostgreSQLBackend()
+    await backend.initialize("postgresql://u:p@h:5432/db", init_callback=init_callback)
+
+    # No name to assert: no wrapper, no extra statement per acquire.
+    assert captured_pool_kwargs["setup"] is init_callback
+    assert captured_pool_kwargs["init"] is init_callback
+
+
+@pytest.mark.asyncio
+async def test_no_application_name_and_no_callback_leaves_hooks_unset(captured_pool_kwargs):
+    backend = PostgreSQLBackend()
+    await backend.initialize("postgresql://u:p@h:5432/db")
+
+    assert captured_pool_kwargs["setup"] is None
+    assert captured_pool_kwargs["init"] is None