fix(server-core): release and connection-test dedicated pre-aggregation drivers

[igorlukanin]

Summary

A deployment with a dedicated pre-aggregation connection leaked that driver's connection pool on every orchestrator teardown, and never connection-tested it.

• Root cause. OrchestratorApi.release() and testConnection() iterated seenDataSources, which holds plain data source names. The driver factory caches a dedicated pre-aggregation driver separately, so that driver was unreachable from either path. addDataSeenSource() is also called from exactly one place — the standalone readiness probe, hardcoded to default — so on most deployments release() reached no internal driver at all.
• The fix. OrchestratorApi now wraps the driver factory it is constructed with and records every driver actually requested, then releases those instances. Because the wrapper is what it passes to QueryOrchestrator, the query path, the pre-aggregation subsystem and the data source probe all record through it — nothing has to guess which drivers exist.
• Release no longer goes back through the factory. It previously called the factory to obtain the driver it was about to close, which after a failed resolution built a fresh driver — opening a connection — purely to close it. It now releases already-resolved instances and skips what never resolved. Each driver is released as it resolves, so one whose initialization hangs cannot stall the others or the orchestrator's own cleanup.
• Teardown completes, then reports. One driver refusing to close no longer strands the rest or the caller's cache clear, but the failure still propagates — shutdown must exit non-zero on it, and a reset must not reload over state it could not free.
• One key rule, one place. driverCacheKey() decides when a pre-aggregation request earns its own connection, and it takes that decision as an argument rather than re-deriving it, so the cache key cannot disagree with the credentials actually used. That fixes a second, order-dependent duplicate: keying on the request flag alone meant a pre-aggregation build arriving before the first query cached separately, and the next query then built another driver for the same connection — which is exactly what a refresh worker does. It also covers the case where a custom driverFactory takes precedence over pre-aggregation env vars, where the two requests resolve to identical credentials and must therefore share one driver. The two compensating aliasing writes this replaces are deleted; the @pre_agg suffix now appears once in the tree instead of twice under two different rules.

Behaviour change worth noting

testConnection() now covers dedicated pre-aggregation connections, so a broken one that was previously invisible to the readiness and liveness probes will now fail them. That is the point of the change, but it is the one way a deployment could newly report unhealthy.

The same change viewed from the load side: a probe now tests every distinct connection a driver has been built for, where before it tested only what addDataSeenSource had announced — in practice just default, and often nothing at all. Requests that share a connection resolve to one driver and are still tested once, so the fan-out is bounded by distinct connections rather than by requests; but on a tenant with many data sources a liveness probe every few seconds now issues N round-trips instead of one. That is the intended direction — a probe that tests one connection out of N is not much of a probe — and it is the reason the union with seenDataSources is kept narrow.

Not in scope: OrchestratorStorage's LRU has no dispose, so an evicted orchestrator is still never released. Closing that would need to avoid dropping pools under queries still running on the evicted orchestrator, which deserves its own change — this one makes it easier, since release() can now neither construct a driver nor be stalled by one.

Test plan

• New test/unit/driver-lifecycle.test.ts (23 tests): both drivers released and tested with a dedicated pre-aggregation connection; named data sources; a single driver for both request orders without one, and for a custom driverFactory overriding the env vars; release never constructing a driver (never-queried, failed-to-resolve, and the external one); an external driver that was used released exactly once without re-entering its factory; a failing release reported without stranding the others, and each of the three callers still completing its own teardown; the record emptied so a second release is a no-op; the readiness probe still forcing the primary connection when only the pre-aggregation driver exists; tracking bounded across repeated requests and probes; the pre-aggregation env-var scan memoized per data source; rollupOnlyMode unchanged.
• Every test red-checked against eight wrong implementations — release-by-seenDataSources (the original bug), no identity dedup, release re-resolving through the factory, dropping the readiness-probe union, a cache key omitting the credentials term, a no-op release, swallowed release errors, and unbounded per-call tracking. Each is caught by a distinct group.
• cubejs-server-core unit suites green: driver-lifecycle, index, OrchestratorApi, OptsHandler, CompilerApi — 99 tests. tsc and eslint clean.
• RefreshScheduler.test.ts fails on unmodified master in this package (order-dependent); verified by stashing, unrelated to this diff and unchanged by it.
• CI must pass

[claude]

Claude finished @igorlukanin's task in 4m 21s —— View job

---

Review: solid fix, one caller-side gap — 1 high, 1 medium, 3 low (4 inline comments)

Full review

• Read the diff (git diff origin/master...HEAD)
• Review OrchestratorApi.ts tracking / release() / testConnection()
• Review utils.ts cache key + server.ts factory keying
• Review test coverage
• Check release() / testConnection() call sites across server-core, api-gateway, cubejs-server
• Post review

Summary

The root-cause analysis is correct and the fix is well-aimed. seenDataSources really is populated from exactly one place (gateway.ts:2927, hardcoded 'default'), so wrapping the factory is the only way to know which drivers exist without duplicating the factory's caching rule. Deriving the cache key from usePreAgg rather than the caller's flag is the right call and removes both compensating aliasing writes plus a genuine order-dependent duplicate on the refresh-worker path. Releasing already-resolved instances instead of re-entering the factory, and deduping by driver identity, are both correct. Promise.allSettled + rethrow is the right shape for teardown. Code comments are unusually good — they explain why, including the alternatives rejected.

Findings

#	Sev	Where	Issue
1	High	OrchestratorApi.ts:394-400	release() can now realistically reject, and all three callers do their cleanup after an unguarded await — storage keeps spent APIs with dead pools, shutdown skips timer cancellation, reset skips reloadEnvVariables/startScheduledRefreshTimer. Needs try/finally.
2	Medium	server.ts:605-606	hasPreAggregationsEnvVars() (full process.env scan + Set alloc) moved above the cache-hit early return, so it now runs on every driver request rather than once per data source. Memoize per dataSource.
3	Low	OrchestratorApi.ts:381	The external driver is still released through its factory, so an unused external driver is constructed + connection-tested purely to be closed — the pattern this PR removes for internal drivers.
4	Low	OrchestratorApi.ts:83	The @pre_agg suffix still appears twice; the driverCacheKey docstring ("cached and tracked under") implies this call site uses it.
5	Low	testConnection()	Probe fan-out grows from 1 internal test to one per requested driver, per probe — no longer just default. Shared connections dedup, so this is bounded by distinct connections, but on a many-data-source tenant a k8s liveness probe every few seconds now issues N round-trips. Intended direction; worth being explicit about in the PR body alongside the pre-agg-coverage note, since it's the same behaviour change viewed from the load side.

Finding 1 is the only one I'd consider blocking. Note it is not introduced by the diff in isolation — the callers were always unguarded — but the diff is what makes that path reachable: before it, release() typically had no internal driver to fail on.

Nit: driverCacheKey lands in a previously-empty core/utils.ts. Colocating with the other driver-resolution helpers (or with hasPreAggregationsEnvVars's consumers) would make it easier to find than a generic utils.

Tests

driver-lifecycle.test.ts is genuinely good — it exercises the real resolveDriver path so isCustomDriverFactory() doesn't silently disable usePreAgg, it covers both request orders, and the "release does not build drivers" group asserts on construction counts, which is what actually distinguishes the fix from the bug. The red-check against eight wrong implementations is the right discipline.

Gaps, in priority order:

1. No test that a caller can finish tearing down when one release() rejects (finding 1). The existing test asserts the sibling driver still closed within one api; the failure mode is one level up, in OrchestratorStorage.releaseConnections.
2. No coverage of the external driver in release() (finding 3) — rollupOnlyMode covers it only for testConnection.
3. No assertion that requestedDrivers is emptied by release(), which the code deliberately does and which finding 1 depends on.

Not verified

Dependencies aren't installed in this checkout (no node_modules), and building query-orchestrator + friends to run Jest wasn't practical here — so I reviewed statically and did not independently reproduce the 91 green tests, tsc, or eslint. lib: ["es2023"] in tsconfig.base.json does cover Promise.allSettled/PromiseRejectedResult, so that part will compile. Leaving the suite to CI.

• branch `igor/core-723-pre-aggregation-drivers-are-never-released-or-connection`

[codecov]

Codecov Report

❌ Patch coverage is 96.59091% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 59.26%. Comparing base (f7ff67a) to head (5281491).

Files with missing lines	Patch %	Lines
...ges/cubejs-server-core/src/core/OrchestratorApi.ts	94.73%	1 Missing and 2 partials ⚠️

❗ There is a different number of reports uploaded between BASE (f7ff67a) and HEAD (5281491). Click for more details.

HEAD has 1 upload less than BASE

Flag	BASE (f7ff67a)	HEAD (5281491)
cubesql	1	0

Additional details and impacted files

@@             Coverage Diff             @@
##           master   #11465       +/-   ##
===========================================
- Coverage   79.46%   59.26%   -20.20%     
===========================================
  Files         480      224      -256     
  Lines       98810    17946    -80864     
  Branches     3636     3644        +8     
===========================================
- Hits        78515    10635    -67880     
+ Misses      19774     6791    -12983     
+ Partials      521      520        -1     

Flag	Coverage Δ
cube-backend	59.26% <96.59%> (+0.15%)	⬆️
cubesql	?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[claude]

Claude finished @igorlukanin's task in 3m 0s —— View job

---

Re-review of 5281491: all four prior findings addressed — 0 high, 0 medium, 2 low (1 inline comment)

Full review

• Read the new commit (git diff origin/master...HEAD)
• Verify prior findings 1–4 addressed
• Re-check the cache-key rule in server.ts against resolveDriver's actual credentials
• Review the new tests
• Resolve stale threads, post review

Prior findings

#	Sev	Status
1	High	Fixed. OrchestratorStorage.releaseConnections uses allSettled + try/finally around storage.clear(), and both CubejsServerCore.releaseConnections and resetInstanceState moved their cleanup into finally. Failure still propagates. Four tests cover it, including that a slow sibling api finishes before the failure is reported — which Promise.all would not have given you.
2	Medium	Fixed. hasPreAggregationsEnvVars memoized per dataSource in preAggregationsEnvVars, back to once per data source.
3	Low	Fixed. recordingExternalDriverFactory records the external driver, the wrapper is swapped into a copy of options so QueryCache/PreAggregations resolve through it too, and release() closes the recorded instance. Both directions tested (never-used → 0 builds; used → released once, no rebuild). Swapping it into options rather than only the class field is the right call — wrapping just this class's reads would have left the orchestrator's own external driver unreleasable.
4	Low	Fixed. driverCacheKey's docstring now states outright that the factory passes the credentials predicate and OrchestratorApi's tracking passes the raw request, so it may hold two keys for one shared driver, and why that's harmless. The call site carries the same note.

Finding 5 (probe fan-out) was informational and the PR body already covers it.

Verified independently

The cache-key change hinges on usePreAgg being the only input that distinguishes the credentials, and it is: resolveDriver receives preAggregations: usePreAgg || false (server.ts:656), the same term the key is derived from — so a key can no longer disagree with the connection the driver was built for. Both collapse cases check out: hasSeparatePreAggEnv === false and isCustomDriverFactory() === true each force usePreAgg false, both requests resolve identically, one entry is correct. The two deleted aliasing writes are exactly what that subsumes.

release() no longer re-enters either factory, and Promise.resolve(driver) is identity for the native promises both factories return, so no extra unhandled-rejection surface — the recorded promise is the same object the original caller already handles.

New findings (both low)

1. server.ts:938 — a throwing finally replaces the pending exception, so a rejecting cancel() (or a throwing startScheduledRefreshTimer()) swaps the release failure for the cleanup one. Still non-zero exit; only the diagnostic changes. Inline.
2. A driver resolved after release() has cleared the map — an in-flight query during shutdown — is recorded into the emptied map and never closed. The pool is about to die with the process, so this matters only for the LRU-eviction case the PR explicitly defers, and it's strictly better than before. Not worth code now; worth remembering when OrchestratorStorage.dispose lands.

Tests

The new groups are aimed at the failure modes rather than the happy path: the OrchestratorStorage test asserts the map is emptied and that the slow sibling finished, which distinguishes allSettled from Promise.all — a Promise.all implementation passes the first assertion and fails the second. "The record is emptied, so a second release is a no-op" pins the invariant the callers' finally blocks rely on. External coverage asserts on build counts, which is the only thing that separates "released what was built" from "built in order to release".

Not verified

Still no node_modules in this checkout, so I reviewed statically and did not re-run the suites, tsc, or eslint. Promise.allSettled/PromiseRejectedResult are covered by lib: ["es2023"]. Leaving the run to CI; codecov reports the patch at 95.8%.

• branch `igor/core-723-pre-aggregation-drivers-are-never-released-or-connection`

[claude]

Low: a finally that throws replaces the pending exception, so if cancel() rejects (or startScheduledRefreshTimer() throws in resetInstanceState) the release failure is silently swapped for the cleanup failure. Both still exit non-zero, so the observable outcome is right — the diagnostic just points at the wrong thing. release() already logs each driver error before rethrowing, so the information isn't lost; noting it only in case you'd rather the release error win.