[NOTICKET][DB] perf(query): DB2 - pooled engine cache for user DBs

DbExecutor built a fresh engine and disposed it on every query, paying a full
TCP+TLS+auth handshake per call (~6-8s measured, dominating slow-path latency).
engine_scope's connect-once-dispose is right for ingestion, wrong for the query
path.

- New UserEngineCache (src/database_client/engine.py): process-wide cache of pooled
engines keyed by client_id + a hash of the decrypted creds (rotation auto-
invalidates the key). Bounded LRU (50) + 600s idle TTL; pool_pre_ping +
pool_recycle=300; small pool (1 + 2 overflow). invalidate(client_id) disposes
eagerly. Scope: postgres/supabase only; other db_types keep the legacy per-call
path, so nothing regresses.
- DbExecutor._run_sync reuses the warm connection (no per-query SET, no dispose).
- Read-only + statement_timeout moved from libpq startup `options` to a per-
connection `connect` event: Neon's transaction pooler rejects
default_transaction_read_only as a startup parameter (caught in a live run) but
accepts it as a SET. Best-effort; authoritative read-only is the SELECT-only
compiler + sqlglot guard. The per-request ownership/active check is unchanged.

Live-measured: warm query_structured 6.6-9.4s -> ~2.5s. First query per process
still cold (DB3 pre-connect would hide it). Tests: tests/database_client/test_engine.py.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

PROGRESS.md

@@ -41,7 +41,7 @@ Verified against code before logging. Severity: **critical** / important / nice-
 | R2 | **Always compile a LIMIT** — `sql.py` now emits a bound for every query: explicit limit honored (clamped to `MAX_RESULT_ROWS=10000`), unbounded queries get `LIMIT cap+1` so an unbounded SELECT can't stream a whole table into memory. `CompiledSql.row_cap` carries the cap; `DbExecutor` caps + flags truncation from it (dropped its own `_ROW_HARD_CAP`). Tests updated (`test_sql.py`, +3 cases); `S608` restored to `tests/**` ruff ignore (was dropped). | **critical** | DB | `[x]` |
 | R3 | **Commit `tests/` + minimal CI** — `tests/` is gitignored; the 200+ tests cited as done exist only on laptops (already caused rename rot). GitHub origin carries tests; HF Space gets the Docker build (already doesn't COPY tests). | **critical (process)** | shared | `[ ]` |
 | DB1 | **In-memory `describe_source`** (request-scoped `MemoizingCatalogReader`, `reader.py`) + **LLM-client hoist** (shared module-level `ChatHandler` in `chat.py`). Measured live: `describe_source` 3.5s→~2.0s (structured read now served from the planner's cached snapshot; only the unstructured read remains a round-trip), catalog reads/request ~5→~2. External `query_structured` handshake unchanged (DB2's job) so total slow path is ~flat until DB2. Tests: `tests/catalog/test_reader.py`. | important | agent | `[x]` |
-| DB2 | **Keyed engine cache** (LRU, small pool, `pool_recycle`, `pool_pre_ping`) replacing per-call `engine_scope`; read-only + statement_timeout via `connect_args options` (zero SET round-trips, read-only-at-birth); invalidate on client update/delete. Extract to `src/database_client/engine.py`. **Land with/after R1.** | important | DB | `[ ]` |
+| DB2 | **Keyed engine cache** — `src/database_client/engine.py::UserEngineCache` (process singleton): pooled engines keyed by `client_id + creds-hash` (rotation auto-invalidates), bounded LRU (50) + 600s idle TTL, `pool_pre_ping` + `pool_recycle=300`. `DbExecutor._run_sync` reuses the warm connection instead of `create_engine→connect→dispose` per query (postgres/supabase only; other db_types keep the legacy path — no regression). **Live-measured: warm `query_structured` 6.6–9.4s → ~2.5s** (the residual is the per-call catalog-DB client fetch + pre-ping, not the external handshake). **Finding:** Neon's transaction pooler REJECTS `default_transaction_read_only` as a libpq startup `option` — caught live; moved read-only + statement_timeout to a per-connection `connect` event (best-effort; authoritative read-only is the SELECT-only compiler + sqlglot guard, see R10). Per-request ownership/active check kept. Proceeded ahead of R1 per owner decision (marginal security delta over the existing no-auth state; auth tracked separately). Tests: `tests/database_client/test_engine.py`. First query/process still cold → DB3. | important | DB | `[x]` |
 | DB3 | **Speculative pre-connect** overlapped with the ~4s Planner call (only if DB2 still shows first-touch handshake). | nice-to-have | DB | `[ ]` |
 | R4 | **Stream the Assembler + emit TaskRunner progress events** — slow path streams nothing for ~20s (`chat_handler.py:321` single chunk); proxy idle-timeout + UX risk. | important | agent | `[ ]` |
 | R5 | **Response cache**: key on `user_id` + catalog version; invalidate on ingest. Today `chat:{room_id}:{message}`, 24h TTL, no user (`chat.py:138`) → cross-room replay + stale answers. | important | B | `[ ]` |

src/database_client/engine.py

@@ -0,0 +1,168 @@
+"""UserEngineCache — pooled, reused SQLAlchemy engines for users' external DBs.
+
+The query path (`DbExecutor`) previously built a fresh engine and tore it down on
+EVERY query (`db_pipeline_service.engine_scope`), paying a full TCP+TLS+auth
+handshake per call (~6-8s measured, dominating slow-path latency). That helper's
+connect-once-then-dispose semantics are correct for the *ingestion* pipeline
+(infrequent, one connection per run) but wrong for the query path (frequent,
+latency-sensitive, repeated to the same DB).
+
+This module caches one pooled engine per external DB so connections stay warm
+across queries. Scope: **postgres / supabase only** (the measured case and the
+`schema` source type). Other db_types fall back to the legacy per-call path in
+`DbExecutor`, so nothing regresses.
+
+Safety / multi-tenancy:
+- Key = client_id + a hash of the decrypted credentials, so a credential rotation
+  produces a new key (the stale engine idle-evicts) — a cached engine never serves
+  rotated creds.
+- Read-only + statement_timeout are pinned at connection establishment via libpq
+  `options` (read-only-at-birth), so they can't be escaped by a reused pooled
+  connection and cost zero per-query round-trips.
+- The caller still re-fetches the DatabaseClient row every query and re-checks
+  ownership + `active` status — caching the engine never bypasses authorization.
+- Bounded LRU + idle TTL cap memory / file descriptors / connections held on the
+  user's DB. `invalidate(client_id)` disposes eagerly on client update/delete.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import threading
+import time
+from collections import OrderedDict
+
+from sqlalchemy import URL, create_engine, event
+from sqlalchemy.engine import Engine
+
+from src.middlewares.logging import get_logger
+
+logger = get_logger("user_engine_cache")
+
+_POSTGRES_LIKE = frozenset({"postgres", "supabase"})
+_STATEMENT_TIMEOUT_MS = 30_000
+
+# Pool sizing is deliberately small: this is a per-user external DB, often with a
+# low max_connections, and we cache many of them. pool_pre_ping drops dead
+# connections; pool_recycle bounds connection age so a serverless user DB can still
+# autosuspend between bursts.
+_POOL_SIZE = 1
+_MAX_OVERFLOW = 2
+_POOL_RECYCLE_SECONDS = 300
+
+# Cache bounds across all users.
+_MAX_ENGINES = 50
+_IDLE_TTL_SECONDS = 600
+
+
+def _creds_fingerprint(credentials: dict) -> str:
+    blob = json.dumps(credentials, sort_keys=True, default=str)
+    return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16]
+
+
+class UserEngineCache:
+    """Process-wide cache of pooled engines for users' external Postgres DBs.
+
+    Thread-safe: `DbExecutor` runs sync DB work in `asyncio.to_thread` worker
+    threads, so concurrent requests can hit this from multiple threads.
+    """
+
+    def __init__(self) -> None:
+        # key -> (engine, last_used_monotonic)
+        self._engines: OrderedDict[str, tuple[Engine, float]] = OrderedDict()
+        self._lock = threading.Lock()
+
+    def get_engine(self, client_id: str, db_type: str, credentials: dict) -> Engine | None:
+        """Return a pooled engine for (client_id, creds), or None if unsupported.
+
+        None means "not a postgres-like DB" — the caller should use its legacy
+        per-call path for those (rare, unmeasured) db_types.
+        """
+        if db_type not in _POSTGRES_LIKE:
+            return None
+
+        key = f"{client_id}:{_creds_fingerprint(credentials)}"
+        now = time.monotonic()
+        with self._lock:
+            self._evict_idle(now)
+            entry = self._engines.get(key)
+            if entry is not None:
+                self._engines[key] = (entry[0], now)
+                self._engines.move_to_end(key)
+                return entry[0]
+
+            engine = self._build_engine(credentials)
+            self._engines[key] = (engine, now)
+            self._engines.move_to_end(key)
+            self._evict_overflow()
+            logger.info("user engine created", client_id=client_id, cached=len(self._engines))
+            return engine
+
+    def invalidate(self, client_id: str) -> None:
+        """Dispose + drop every cached engine for a client (creds rotated/deleted)."""
+        with self._lock:
+            stale = [k for k in self._engines if k.startswith(f"{client_id}:")]
+            for k in stale:
+                engine, _ = self._engines.pop(k)
+                engine.dispose()
+            if stale:
+                logger.info("user engine invalidated", client_id=client_id, disposed=len(stale))
+
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _build_engine(credentials: dict) -> Engine:
+        # Mirrors db_pipeline_service.connect()'s postgres URL shape, plus a real pool.
+        query = {"sslmode": credentials["ssl_mode"]} if credentials.get("ssl_mode") else {}
+        url = URL.create(
+            drivername="postgresql+psycopg2",
+            username=credentials["username"],
+            password=credentials["password"],
+            host=credentials["host"],
+            port=credentials["port"],
+            database=credentials["database"],
+            query=query,
+        )
+        engine = create_engine(
+            url,
+            pool_size=_POOL_SIZE,
+            max_overflow=_MAX_OVERFLOW,
+            pool_recycle=_POOL_RECYCLE_SECONDS,
+            pool_pre_ping=True,
+        )
+
+        # Apply read-only + statement_timeout once per PHYSICAL connection via a
+        # connect event (not per query, so the pooling latency win stays). These are
+        # ordinary SET commands, NOT libpq startup `options` — Neon's transaction
+        # pooler rejects `default_transaction_read_only` as a startup parameter but
+        # accepts it as a SET. Best-effort: the authoritative read-only guarantee is
+        # the compiler (SELECT-only) + the sqlglot DML guard; statement_timeout is
+        # backed by the executor's asyncio.wait_for. So a failure here must not break
+        # the connection.
+        @event.listens_for(engine, "connect")
+        def _init_session(dbapi_conn, _record):  # noqa: ANN001
+            try:
+                cur = dbapi_conn.cursor()
+                cur.execute(f"SET statement_timeout = {_STATEMENT_TIMEOUT_MS}")
+                cur.execute("SET default_transaction_read_only = on")
+                cur.close()
+            except Exception as exc:  # noqa: BLE001 — best-effort session hardening
+                logger.warning("session init SET failed", error=str(exc))
+
+        return engine
+
+    def _evict_idle(self, now: float) -> None:
+        stale = [k for k, (_, ts) in self._engines.items() if now - ts > _IDLE_TTL_SECONDS]
+        for k in stale:
+            engine, _ = self._engines.pop(k)
+            engine.dispose()
+
+    def _evict_overflow(self) -> None:
+        while len(self._engines) > _MAX_ENGINES:
+            _, (engine, _) = self._engines.popitem(last=False)  # LRU = oldest end
+            engine.dispose()
+
+
+# Process-wide singleton consumed by DbExecutor.
+user_engine_cache = UserEngineCache()

src/query/executor/db.py

@@ -29,6 +29,7 @@ from sqlalchemy import text
 
 from ...catalog.models import Catalog, Source
 from ...database_client.database_client_service import database_client_service
+from ...database_client.engine import user_engine_cache
 from ...db.postgres.connection import AsyncSessionLocal
 from ...middlewares.logging import get_logger
 from ...pipeline.db_pipeline import db_pipeline_service
@@ -41,7 +42,6 @@ logger = get_logger("db_executor")
 
 _QUERY_TIMEOUT_SECONDS = 30
 _DBCLIENT_PREFIX = "dbclient://"
-_POSTGRES_LIKE = frozenset({"postgres", "supabase"})
 
 
 class DbExecutor(BaseExecutor):
@@ -85,7 +85,9 @@ class DbExecutor(BaseExecutor):
             creds = decrypt_credentials_dict(client.credentials)
 
             columns, rows = await asyncio.wait_for(
-                asyncio.to_thread(self._run_sync, client.db_type, creds, compiled),
+                asyncio.to_thread(
+                    self._run_sync, client_id, client.db_type, creds, compiled
+                ),
                 timeout=_QUERY_TIMEOUT_SECONDS,
             )
 
@@ -189,16 +191,22 @@ class DbExecutor(BaseExecutor):
             )
 
     @staticmethod
-    def _run_sync(db_type: str, creds: dict, compiled: CompiledSql) -> tuple[list[str], list[dict]]:
-        with db_pipeline_service.engine_scope(db_type, creds) as engine:
+    def _run_sync(
+        client_id: str, db_type: str, creds: dict, compiled: CompiledSql
+    ) -> tuple[list[str], list[dict]]:
+        engine = user_engine_cache.get_engine(client_id, db_type, creds)
+        if engine is not None:
+            # Pooled, reused engine (postgres-like). Read-only + statement_timeout
+            # are pinned at connection establishment (connect_args options), so no
+            # per-query SET round-trips and no dispose — the connection returns to
+            # the pool warm for the next query.
             with engine.connect() as conn:
-                if db_type in _POSTGRES_LIKE:
-                    # session-level read-only + per-statement timeout (ms)
-                    conn.execute(text("SET default_transaction_read_only = on"))
-                    conn.execute(
-                        text(f"SET statement_timeout = {_QUERY_TIMEOUT_SECONDS * 1000}")
-                    )
                 result = conn.execute(text(compiled.sql), compiled.params)
-                columns = list(result.keys())
-                rows = [dict(row) for row in result.mappings()]
-                return columns, rows
+                return list(result.keys()), [dict(row) for row in result.mappings()]
+
+        # Legacy per-call path for non-postgres db_types (connect once, dispose).
+        # These never set read-only/timeout before, so behavior is unchanged.
+        with db_pipeline_service.engine_scope(db_type, creds) as eng:
+            with eng.connect() as conn:
+                result = conn.execute(text(compiled.sql), compiled.params)
+                return list(result.keys()), [dict(row) for row in result.mappings()]