Skip to content

fix(lakebase): recycle pool connections and pass through dropped pool options - #719

Open
billmcmillin wants to merge 1 commit into
databricks:mainfrom
billmcmillin:fix/lakebase-connection-lifetime
Open

billmcmillin wants to merge 1 commit into
databricks:mainfrom
billmcmillin:fix/lakebase-connection-lifetime

Conversation

@billmcmillin

Copy link
Copy Markdown

Fixes #718.

Summary

createLakebasePool() drops most pg pool options, including maxLifetimeSeconds, so pool connections are never recycled. Access granted through Unity Catalog claims in a Lakebase OAuth token stops working 4 hours after the token is minted, while the connection stays open. A busy pool keeps its connections past that point, and from then on every query that relies on that access fails with 42501 permission denied, until the app restarts.

This PR passes the dropped options through and recycles connections every 30 minutes by default.

Root cause

getLakebasePgConfig() builds the pg.Pool config from a fixed list of 11 keys:

return {
  host, port, user, database, password, ssl, max,
  idleTimeoutMillis, connectionTimeoutMillis,
  statement_timeout, idle_in_transaction_session_timeout,
};

LakebasePoolConfig extends pg's PoolConfig, so any other option type-checks and is then silently dropped. Two of the consequences:

  • maxLifetimeSeconds is always pg-pool's default of 0 (never recycle). In a live soak on main, pool connections were 4.3 hours old by the end of the run.
  • application_name, keepAlive, query_timeout and options never reach the server. On main, every soak connection reported an empty application_name.

Lakebase checks the token only at login, so connections on native Postgres grants keep working however old they get. Claims access doesn't: on main, it lapsed 240.5 minutes after the token was minted, on every connection, while liveness queries on the same connections kept succeeding.

Changes

packages/lakebase/src/pool-config.ts (the fix)

  • Pass maxLifetimeSeconds, application_name, keepAlive, query_timeout and options through to pg.Pool.
  • The key list stays explicit rather than spreading the caller's config. The config also carries workspaceClient, endpoint, logger and telemetry, which must not reach pg. And host, max and the other pool options are normalised by parsePoolConfig() first, which a spread could overwrite. The cost is that the next pg option a caller needs has to be added to the list too.
  • Default maxLifetimeSeconds to 1800 (30 minutes). The default only has to be well under the safe limit below. Recycling is cheap: a recycled connection logs in again with the cached token, so it doesn't mint a new one. The value is resolved with ??, so an explicit 0 (never recycle) is still honoured.
  • New resolveMaxLifetimeSeconds() logs a warning when the lifetime is 0, or above 9000 seconds (2.5 hours), and returns it unchanged; it doesn't cap. The limit is the 4-hour lapse, less up to an hour of token age at login (Lakebase rejects expired tokens, so a token is at most an hour old when a connection opens), less a 30-minute margin for a client held past its lifetime, login time and clock skew.
  • A negative or non-finite maxLifetimeSeconds logs a warning and falls back to the default. Otherwise pg-pool would close every connection as soon as it opens (negative), or never recycle (NaN, which it turns into 0). Before this PR the value was dropped, so apps passing one work today. That's why this warns rather than throws.

packages/lakebase/src/types.ts

  • Document maxLifetimeSeconds on LakebasePoolConfig. It's inherited from PoolConfig; it's redeclared only to carry the docs.
  • The docs cover: what it does, why it matters for Lakebase, which values warn, the default, and that pg-pool recycles only idle connections, so a client that's checked out and held isn't closed until it's released.
  • The wording says Unity Catalog access "can stop working on a connection that stays open for hours" rather than stating 4 hours. The 4 hours is measured, not documented by Lakebase.

docs/docs/api/appkit/Interface.LakebasePoolConfig.md

  • Regenerated by pnpm docs:build from the doc comment above.

packages/lakebase/src/__tests__/connection-lifetime.test.ts (new)

  • Builds real pools with createLakebasePool() and a static password, then reads pool.options to see what reached pg.Pool. Construction is lazy, so no connection is opened and no credentials are requested.
  • Clears the PG*, LAKEBASE_ENDPOINT and DATABRICKS_CLIENT_ID environment variables before each test, so a developer's shell can't change the results.
  • Works out the safe limit itself (4 h − 1 h − 30 min) instead of importing it, so a wrong constant in the code fails the tests.
  • Checks the default is in range (above 0, at most 2.5 h), not that it equals 1800, so the owners can change the default without editing tests.
  • The no-warning case includes 2 hours on purpose. It pins the threshold to the claims lapse, not to the token's 1-hour expiry.
  • Invalid values are compared with the lifetime of a pool built with no setting, not with a hard-coded 1800.

packages/lakebase/src/__tests__/pool-timeouts.test.ts

  • One test.each block checks each of the five options reaches pool.options. It follows the pattern of the timeout test already in this file.

tools/lakebase-soak.ts (new, not published, not run in CI)

  • A manual soak harness. It runs a real pool against a sandbox Lakebase endpoint for hours and checks connection age, token age at login, application_name and table access. It imports the package from source, so the same command compares main with a branch.
  • Each tick, every active query slot runs a query that reports its connection's age, server PID and application_name, and optionally reads a probe table. Slots are added on a ramp, so connections open throughout a token's life. It disables the idle timeout, so connections live long enough to test recycling, and it holds one connection checked out for the whole run, to show the held-client limit.
  • It wraps the password callback to count logins and token mints, and reads each token's age from its JWT. Tokens are never written out.
  • It writes events.jsonl and summary.json, and exits 0 only if every check passes:
    • no stalls, auth errors or query errors;
    • connections recycled within the lifetime;
    • token age at login plus the lifetime stays under 3.5 hours;
    • application_name is on every connection;
    • the probe table stays readable on pool connections (the held connection is reported separately, since it's expected to lose access at 4 hours).
  • Options:
    • --claims-table mints tokens with READ_ONLY claims for a UC table.
    • --ttl-seconds mints shorter-lived claims tokens.
    • --probe-table names the table to read.
  • A run that stalls (sleep, network loss) fails rather than passing on missing data. On macOS the harness keeps the machine awake while it runs.
  • It's 739 of this PR's 1,010 added lines. It can move to a separate PR if reviewers would rather not keep it in tools/.

token-refresh.ts is unchanged, so the pool's token mint rate doesn't change.

Testing

New tests in connection-lifetime.test.ts, plus a pass-through block in pool-timeouts.test.ts that follows that file's existing pattern. 11 of the 12 fail on main and pass with the fix:

Test main This PR
Default lifetime is above 0 and at most 2.5 h 0 (never recycle) 1800 s
Explicit 0 is kept and warns Kept, no warning Kept, warns
A value above 2.5 h warns No warning Warns
A negative, NaN or infinite value uses the default and warns Dropped, no warning Default, warns
No warning at the default, 30 min, 2 h or exactly 2.5 h Passes Passes
application_name, keepAlive, query_timeout, options, maxLifetimeSeconds reach pool.options Dropped Passed through
  • All 75 Lakebase tests pass. pnpm test passes apart from one mock-workspace-client test that already fails on main.
  • pnpm build, pnpm docs:build, pnpm check and pnpm -r typecheck pass.

Live 260-minute soaks on a sandbox Lakebase Autoscaling project (1 CU), as a service principal whose only access to the probe table comes through Unity Catalog, with tokens that request READ_ONLY claims for it:

main This PR
Pool connection age Up to 4.3 h Up to 1800 s
application_name Empty Set on every connection
Table reads on pool connections after 4 h 42501 from minute 240.5, on every connection 0 denied (304 reads after minute 240; 3,769 in the run)
Connection held for the whole run (never released) 42501 from minute 240.5 42501 from minute 240.5, as documented
Token mints About 1 an hour About 1 an hour (5 in 260 minutes)

The lapse isn't tied to the token's expiry. With one-hour tokens, reads kept working for 15 minutes past expiry, and with 5-minute tokens for 29 minutes, so there's no shorter way to test it than a run of more than 4 hours.

Notes

  • Behavior change: pools now recycle connections every 30 minutes. Options that callers already passed, such as query_timeout and keepAlive, now take effect where they were ignored before.
  • A client checked out for more than about 3.5 hours still loses claims access, because pg-pool recycles only on release. This is documented on maxLifetimeSeconds.
  • AppKit's logger prints the new warnings. Standalone @databricks/lakebase logs only errors by default, so there they appear only when a logger with warn is configured. Each pool warns once when it's created, so AppKit's per-user pools repeat the warning per user.
  • getLakebaseOrmConfig() now includes maxLifetimeSeconds. Whether TypeORM and Sequelize pass it to their own pools hasn't been checked.
  • createLakebasePool can't request Unity Catalog claims today. This fix is a prerequisite for adding that, but the feature itself is separate.

This pull request and its description were written by Isaac.

… options

getLakebasePgConfig() built the pg.Pool config from a fixed list of 11
keys, so maxLifetimeSeconds, application_name, keepAlive, query_timeout
and options were silently dropped. Without a lifetime, pool connections
were never recycled. Access granted through Unity Catalog claims in a
Lakebase OAuth token stops working 4 hours after the token is minted
while the connection stays open, so a busy pool kept failing with
"permission denied" until the app restarted.

- Pass the five options through to pg.Pool, keeping the explicit key
  list.
- Default maxLifetimeSeconds to 1800 (30 minutes); an explicit 0 is
  kept.
- Warn when maxLifetimeSeconds is 0 or above 9000 (2.5 hours): the
  4-hour lapse, less up to an hour of token age at login, less a
  30-minute margin.
- Ignore a negative or non-finite maxLifetimeSeconds, with a warning,
  and use the default. pg-pool would otherwise close every connection
  as soon as it opens (negative) or never recycle (NaN).
- Document maxLifetimeSeconds, including that held clients are only
  recycled on release, and regenerate the API docs.
- Add unit tests for the default, the warnings and the pass-through.
- Add tools/lakebase-soak.ts, a soak harness that runs a real pool
  against a sandbox endpoint and checks connection age, token age at
  login, application_name and table access.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Bill McMillin <bill.mcmillin@databricks.com>
@billmcmillin
billmcmillin requested a review from a team as a code owner November 6, 2026 19:18

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant