From dc1c9aff15e0281faa2fdc4398cd47a7dde2f653 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Fri, 14 Aug 2026 15:03:54 +0200 Subject: [PATCH] fix(db): re-assert the DSN's application_name on every pool acquire A DSN like postgresql://...?application_name=my-worker labels the session correctly under psql and on the first connection, then reports an empty application_name in pg_stat_activity for the rest of the service's life once a connection pooler sits in front of PostgreSQL. asyncpg does forward the DSN's application_name in the startup packet (it passes unrecognized DSN query parameters through as server_settings), so a direct connection is attributed correctly. The pool, however, runs RESET ALL on release. Direct to PostgreSQL that is harmless - RESET ALL restores the startup-packet value. Behind pgbouncer the server connection's startup packet is the pooler's own, with no application_name; pgbouncer applies the client's name with a SET when it links client to server, so RESET ALL resets it to empty and pgbouncer - which believes the value is already applied - never re-issues it. Only the first acquire on each server connection is attributed. Re-assert the name from the pool's setup hook, which asyncpg runs on every acquire after RESET ALL. This is the same mechanism the backend already relies on to keep hnsw.ef_search and the other session GUCs applied across connection reuse. set_config() rather than SET because the name is operator-supplied and SET does not accept bind parameters. Verified against pgbouncer 1.25.1 -> PostgreSQL 16: before, acquires 2..n report ''; after, every acquire reports the configured name, including under concurrent load across multiple server connections. --- .../hindsight_api/engine/db/postgresql.py | 62 +++++++- .../tests/test_dsn_application_name.py | 137 ++++++++++++++++++ 2 files changed, 195 insertions(+), 4 deletions(-) create mode 100644 hindsight-api-slim/tests/test_dsn_application_name.py diff --git a/hindsight-api-slim/hindsight_api/engine/db/postgresql.py b/hindsight-api-slim/hindsight_api/engine/db/postgresql.py index 552fa22687..9df3be4d49 100644 --- a/hindsight-api-slim/hindsight_api/engine/db/postgresql.py +++ b/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}, " diff --git a/hindsight-api-slim/tests/test_dsn_application_name.py b/hindsight-api-slim/tests/test_dsn_application_name.py new file mode 100644 index 0000000000..c0ebf62f9b --- /dev/null +++ b/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