perf(data-warehouse): reuse the duckgres session across a group's live batches

[EDsCODE]

Problem

After the co-claim fix (#72985), the duckgres sink's throughput ceiling is per-batch session cost. Production phase attribution (2026-07-22, consumer phase histogram + duckgres CP query-log correlation):

Cost	Where
~0.5-1s	worker session create (fresh connection per batch)
~5-19s	extended-protocol describe probe cold-enumerating a 2,000+-table DuckLake catalog (measured +5.3s cold vs warm, empirically)
~12-19s	the transaction's first-write metadata touch on the target table (14k+ files)
~80-100ms	the actual 5,000-row MERGE
<0.7s	marker INSERT + DuckLake commit

Every batch paid the cold costs because process_batch opens a fresh connection each time — ~30-40s per batch, ~7-23/hr per group against a producer emitting ~74/hr.

Change

A (team_id, schema_id)-keyed session cache in the duckgres processor:

• The sink's group lease already serializes each key to one pod task, so cached entries are never used concurrently; the lock guards only the dict.
• Preamble (httpfs load + extract-read secret) runs once per cached session.
• Entries are dropped on any processing error (the connection may hold aborted-transaction state), after 90s idle, and at a 600s absolute age cap — threaded through reuse so the embedded session credentials can never outlive their validity even under steady reuse.
• Backfill batches keep the existing fresh-connection lifecycle (rare, long, already chunk-coalesced).
• Trade-off, documented in code: a cached session pins the group's worker for up to 90s after the group drains — bounded, and the next claim of that group wants the same worker warm anyway.

Expected effect: warm batches at ~2s ⇒ several hundred batches/hour/group; combined with #72985 the sink comfortably outruns the producer, and one-snapshot-per-batch churn slows as backlogs stop accumulating.

How did you test this code?

I'm an agent; checks actually run:

• TDD (each watched failing first): reuse across consecutive batches (one connect, one preamble), per-group keying, drop-on-error, idle/age expiry, and an age-cap-survives-steady-reuse test that caught a real bug in my first implementation (reuse was resetting the age clock).
• Full test_processor.py + test_consumer.py + test_jobs_db.py: 98 passed.
• Pre-commit hooks (ruff, format, ty) green.

Rollback

Revert; behavior returns to fresh-connection-per-batch (slow but correct).

🤖 Agent context

Authored by Claude Code (Fable 5). Scope note: an earlier draft coalesced N batches into one DuckLake transaction, but the production trace showed the commit itself is cheap (<0.7s) while the costs are per-session — so this PR ships the smaller mechanism that captures the measured win without multi-batch status-consistency risk. Snapshot-growth hygiene (77k snapshots on the dogfood org) is a separate follow-up on the duckgres/compaction side.

🤖 Generated with Claude Code

[trunk-io]

Merging to master in this repository is managed by Trunk.

• To merge this pull request, check the box to the left or comment /trunk merge below.

After your PR is submitted to the merge queue, this comment will be automatically updated with its status. If the PR fails, failure details will also be posted here

[hex-security-app]

Bound cached Duckgres sessions by the org connection budget

store() retains a separate connection for every (team_id, schema_id) but never closes it when its 90-second TTL elapses unless another cache operation happens. This bypasses the queue's sink_max_concurrency connection cap: an org user can trigger live batches across many schemas, leaving completed groups' sessions open and exhausting that org's Duckgres max_connections, denying warehouse access to its other projects. The cache is also not bounded while the worker is idle, so its final set of connections persists indefinitely.

Prompt To Fix With AI

Redesign the session cache so retained idle sessions cannot exceed the DuckgresServer sink connection budget and are closed independently when their idle TTL expires (for example, use a bounded cache with scheduled eviction/cleanup). Ensure cache retention is accounted for by the per-org connection budget, rather than only active group leases, and close all sessions on worker shutdown.

---

_{Severity: medium | Confidence: 92% | React with 👍 if useful or 👎 if not}

[greptile-apps]

Prompt To Fix All With AI

Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 1
products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/duckgres/processor.py:145-151
**Idle Sessions Never Expire**

Expiration only runs during a later `acquire()` or `store()`. When groups drain after their final batches, their cached connections remain open past the 90-second idle limit; if enough groups drain together, these sessions can pin all Duckgres workers and prevent a new group from connecting, so cleanup needs an independent timer or group-release hook.

_{Reviews (1): Last reviewed commit: "perf(data-warehouse): reuse the duckgres..." | Re-trigger Greptile}

[greptile-apps]

Idle Sessions Never Expire

Expiration only runs during a later acquire() or store(). When groups drain after their final batches, their cached connections remain open past the 90-second idle limit; if enough groups drain together, these sessions can pin all Duckgres workers and prevent a new group from connecting, so cleanup needs an independent timer or group-release hook.

Prompt To Fix With AI

This is a comment left during a code review.
Path: products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/duckgres/processor.py
Line: 145-151

Comment:
**Idle Sessions Never Expire**

Expiration only runs during a later `acquire()` or `store()`. When groups drain after their final batches, their cached connections remain open past the 90-second idle limit; if enough groups drain together, these sessions can pin all Duckgres workers and prevent a new group from connecting, so cleanup needs an independent timer or group-release hook.

How can I resolve this? If you propose a fix, please make it concise.

[veria-ai]

PR overview

This PR updates the Duckgres-backed data warehouse import pipeline to reuse a Duckgres session across a group's live batches, reducing session setup overhead during ingestion. The touched processor code manages cached sessions for warehouse source batch processing.

One issue remains open after one earlier issue was addressed. The remaining concern is that cached idle sessions can accumulate outside the existing per-organization connection budget, allowing a user to churn through groups and temporarily occupy more worker sessions than intended. This could block subsequent ingestion for that organization until sessions expire or are evicted, so the PR still carries a limited resource-exhaustion risk.

Open issues (1)

• Medium: Cached sessions bypass the connection budget — products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/duckgres/processor.py:156

Fixed/addressed: 1 · PR risk: 5/10

[stamphog]

Note

🤖 stamphog reviewed ca519b9d154d64ca26d4594d5c7dc2978ef3927c — verdict: REFUSED

This touches the data-warehouse ingestion path with a security-relevant unresolved concern: hex-security-app and veria-ai both confirm, on the current head, that the session cache can still retain more connections than the org's connection budget allows, risking exhaustion of a customer's warehouse connections — the author's fixes addressed idle-expiry and org-transfer but not this capacity-bound issue.

• Author wrote 100% of the modified lines and has 21 merged PRs in these paths (familiarity STRONG).
• veria-ai[bot] reviewed the current head.
• Unresolved substantive concern from @hex-security-app (current head): the session cache has no per-org capacity bound tied to sink_max_concurrency, so sequentially-drained groups can accumulate cached connections beyond the org's connection budget and exhaust Duckgres connections for that org.
• @veria-ai's current-head review independently confirms this same open issue ('Medium: Cached sessions bypass the connection budget') and rates it a limited but real resource-exhaustion risk.
• Author (EDsCODE) is not on the owning team @PostHog/team-warehouse-sources, though STRONG familiarity would otherwise cover that gap absent outstanding concerns.

Gate mechanics and policy version

Gate		Result
prerequisites	✓	all clear
deny-list	✓	no deny categories matched
size	✓	198L, 1F substantive, 349L/2F incl. docs/generated/snapshots — within ceiling
tier	✓	T1-agent / T1d-complex (349L, 2F, single-area, perf)
stamphog 2.0.0b3		.stamphog/policy.yml @ b2a037d · reviewed head ca519b9

Updated in place — this replaces 2 earlier stamphog review(s) on this PR.

[hex-security-app]

I reviewed the current changes and found no new security issues to report.

[greptile-apps]

The concern is valid. The docstring says "bounded…for up to IDLE_TTL after the group drains," but that bound is only enforced lazily on a subsequent acquire()/store() call. If a group drains and no further batches ever arrive (for any key), the connections sit open indefinitely. The _session_cache module-level singleton has no background sweeper and no shutdown hook.

The fix is straightforward — add a daemon sweeper thread and an atexit cleanup:

    IDLE_TTL_SECONDS = 90.0
    MAX_AGE_SECONDS = 600.0

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._entries: dict[tuple[int, str], tuple[psycopg.Connection[Any], float, float]] = {}
        self._start_sweeper()

    def _start_sweeper(self) -> None:
        def _sweep() -> None:
            while True:
                time.sleep(self.IDLE_TTL_SECONDS)
                self._evict_stale()

        t = threading.Thread(target=_sweep, daemon=True, name="duckgres-session-sweeper")
        t.start()

    def _evict_stale(self) -> None:
        now = time.monotonic()
        with self._lock:
            stale = [
                k
                for k, (_, created_at, used) in self._entries.items()
                if now - used > self.IDLE_TTL_SECONDS or now - created_at > self.MAX_AGE_SECONDS
            ]
            evicted = [self._entries.pop(k) for k in stale]
        for conn, _, _ in evicted:
            self._close_quietly(conn)

And register a shutdown hook below the _session_cache = _DuckgresSessionCache() line:

import atexit
atexit.register(_session_cache.clear)

The daemon thread wakes every IDLE_TTL_SECONDS and closes any session that has exceeded either TTL. Because it's a daemon thread it won't block process exit, and atexit.clear() handles graceful shutdown. This makes the "bounded" claim in the docstring actually true regardless of whether future batches arrive.

You'll also want a test for this: mock time.monotonic past the idle TTL, call _evict_stale() directly (or advance a fake clock and let the sweeper fire), and assert the connection was closed without any subsequent acquire() call.

[EDsCODE]

Addressed all three review threads in the latest commit:

1. Idle sessions never expire (@greptile-apps P1, @hex-security-app, the stamphog refusal) — confirmed real; the "bounded" docstring claim was false without traffic. Now: a daemon sweeper thread (lazy-started on first store() so importing processes never spawn it; its sleep is bound to the real time.sleep at import so test-patched clocks can't hot-loop it) runs _evict_stale() every 30s, closing entries past the 90s idle TTL or the 600s age cap independently of traffic, plus atexit cleanup for graceful shutdown. Worker pinning after a group drains is now hard-bounded at IDLE_TTL + SWEEP_INTERVAL ≈ 2 minutes. Tested by driving _evict_stale() directly under a controlled clock with zero intervening cache traffic (test_evict_stale_closes_idle_sessions_without_new_traffic), plus a sweeper-started assertion.
2. Cached session survives an org transfer (@veria-ai) — taken as cheap hardening even though project→org transfer isn't a supported product flow today: the cache key now includes the team's current organization_id, resolved on every batch (same per-batch ORM cost the old fresh-connection code paid). A moved team keys to a fresh connection against the new org's warehouse; the orphaned entry is closed by the sweeper within ~2 minutes. Tested by flipping the resolved org between batches (test_org_change_invalidates_the_cached_session).

All watched failing first; full processor + consumer suites green (48).

🤖 Generated with Claude Code

[hex-security-app]

I verified the sweeper and atexit cleanup remove the indefinite-idle-session behavior, and the current-org cache key prevents reuse after an org change. I cannot resolve my finding yet: the cache still has no per-org capacity bound or accounting against sink_max_concurrency. Sequentially completed groups can each retain a connection for 90–120 seconds after their lease is released, so they can accumulate beyond the fleet-wide active-group budget and exhaust Duckgres connections during that window.

[greptile-apps]

Both fixes look correct. Walking through each:

Sweeper (P1): The daemon thread properly addresses the root issue. The _real_sleep capture before any test patch is the right trick — the sweeper's sleep is bound at import time so test-patched time.monotonic can't cause a hot-loop. atexit covers graceful shutdown. Worker pin is now hard-bounded at IDLE_TTL + SWEEP_INTERVAL = 120s max as claimed.

Org key inclusion: Resolved per-batch org_id as part of the cache key is exactly the right fix — a transferred team keys to a fresh connection immediately on the next batch, the orphaned entry closes within the sweep window. The test flipping organization_id between batches covers it cleanly.

One minor observation on the cold-path: a cache miss now does two Team ORM lookups — one in process_batch to compute org_id, and one inside _connect_to_duckgres to resolve the same field for the config. Not a correctness issue (it's the cache-miss path), but if you want to avoid it you could thread org_id into _connect_to_duckgres. Fine to leave as-is.

Implementation looks solid. LGTM.

Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.

[veria-ai]

Medium: Cached sessions bypass the connection budget

Each completed group adds an entry here for up to 120 seconds, but sink_max_concurrency only limits active group leases. A user who schedules batches across many schemas can churn through groups and retain more sessions than the organization budget permits, occupying the one-session workers and blocking subsequent ingestion. Enforce a per-organization cache limit no greater than its connection budget, evicting or closing an existing idle session before storing another.

[EDsCODE]

Both follow-ups addressed:

• @hex-security-app (per-org capacity bound): the cache now enforces MAX_SESSIONS_PER_ORG = 4 with LRU eviction on store — a burst of sequentially drained groups can retain at most 4 of an org's connections per consumer pod (12 fleet-wide at 3 pods, versus e.g. the dogfood org's 64-worker cap), on top of the existing 90s idle / 600s age / sweeper bounds. Retention can no longer scale with the number of schemas an org churns through. Test: test_per_org_session_cap_evicts_the_oldest_entry (5th schema's store closes the least-recently-used entry, others untouched).
• @greptile-apps (double Team lookup): org resolution is hoisted to one per-batch ORM read that feeds both the cache key and _connect_to_duckgres (which now takes the resolved org id); the backfill path uses the same hoisted value. Test: test_cache_miss_resolves_the_team_once.

Both watched failing first; processor + consumer suites green (50).

🤖 Generated with Claude Code

[greptile-apps]

Both fixes look solid. The implementation is clean and the tests are thorough.

MAX_SESSIONS_PER_ORG = 4 — the math checks out. 4 per pod × 3 pods = 12 fleet-wide against a 64-worker org cap is a reasonable bound. The LRU eviction in store() is correct — it sorts by last_used (index [2]) ascending and pops the oldest.

One minor nit on store(): the previous guard closes a conn if a key collision stores a different object. Since acquire() always pop()s the entry before returning, this can only fire in a pathological double-store scenario that the single-writer-per-group guarantee makes impossible in practice. It's fine as defensive code, just slightly surprising on first read — not a blocker.

test_org_change_invalidates_the_cached_session correctly exercises org-a → org-b by mutating organization_id between calls, which confirms the composite key.

test_cache_miss_resolves_the_team_once confirms the single ORM read; the backfill path using the same hoisted org_id means both paths are covered.

LGTM — both follow-ups addressed correctly.

[hex-security-app]

Thanks — I verified the LRU eviction now bounds each process-local cache to four entries and closes the evicted connection. However, this does not fully address my capacity finding: sink_max_concurrency is a configurable, fleet-wide per-org limit (default 4), whereas MAX_SESSIONS_PER_ORG = 4 applies independently in every consumer pod and is not derived from that limit. With three pods, an org can retain 12 idle sessions (plus active leased sessions) even when its configured sink budget is 4, and configurations below 4 are exceeded by a single pod. The cache allowance needs to be accounted against the org's actual fleet-wide connection budget (or otherwise reserved explicitly) before I can mark the finding fixed.

[github-actions]

🤖 CI report

⚠️ Backend snapshots — 11 updated (11 modified, 0 added, 0 deleted)

Query snapshots: Backend query snapshots updated

Changes: 11 snapshots (11 modified, 0 added, 0 deleted)

What this means:

• Query snapshots have been automatically updated to match current output
• These changes reflect modifications to database queries or schema

Next steps:

• Review the query changes to ensure they're intentional
• If unexpected, investigate what caused the query to change

Review snapshot changes →