From 0ab349f35de6282d16677d55dad11b9f37ab1ee2 Mon Sep 17 00:00:00 2001 From: Bill McMillin Date: Fri, 6 Nov 2026 10:13:03 -0500 Subject: [PATCH 1/3] fix(lakebase): recycle pool connections and pass through dropped pool 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 Signed-off-by: Bill McMillin --- .../appkit/Interface.LakebasePoolConfig.md | 33 + .../src/__tests__/connection-lifetime.test.ts | 134 ++++ .../src/__tests__/pool-timeouts.test.ts | 27 + packages/lakebase/src/pool-config.ts | 59 ++ packages/lakebase/src/types.ts | 18 + tools/lakebase-soak.ts | 739 ++++++++++++++++++ 6 files changed, 1010 insertions(+) create mode 100644 packages/lakebase/src/__tests__/connection-lifetime.test.ts create mode 100644 tools/lakebase-soak.ts diff --git a/docs/docs/api/appkit/Interface.LakebasePoolConfig.md b/docs/docs/api/appkit/Interface.LakebasePoolConfig.md index d30f64e84..858975200 100644 --- a/docs/docs/api/appkit/Interface.LakebasePoolConfig.md +++ b/docs/docs/api/appkit/Interface.LakebasePoolConfig.md @@ -73,6 +73,39 @@ const pool = createLakebasePool({ *** +### maxLifetimeSeconds? + +```ts +optional maxLifetimeSeconds: number; +``` + +Maximum time, in seconds, a connection may stay open before the pool +closes it and opens a new one. Passed to `pg.Pool`. + +Lakebase checks the OAuth token only when a connection opens, and access +granted through Unity Catalog claims in the token can stop working on a +connection that stays open for hours. Recycling connections keeps that +from happening. Values up to 9000 (2.5 hours) are safe; larger values, and +`0` (never recycle), log a warning. A negative or non-numeric value logs a +warning and the default is used instead. + +pg-pool recycles only idle connections: a client you check out and hold +isn't closed until you release it. + +#### Default + +```ts +1800 (30 minutes) +``` + +#### Overrides + +```ts +PoolConfig.maxLifetimeSeconds +``` + +*** + ### sslMode? ```ts diff --git a/packages/lakebase/src/__tests__/connection-lifetime.test.ts b/packages/lakebase/src/__tests__/connection-lifetime.test.ts new file mode 100644 index 000000000..edabffa48 --- /dev/null +++ b/packages/lakebase/src/__tests__/connection-lifetime.test.ts @@ -0,0 +1,134 @@ +import pg from "pg"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { createLakebasePool } from "../pool"; +import type { LakebasePoolConfig, Logger } from "../types"; + +// Why connections need a lifetime. Lakebase checks a token only at login and +// rejects expired ones, so a token is at most an hour old when a connection +// opens. An open connection keeps working after that, but access granted +// through Unity Catalog claims in the token stops working 4 hours after the +// token was minted (measured on a live endpoint: every read on the old +// connections then fails with 42501 while the connections stay up). So every +// connection has to close before its token is 4 hours old: +// +// token age at login (≤ 1 h) + connection lifetime < 4 h − margin +// +// The margin covers a connection checked out past its lifetime (pg-pool only +// recycles idle clients), login time and clock skew. A lifetime within +// MAX_SAFE_LIFETIME_MS is safe whatever token the cache hands out. +const HOUR_MS = 60 * 60_000; +const TOKEN_LIFETIME_MS = HOUR_MS; +const CLAIMS_LAPSE_MS = 4 * HOUR_MS; +const SAFETY_MARGIN_MS = HOUR_MS / 2; +const MAX_SAFE_LIFETIME_MS = + CLAIMS_LAPSE_MS - TOKEN_LIFETIME_MS - SAFETY_MARGIN_MS; + +// config.ts falls back to these (PGPORT and PGSSLMODE are read only from the +// environment), so clear them to keep a developer's shell from changing the +// result. +const ENV_KEYS = [ + "PGHOST", + "PGPORT", + "PGDATABASE", + "PGUSER", + "PGSSLMODE", + "LAKEBASE_ENDPOINT", + "DATABRICKS_CLIENT_ID", +] as const; + +beforeEach(() => { + for (const key of ENV_KEYS) { + vi.stubEnv(key, undefined); + } +}); + +const pools: pg.Pool[] = []; + +afterEach(async () => { + vi.unstubAllEnvs(); + await Promise.all(pools.splice(0).map((pool) => pool.end())); +}); + +// Real factory and pg.Pool, so `pool.options` shows what reached pg. +// Construction is lazy, so no connection is opened. +function makePool(config: Partial = {}): pg.Pool { + const pool = createLakebasePool({ + host: "ep-test.database.us-east-1.databricks.com", + database: "databricks_postgres", + user: "test-user@example.com", + password: "test-only", + ...config, + }); + pools.push(pool); + return pool; +} + +/** The lifetime that reached pg.Pool; pg-pool treats 0 as "never recycle". */ +function lifetimeMs(pool: pg.Pool): number { + return (pool.options.maxLifetimeSeconds ?? 0) * 1000; +} + +describe("connection lifetime", () => { + test("recycles connections by default, early enough that claims never lapse", () => { + const pool = makePool(); + + // The exact default is the owners' call; it only has to be safe for a + // connection that logs in with a nearly expired token. + expect(lifetimeMs(pool)).toBeGreaterThan(0); + expect(lifetimeMs(pool)).toBeLessThanOrEqual(MAX_SAFE_LIFETIME_MS); + }); +}); + +describe("connection lifetime warnings", () => { + function makeLogger(): Logger & { warn: ReturnType } { + return { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn() }; + } + + // Policy check: an explicit 0 is a deliberate opt-out, so it's kept, but it + // should never be silent. + test("keeps an explicit lifetime of 0 and warns that connections are never recycled", () => { + const logger = makeLogger(); + const pool = makePool({ maxLifetimeSeconds: 0, logger }); + + expect(lifetimeMs(pool)).toBe(0); + expect(logger.warn).toHaveBeenCalled(); + }); + + // pg-pool would close every connection as soon as it opens for a negative + // lifetime, and never recycle for NaN (it turns NaN into 0), so these must + // not reach pg.Pool. + test.each([ + ["negative", -5], + ["NaN", Number.NaN], + ["infinite", Number.POSITIVE_INFINITY], + ])( + "uses the default lifetime and warns when the lifetime is %s", + (_name, value) => { + const logger = makeLogger(); + const pool = makePool({ maxLifetimeSeconds: value, logger }); + + expect(lifetimeMs(pool)).toBe(lifetimeMs(makePool())); + expect(logger.warn).toHaveBeenCalled(); + }, + ); + + test("warns when the lifetime is too long to recycle connections before claims lapse", () => { + const logger = makeLogger(); + makePool({ maxLifetimeSeconds: MAX_SAFE_LIFETIME_MS / 1000 + 1, logger }); + + expect(logger.warn).toHaveBeenCalled(); + }); + + test("doesn't warn for the default lifetime or a safe one", () => { + const logger = makeLogger(); + makePool({ logger }); + makePool({ maxLifetimeSeconds: 30 * 60, logger }); + // Still safe, so a warning here would be noise. This pins the threshold to + // the claims lapse, not to the token's 1-hour expiry. + makePool({ maxLifetimeSeconds: 2 * 60 * 60, logger }); + makePool({ maxLifetimeSeconds: MAX_SAFE_LIFETIME_MS / 1000, logger }); + + expect(logger.warn).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/lakebase/src/__tests__/pool-timeouts.test.ts b/packages/lakebase/src/__tests__/pool-timeouts.test.ts index 59173b898..4b23408ca 100644 --- a/packages/lakebase/src/__tests__/pool-timeouts.test.ts +++ b/packages/lakebase/src/__tests__/pool-timeouts.test.ts @@ -35,3 +35,30 @@ describe("PostgreSQL pool timeouts", () => { }, ); }); + +// LakebasePoolConfig extends pg's PoolConfig, so these options always +// type-check. These tests catch getLakebasePgConfig() copying only a fixed list +// of keys and silently dropping the rest. +describe("PostgreSQL pool option pass-through", () => { + test.each([ + ["application_name", "my-lakebase-app"], + ["keepAlive", true], + ["query_timeout", 15_000], + ["options", "-c search_path=app"], + ["maxLifetimeSeconds", 900], + ] as const)("passes %s through to pg.Pool", async (key, value) => { + const pool = createLakebasePool({ + host: "localhost", + database: "test", + user: "test", + password: "test-only", + sslMode: "disable", + [key]: value, + }); + try { + expect(pool.options[key]).toEqual(value); + } finally { + await pool.end(); + } + }); +}); diff --git a/packages/lakebase/src/pool-config.ts b/packages/lakebase/src/pool-config.ts index 01b11c2a6..429419502 100644 --- a/packages/lakebase/src/pool-config.ts +++ b/packages/lakebase/src/pool-config.ts @@ -5,6 +5,60 @@ import { type DriverTelemetry, initTelemetry } from "./telemetry"; import { createTokenRefreshCallback } from "./token-refresh"; import type { LakebasePoolConfig, Logger } from "./types"; +/** Default connection lifetime: 30 minutes. */ +const DEFAULT_MAX_LIFETIME_SECONDS = 30 * 60; + +/** + * Longest connection lifetime that's safe for OAuth tokens. Access granted + * through Unity Catalog claims stops working 4 hours after a token is minted + * (observed on a live endpoint; the connection stays open, and queries fail + * with "permission denied"). Lakebase rejects expired tokens at login, so a + * token is at most an hour old when a connection opens. That leaves 3 hours, + * less a 30-minute margin for a client held past its lifetime, login time and + * clock skew. + */ +const MAX_SAFE_LIFETIME_SECONDS = 4 * 3600 - 3600 - 30 * 60; + +/** + * Resolve the connection lifetime, keeping an explicit 0 (never recycle), and + * warn when connections could outlive a token's Unity Catalog claims. + * + * pg-pool closes a connection right after it opens when the lifetime is + * negative, and never recycles it when the lifetime is NaN, so an invalid + * value falls back to the default instead of reaching pg-pool. + */ +function resolveMaxLifetimeSeconds( + userConfig: Partial, + logger?: Logger, +): number { + const maxLifetimeSeconds = + userConfig.maxLifetimeSeconds ?? DEFAULT_MAX_LIFETIME_SECONDS; + + if (!Number.isFinite(maxLifetimeSeconds) || maxLifetimeSeconds < 0) { + logger?.warn( + "maxLifetimeSeconds must be a number of seconds, 0 or more, but is %s. Using the default of %d seconds.", + String(maxLifetimeSeconds), + DEFAULT_MAX_LIFETIME_SECONDS, + ); + return DEFAULT_MAX_LIFETIME_SECONDS; + } + + if (maxLifetimeSeconds === 0) { + logger?.warn( + "maxLifetimeSeconds is 0, so pool connections are never recycled. Access granted through Unity Catalog claims in the OAuth token can stop working on connections that stay open for hours. Set a lifetime of at most %d seconds.", + MAX_SAFE_LIFETIME_SECONDS, + ); + } else if (maxLifetimeSeconds > MAX_SAFE_LIFETIME_SECONDS) { + logger?.warn( + "maxLifetimeSeconds is %d. Access granted through Unity Catalog claims in the OAuth token can stop working on connections that live this long. Use at most %d seconds.", + maxLifetimeSeconds, + MAX_SAFE_LIFETIME_SECONDS, + ); + } + + return maxLifetimeSeconds; +} + /** * Map an SSL mode string to the corresponding `pg` SSL configuration. * @@ -77,6 +131,11 @@ export function getLakebasePgConfig( statement_timeout: userConfig.statement_timeout, idle_in_transaction_session_timeout: userConfig.idle_in_transaction_session_timeout, + maxLifetimeSeconds: resolveMaxLifetimeSeconds(userConfig, logger), + application_name: userConfig.application_name, + keepAlive: userConfig.keepAlive, + query_timeout: userConfig.query_timeout, + options: userConfig.options, }; } diff --git a/packages/lakebase/src/types.ts b/packages/lakebase/src/types.ts index d123e588c..4c668c44f 100644 --- a/packages/lakebase/src/types.ts +++ b/packages/lakebase/src/types.ts @@ -80,6 +80,24 @@ export interface LakebasePoolConfig extends PoolConfig { */ sslMode?: "require" | "disable" | "prefer"; + /** + * Maximum time, in seconds, a connection may stay open before the pool + * closes it and opens a new one. Passed to `pg.Pool`. + * + * Lakebase checks the OAuth token only when a connection opens, and access + * granted through Unity Catalog claims in the token can stop working on a + * connection that stays open for hours. Recycling connections keeps that + * from happening. Values up to 9000 (2.5 hours) are safe; larger values, and + * `0` (never recycle), log a warning. A negative or non-numeric value logs a + * warning and the default is used instead. + * + * pg-pool recycles only idle connections: a client you check out and hold + * isn't closed until you release it. + * + * @default 1800 (30 minutes) + */ + maxLifetimeSeconds?: number; + /** * Telemetry configuration * diff --git a/tools/lakebase-soak.ts b/tools/lakebase-soak.ts new file mode 100644 index 000000000..ffc25b154 --- /dev/null +++ b/tools/lakebase-soak.ts @@ -0,0 +1,739 @@ +#!/usr/bin/env tsx +/** + * Soak test for @databricks/lakebase connection lifetime vs OAuth token expiry. + * + * Runs a real pool against a real Lakebase Autoscaling endpoint for longer than + * one token lifetime (default 75 minutes) and checks that connections are + * recycled before the token they logged in with expires. Use your own sandbox + * project; never point this at a shared database. + * + * It imports the package from source, so it tests whatever is checked out: + * run it on `main` for the "before" numbers and on the fix branch for "after". + * + * What it does + * - Opens query "slots" on a ramp (one more every --ramp-minutes) so new + * connections are created throughout the run, including late in a token's + * life. Each tick, every active slot runs a short query in parallel, which + * keeps them on separate connections. + * - Disables the pool's idle timeout. The package default (30 s) equals the + * 30 s query interval, so idle connections would close between ticks, every + * query would get a fresh connection, and the run would pass on buggy code. + * - Holds one extra connection checked out for the whole run (unless + * --skip-held). pg-pool only recycles a connection once it is released, so + * this one deliberately outlives its token; it's reported separately. + * - Wraps the pool's password callback to count connections and token mints + * and to read each token's remaining life from its JWT `exp` claim. Tokens + * are never written anywhere; they are identified by mint number only. + * + * Pass criteria: no stalls (a gap between successful queries longer than + * three intervals, as when the machine sleeps or loses its network, makes the + * run invalid), no auth or query errors, every pool connection younger than + * the lifetime, every connection recycled before its token is 4 hours old + * (when Unity Catalog claims lapse; oldest token age at login + lifetime under + * 3.5 hours), the custom application_name on every connection, and a run that + * crosses the one-hour token boundary. With --probe-table, every read on a + * pool connection must also succeed; the held connection's reads are reported + * separately and don't count, since it's expected to lose Unity Catalog access + * 4 hours after its token was minted. Mints per pool-hour is reported so the + * load on the credential service can be reviewed. Telemetry span completion is not + * checked here. + * + * Usage + * export DATABRICKS_CONFIG_PROFILE= # Databricks auth + * export PGHOST= PGDATABASE=databricks_postgres + * export LAKEBASE_ENDPOINT=projects/

/branches//endpoints/ + * export PGUSER= + * pnpm exec tsx tools/lakebase-soak.ts [--minutes 75] [--lifetime-seconds 1800] + * + * To see whether access granted through Unity Catalog lapses when the token + * expires (as opposed to native Postgres GRANTs, which don't), run as a service + * principal whose only access to a table is UC-derived, such as a synced table + * it created, and pass --probe-table . Authenticate the service + * principal with DATABRICKS_HOST, DATABRICKS_CLIENT_ID and + * DATABRICKS_CLIENT_SECRET (and unset DATABRICKS_CONFIG_PROFILE); set PGUSER + * to its application ID. + * + * UC-only access works only with a token that requests claims, which + * createLakebasePool can't do. Add --claims-table to + * mint each token with READ_ONLY claims for that UC table instead (keeping the + * package's current cache rule); pair it with --probe-table for the same + * table's Postgres name, e.g. + * --claims-table main.app.items --probe-table app.items + * Add --ttl-seconds 300 to mint 5-minute claims tokens, so a short run can + * show whether access on open connections lapses at, or some multiple of, + * the token's own expiry. + * + * On macOS the harness runs `caffeinate -ims` for its own lifetime, which + * keeps the machine awake while it's on AC power. Closing the lid still + * sleeps a laptop with no external display, so keep it open, plugged in and + * on a network the workspace's IP access list allows (for example the VPN). + * + * Omit --lifetime-seconds to test the package default. Results go to + * --out (default tmp/lakebase-soak//, which is gitignored): events.jsonl and summary.json. + * Exit code is 0 when every check passes, 1 otherwise. Ctrl-C stops early and + * still writes a summary, marked incomplete. + */ + +import { spawn } from "node:child_process"; +import { appendFileSync, mkdirSync, writeFileSync } from "node:fs"; +import { join } from "node:path"; +import { setTimeout as sleep } from "node:timers/promises"; +import { parseArgs } from "node:util"; + +import { + createLakebasePool, + generateDatabaseCredential, + getWorkspaceClient, + RequestedClaimsPermissionSet, +} from "../packages/lakebase/src/index"; + +type Pool = ReturnType; + +const { values: args } = parseArgs({ + options: { + minutes: { type: "string", default: "75" }, + "interval-seconds": { type: "string", default: "30" }, + "lifetime-seconds": { type: "string" }, + slots: { type: "string", default: "8" }, + "ramp-minutes": { type: "string", default: "7" }, + "query-seconds": { type: "string", default: "1" }, + "app-name": { type: "string", default: "appkit-lakebase-soak" }, + "skip-held": { type: "boolean", default: false }, + "probe-table": { type: "string" }, + "claims-table": { type: "string" }, + "ttl-seconds": { type: "string" }, + out: { type: "string" }, + }, +}); + +// Optional table to read on every probe, as schema.table. Use one the identity +// can reach only through Unity Catalog (for example a synced table it created) +// to see whether that access lapses when the token expires. +const PROBE_TABLE = args["probe-table"] + ?.split(".") + .map((part) => `"${part.replaceAll('"', '""')}"`) + .join("."); + +const MINUTES = Number(args.minutes); +const INTERVAL_MS = Number(args["interval-seconds"]) * 1000; +const SLOTS = Number(args.slots); +const RAMP_MS = Number(args["ramp-minutes"]) * 60_000; +const QUERY_SECONDS = Number(args["query-seconds"]); +const APP_NAME = args["app-name"]; +const REQUESTED_LIFETIME_S = + args["lifetime-seconds"] === undefined + ? undefined + : Number(args["lifetime-seconds"]); +// Shorter-lived claims tokens (300-3600 s) let a short run probe what happens +// after expiry. Only used with --claims-table. +const TTL_SECONDS = + args["ttl-seconds"] === undefined ? undefined : Number(args["ttl-seconds"]); +if (TTL_SECONDS !== undefined) { + if (!args["claims-table"]) { + throw new Error("--ttl-seconds only applies with --claims-table"); + } + if (!(TTL_SECONDS >= 300 && TTL_SECONDS <= 3600)) { + throw new Error("--ttl-seconds must be between 300 and 3600"); + } +} +// The token lifetime the run measures against: the requested TTL, or Lakebase's +// default of one hour. +const TOKEN_BOUNDARY_MINUTES = (TTL_SECONDS ?? 3600) / 60; +// A connection whose lifetime timer fires mid-query is closed when the query +// finishes, so allow the query duration plus a little scheduling slack. +const AGE_SLACK_S = QUERY_SECONDS + 5; +// Longer than this without a successful query means the run stalled (the +// machine slept, or lost the network), so its results don't cover that time. +const MAX_GAP_S = Math.max(3 * (INTERVAL_MS / 1000), 90); + +const OUT_DIR = + args.out ?? + join("tmp", "lakebase-soak", new Date().toISOString().replace(/[:.]/g, "-")); +mkdirSync(OUT_DIR, { recursive: true }); +const EVENTS_FILE = join(OUT_DIR, "events.jsonl"); + +const startedAt = Date.now(); + +// Keep macOS awake for as long as this process runs (AC power only; a closed +// lid still sleeps). caffeinate exits by itself when this process does. +if (process.platform === "darwin") { + spawn("caffeinate", ["-ims", "-w", String(process.pid)], { + stdio: "ignore", + detached: true, + }) + .on("error", () => + console.warn("caffeinate not available; the machine may sleep"), + ) + .unref(); +} + +const elapsedS = () => Math.round((Date.now() - startedAt) / 1000); + +function record(event: Record) { + appendFileSync( + EVENTS_FILE, + `${JSON.stringify({ t: new Date().toISOString(), elapsed_s: elapsedS(), ...event })}\n`, + ); +} + +/** + * Seconds until the token's JWT `exp` and since its `iat`, or undefined if it + * isn't a JWT. Without `iat`, the age is worked out from the token lifetime. + */ +function tokenTimes( + token: string, +): { remainingS: number; ageS: number } | undefined { + try { + const payload = JSON.parse( + Buffer.from(token.split(".")[1] ?? "", "base64url").toString("utf8"), + ) as { exp?: unknown; iat?: unknown }; + if (typeof payload.exp !== "number") return undefined; + const nowS = Date.now() / 1000; + const remainingS = Math.round(payload.exp - nowS); + const ageS = + typeof payload.iat === "number" + ? Math.round(nowS - payload.iat) + : TOKEN_BOUNDARY_MINUTES * 60 - remainingS; + return { remainingS, ageS }; + } catch { + return undefined; + } +} + +// Access granted through Unity Catalog claims stops working 4 hours after a +// token was minted, while the connection stays up (measured on a live +// endpoint). A connection is safe if it's recycled before then, with the same +// margin the unit tests use for checkout overrun, login time and clock skew. +const CLAIMS_LAPSE_S = 4 * 60 * 60; +const SAFETY_MARGIN_S = 30 * 60; + +function isAuthError(error: unknown): boolean { + const { code, message } = (error ?? {}) as { + code?: string; + message?: string; + }; + return ( + code === "28P01" || + code === "28000" || + /password|authenticat|token|expired|permission/i.test(message ?? "") + ); +} + +// ── Pool under test ────────────────────────────────────────────────── + +const pool: Pool = createLakebasePool({ + application_name: APP_NAME, + ...(REQUESTED_LIFETIME_S === undefined + ? {} + : { maxLifetimeSeconds: REQUESTED_LIFETIME_S }), + idleTimeoutMillis: 0, + max: SLOTS + 1, + logger: { + debug: () => {}, + info: () => {}, + warn: (message, ...rest) => + record({ type: "package_warning", message, detail: String(rest) }), + error: (message, ...rest) => + record({ type: "package_error", message, detail: String(rest) }), + }, +}); +// What actually reached pg.Pool. 0 means pg never recycles connections. +const effectiveLifetimeS = Number(pool.options.maxLifetimeSeconds ?? 0); + +let connects = 0; +let mints = 0; +let lastToken: string | undefined; +const connectTokenRemaining: number[] = []; +const connectTokenAge: number[] = []; +const packagePassword = pool.options.password; +if (typeof packagePassword !== "function") { + throw new Error( + "Expected an OAuth password callback; set LAKEBASE_ENDPOINT and don't set a static password.", + ); +} + +/** + * Tokens minted with Unity Catalog claims for one table. createLakebasePool + * can't request claims, and access that comes only through UC needs them, so + * this stands in for the package's callback when --claims-table is set. It + * copies the package's current cache rule (reuse until 2 minutes before + * expiry, one mint in flight at a time) so the run shows how today's pool + * would behave if it requested claims. + */ +function createClaimsPassword( + table: string, + ttlSeconds?: number, +): () => Promise { + const endpoint = process.env.LAKEBASE_ENDPOINT; + if (!endpoint) throw new Error("--claims-table needs LAKEBASE_ENDPOINT"); + const workspaceClient = getWorkspaceClient({}); + let cached: { token: string; expiresAt: number } | undefined; + let inFlight: Promise | undefined; + return async () => { + if (cached && Date.now() < cached.expiresAt - 2 * 60_000) { + return cached.token; + } + inFlight ??= generateDatabaseCredential(workspaceClient, { + endpoint, + claims: [ + { + permission_set: RequestedClaimsPermissionSet.READ_ONLY, + resources: [{ table_name: table }], + }, + ], + // The credential API accepts a ttl of 300-3600 s; the package's request + // type doesn't model it, so it's added here for --ttl-seconds. + ...(ttlSeconds === undefined ? {} : { ttl: `${ttlSeconds}s` }), + } as Parameters[1]) + .then((credential) => { + cached = { + token: credential.token, + expiresAt: new Date(credential.expire_time).getTime(), + }; + return credential.token; + }) + .finally(() => { + inFlight = undefined; + }); + return inFlight; + }; +} + +const originalPassword = args["claims-table"] + ? createClaimsPassword(args["claims-table"], TTL_SECONDS) + : (packagePassword as () => Promise); +pool.options.password = async () => { + const token = await originalPassword(); + connects++; + const newToken = token !== lastToken; + if (newToken) { + mints++; + lastToken = token; + } + const times = tokenTimes(token); + if (times) { + connectTokenRemaining.push(times.remainingS); + connectTokenAge.push(times.ageS); + } + record({ + type: "connect", + connect: connects, + token_mint: mints, + new_token: newToken, + token_remaining_s: times?.remainingS, + token_age_s: times?.ageS, + }); + return token; +}; +pool.on("error", (error) => + record({ + type: "pool_error", + code: (error as { code?: string }).code, + message: error.message, + }), +); + +// ── Load ───────────────────────────────────────────────────────────── + +// A session can always see its own row in pg_stat_activity. +const PROBE_SQL = `SELECT a.pid, + extract(epoch FROM now() - a.backend_start)::float8 AS age_s, + a.application_name AS app, + pg_sleep($1::float8)::text AS slept +FROM pg_stat_activity a +WHERE a.pid = pg_backend_pid()`; + +type Probe = { pid: number; age_s: number; app: string }; +type Runner = { query: Pool["query"] }; +type Client = Runner & { + release(error?: Error | boolean): void; + on(event: "error", listener: (error: Error) => void): unknown; + off(event: "error", listener: (error: Error) => void): unknown; +}; + +// pg-pool only handles errors on idle clients. A checked-out client whose +// connection the server drops (restart, failover, scale-to-zero, the 24-hour +// limit) emits "error", and with no listener that crashes the process. Apps +// holding a client must listen; so does the harness, to keep running and +// count these. +const clientErrors = { count: 0, first: "" }; +function clientErrorListener(label: string) { + return (error: Error) => { + const { code } = error as Error & { code?: string }; + clientErrors.count++; + clientErrors.first ||= `${code ?? "no code"}: ${error.message}`; + record({ type: "client_error", slot: label, code, message: error.message }); + }; +} + +const maxAgeByPid = new Map(); +const appNames = new Set(); +let queriesOk = 0; +let queriesFailed = 0; +let authErrors = 0; +let lastSuccessAt = startedAt; +let longestGapS = 0; +const stalls = { + count: 0, + totalS: 0, + firstAtS: undefined as number | undefined, +}; +let maxPoolSize = 0; +const held = { maxAgeS: 0, errors: 0, errorsAfterBoundary: 0, firstError: "" }; + +// Pool connections and the held connection are counted separately: the held +// one is never recycled, so it's expected to lose claims access at 4 hours, +// and only the pool connections decide whether the fix works. +function newTableCounts() { + return { + ok: 0, + denied: 0, + deniedAfterBoundary: 0, + otherErrors: 0, + firstDeniedAtS: undefined as number | undefined, + lastOkAt: startedAt, + deniedConnectionAgesS: [] as number[], + firstError: "", + }; +} +const tableProbe = { pool: newTableCounts(), held: newTableCounts() }; + +function tableCountsSummary(counts: ReturnType) { + return { + ok: counts.ok, + permission_denied: counts.denied, + permission_denied_after_token_boundary: counts.deniedAfterBoundary, + other_errors: counts.otherErrors, + first_denied_at_s: counts.firstDeniedAtS ?? null, + youngest_denied_connection_age_s: counts.deniedConnectionAgesS.length + ? Math.min(...counts.deniedConnectionAgesS) + : null, + first_error: counts.firstError || null, + }; +} + +// Runs on the same connection as the age probe, so a permission error can be +// tied to how old that connection (and the token it logged in with) was. +async function probeTable(client: Client, label: string, row: Probe) { + if (!PROBE_TABLE) return; + const counts = label === "held" ? tableProbe.held : tableProbe.pool; + try { + await client.query(`SELECT count(*) FROM ${PROBE_TABLE}`); + counts.ok++; + counts.lastOkAt = Date.now(); + record({ + type: "table_query", + slot: label, + pid: row.pid, + age_s: Math.round(row.age_s), + }); + } catch (error) { + const { code, message } = error as { code?: string; message?: string }; + if (code === "42501") { + counts.denied++; + counts.firstDeniedAtS ??= elapsedS(); + if (elapsedS() > TOKEN_BOUNDARY_MINUTES * 60) { + counts.deniedAfterBoundary++; + } + counts.deniedConnectionAgesS.push(Math.round(row.age_s)); + } else { + counts.otherErrors++; + } + counts.firstError ||= message ?? String(error); + record({ + type: "table_error", + slot: label, + pid: row.pid, + age_s: Math.round(row.age_s), + code, + message, + }); + } +} + +// Pool slots check a connection out for the probe and return it afterwards; +// the held connection is passed in and never returned. +async function probe(label: string, heldConnection?: Client) { + let client = heldConnection; + let broken: Error | undefined; + const onError = clientErrorListener(label); + try { + if (!client) { + client = (await pool.connect()) as Client; + client.on("error", onError); + } + const { rows } = await client.query(PROBE_SQL, [QUERY_SECONDS]); + const row = rows[0] as Probe; + appNames.add(row.app); + if (label === "held") { + held.maxAgeS = Math.max(held.maxAgeS, row.age_s); + } else { + maxAgeByPid.set( + row.pid, + Math.max(maxAgeByPid.get(row.pid) ?? 0, row.age_s), + ); + } + queriesOk++; + noteGap(Date.now()); + lastSuccessAt = Date.now(); + record({ + type: "query", + slot: label, + pid: row.pid, + age_s: Math.round(row.age_s), + app: row.app, + }); + await probeTable(client, label, row); + } catch (error) { + broken = error as Error; + const { code, message } = error as { code?: string; message?: string }; + const auth = isAuthError(error); + if (label === "held") { + held.errors++; + if (elapsedS() > TOKEN_BOUNDARY_MINUTES * 60) held.errorsAfterBoundary++; + held.firstError ||= message ?? String(error); + } else { + queriesFailed++; + if (auth) authErrors++; + } + record({ type: "query_error", slot: label, code, auth, message }); + } finally { + // Like pool.query: a connection that errored is discarded, not reused. + if (!heldConnection) { + client?.off("error", onError); + client?.release(broken); + } + } +} + +// Called with the time of each success, and once at the end so a run that +// stops succeeding partway through still counts the trailing gap. +function noteGap(now: number) { + const gapS = (now - lastSuccessAt) / 1000; + longestGapS = Math.max(longestGapS, gapS); + if (gapS > MAX_GAP_S) { + stalls.count++; + stalls.totalS += gapS; + stalls.firstAtS ??= Math.round((lastSuccessAt - startedAt) / 1000); + record({ + type: "stall", + from_s: Math.round((lastSuccessAt - startedAt) / 1000), + gap_s: Math.round(gapS), + }); + } +} + +let stopRequested = false; +process.on("SIGINT", () => { + if (stopRequested) process.exit(130); + stopRequested = true; + console.log("\nStopping after this tick (Ctrl-C again to quit immediately)…"); +}); + +record({ + type: "start", + minutes: MINUTES, + interval_s: INTERVAL_MS / 1000, + requested_lifetime_s: REQUESTED_LIFETIME_S ?? "package default", + effective_lifetime_s: effectiveLifetimeS, + slots: SLOTS, + ramp_minutes: RAMP_MS / 60_000, + held_connection: !args["skip-held"], + token_source: args["claims-table"] + ? `claims for ${args["claims-table"]}` + : "package", +}); +console.log( + `Soak: ${MINUTES} min, lifetime reaching pg.Pool = ${effectiveLifetimeS}s` + + `${effectiveLifetimeS === 0 ? " (never recycled)" : ""}, ${SLOTS} slots, results in ${OUT_DIR}`, +); + +// Fail fast with a readable message if the first connection can't be made +// (wrong PGHOST, expired Databricks login, missing permissions, …). +let heldClient: Client | undefined; +try { + await pool.query("SELECT 1"); + heldClient = args["skip-held"] + ? undefined + : ((await pool.connect()) as Client); + heldClient?.on("error", clientErrorListener("held")); +} catch (error) { + const { code, message } = error as { code?: string; message?: string }; + record({ type: "startup_error", code, message }); + console.error( + `Could not connect to Lakebase: ${message ?? String(error)}${code ? ` (${code})` : ""}`, + ); + await pool.end().catch(() => {}); + process.exit(1); +} +const endAt = startedAt + MINUTES * 60_000; +for (let tick = 0; Date.now() < endAt && !stopRequested; tick++) { + const activeSlots = Math.min( + SLOTS, + 1 + Math.floor((Date.now() - startedAt) / RAMP_MS), + ); + await Promise.all([ + ...Array.from({ length: activeSlots }, (_, i) => probe(`slot-${i}`)), + ...(heldClient ? [probe("held", heldClient)] : []), + ]); + maxPoolSize = Math.max(maxPoolSize, pool.totalCount); + record({ + type: "pool", + total: pool.totalCount, + idle: pool.idleCount, + waiting: pool.waitingCount, + }); + console.log( + `[${String(Math.floor(elapsedS() / 60)).padStart(3)} min] slots=${activeSlots} ` + + `connections=${pool.totalCount} connects=${connects} mints=${mints} ` + + `ok=${queriesOk} failed=${queriesFailed} auth_errors=${authErrors}`, + ); + const nextTickAt = startedAt + (tick + 1) * INTERVAL_MS; + if (Date.now() - lastSuccessAt > MAX_GAP_S * 1000) { + console.warn( + ` no successful query for ${Math.round((Date.now() - lastSuccessAt) / 1000)} s; this stretch of the run is invalid`, + ); + } + await sleep(Math.max(0, nextTickAt - Date.now())); +} + +heldClient?.release(); +const waitingAtEnd = pool.waitingCount; +await pool.end(); + +// ── Summary ────────────────────────────────────────────────────────── + +const endedAt = Date.now(); +noteGap(endedAt); +const durationH = (endedAt - startedAt) / 3_600_000; +const ages = [...maxAgeByPid.values()].sort((a, b) => a - b); +const pct = (p: number) => + ages.length + ? Math.round(ages[Math.min(ages.length - 1, Math.floor(p * ages.length))]) + : null; +const maxAge = ages.at(-1) ?? 0; +const minTokenAtConnect = connectTokenRemaining.length + ? Math.min(...connectTokenRemaining) + : undefined; +const maxTokenAgeAtConnect = connectTokenAge.length + ? Math.max(...connectTokenAge) + : undefined; + +const checks = { + // A sleeping or disconnected machine makes no queries, so it can't see + // errors either; without this, a stalled run could still pass. + no_stalls: stalls.count === 0, + run_crossed_token_boundary: durationH * 60 > TOKEN_BOUNDARY_MINUTES, + no_auth_errors: authErrors === 0, + no_query_errors: queriesFailed === 0, + connections_recycled: + effectiveLifetimeS > 0 && maxAge <= effectiveLifetimeS + AGE_SLACK_S, + // The oldest token any connection logged in with, plus the lifetime, must + // stay clear of the 4-hour claims lapse. The held connection is excluded by + // design: it's never recycled, to show what a long-held client does. + connections_close_before_claims_lapse: + effectiveLifetimeS > 0 && + maxTokenAgeAtConnect !== undefined && + maxTokenAgeAtConnect + effectiveLifetimeS < + CLAIMS_LAPSE_S - SAFETY_MARGIN_S, + application_name_on_every_connection: + appNames.size === 1 && appNames.has(APP_NAME), + no_clients_waiting_at_end: waitingAtEnd === 0, + ...(PROBE_TABLE + ? { + probe_table_readable_throughout: + tableProbe.pool.ok > 0 && + (endedAt - tableProbe.pool.lastOkAt) / 1000 <= MAX_GAP_S && + tableProbe.pool.denied === 0 && + tableProbe.pool.otherErrors === 0, + } + : {}), +}; +const passed = Object.values(checks).every(Boolean); + +const summary = { + passed, + incomplete: stopRequested, + checks, + duration_minutes: Math.round(durationH * 60), + token_source: args["claims-table"] + ? `UC claims for ${args["claims-table"]} (harness-minted, package cache rule, ttl ${TTL_SECONDS ?? 3600}s)` + : "package (createLakebasePool, no claims)", + lifetime: { + requested_s: REQUESTED_LIFETIME_S ?? "package default", + effective_s: effectiveLifetimeS, + }, + mints: { + total: mints, + per_pool_hour: Number((mints / durationH).toFixed(2)), + // Excludes the first mint at startup, which every pool pays once; this is + // the steady-state refresh rate (about 1 per hour with today's cache). + refreshes_per_pool_hour: Number( + (Math.max(0, mints - 1) / durationH).toFixed(2), + ), + note: "Token mints per pool-hour, for reviewing load on the credential service.", + }, + connections_opened: connects, + token_at_connect_s: { + min_remaining: minTokenAtConnect ?? "unknown (token is not a JWT)", + max_age: maxTokenAgeAtConnect ?? "unknown (token is not a JWT)", + readings: connectTokenAge.length, + safe_limit_for_age_plus_lifetime: CLAIMS_LAPSE_S - SAFETY_MARGIN_S, + }, + pool_connection_age_s: { + connections: ages.length, + min: pct(0), + median: pct(0.5), + p95: pct(0.95), + max: ages.length ? Math.round(maxAge) : null, + }, + queries: { + ok: queriesOk, + failed: queriesFailed, + auth_errors: authErrors, + success_rate: Number( + (queriesOk / Math.max(1, queriesOk + queriesFailed)).toFixed(4), + ), + longest_gap_between_successes_s: Math.round(longestGapS), + }, + stalls: { + note: `Gaps of more than ${MAX_GAP_S} s with no successful query (sleep, network loss); the run doesn't cover them.`, + count: stalls.count, + total_s: Math.round(stalls.totalS), + first_at_s: stalls.firstAtS ?? null, + }, + max_pool_size: maxPoolSize, + application_names_seen: [...appNames], + held_connection: args["skip-held"] + ? "skipped" + : { + note: "Checked out for the whole run, so never recycled; shows what a long-held connection does after its token expires.", + max_age_s: Math.round(held.maxAgeS), + errors: held.errors, + errors_after_token_boundary: held.errorsAfterBoundary, + first_error: held.firstError || null, + }, + client_errors: { + note: "Errors emitted by checked-out clients (server dropped the connection). Without a listener these crash a Node app.", + count: clientErrors.count, + first: clientErrors.first || null, + }, + probe_table: PROBE_TABLE + ? { + table: args["probe-table"], + note: "permission_denied counts SQLSTATE 42501 (UC-derived access lapsing). Only pool connections decide probe_table_readable_throughout; the held connection is never recycled, so it's expected to lose access 4 hours after its token was minted.", + pool_connections: tableCountsSummary(tableProbe.pool), + held_connection: args["skip-held"] + ? "skipped" + : tableCountsSummary(tableProbe.held), + } + : "not set", +}; + +writeFileSync( + join(OUT_DIR, "summary.json"), + `${JSON.stringify(summary, null, 2)}\n`, +); +console.log(JSON.stringify(summary, null, 2)); +console.log(`\n${passed ? "PASS" : "FAIL"}: details in ${OUT_DIR}`); +process.exit(passed ? 0 : 1); From 888448ffb07cd29e67f4267bf6f8cf5b25ee0db4 Mon Sep 17 00:00:00 2001 From: Bill McMillin Date: Thu, 12 Nov 2026 08:25:22 -0500 Subject: [PATCH 2/3] fix(lakebase): use the default lifetime when it overflows Node's timer pg-pool passes maxLifetimeSeconds * 1000 to setTimeout. Above 2^31 - 1 ms (about 24.8 days), Node fires the timer after 1 ms, so every connection closed after a single use. Treat those values like Infinity: warn and use the default. Compare in milliseconds so 2147483 s is kept and fractional values just past the limit are rejected. Co-authored-by: Isaac Signed-off-by: Bill McMillin --- .../appkit/Interface.LakebasePoolConfig.md | 5 +++-- .../src/__tests__/connection-lifetime.test.ts | 16 +++++++++++++- packages/lakebase/src/pool-config.ts | 21 +++++++++++++++---- packages/lakebase/src/types.ts | 5 +++-- 4 files changed, 38 insertions(+), 9 deletions(-) diff --git a/docs/docs/api/appkit/Interface.LakebasePoolConfig.md b/docs/docs/api/appkit/Interface.LakebasePoolConfig.md index 858975200..51f5f49cf 100644 --- a/docs/docs/api/appkit/Interface.LakebasePoolConfig.md +++ b/docs/docs/api/appkit/Interface.LakebasePoolConfig.md @@ -86,8 +86,9 @@ Lakebase checks the OAuth token only when a connection opens, and access granted through Unity Catalog claims in the token can stop working on a connection that stays open for hours. Recycling connections keeps that from happening. Values up to 9000 (2.5 hours) are safe; larger values, and -`0` (never recycle), log a warning. A negative or non-numeric value logs a -warning and the default is used instead. +`0` (never recycle), log a warning. A negative or non-numeric value, or +one above 2147483 (about 24.8 days, the longest timer Node supports), logs +a warning and the default is used instead. pg-pool recycles only idle connections: a client you check out and hold isn't closed until you release it. diff --git a/packages/lakebase/src/__tests__/connection-lifetime.test.ts b/packages/lakebase/src/__tests__/connection-lifetime.test.ts index edabffa48..9db43d460 100644 --- a/packages/lakebase/src/__tests__/connection-lifetime.test.ts +++ b/packages/lakebase/src/__tests__/connection-lifetime.test.ts @@ -97,11 +97,15 @@ describe("connection lifetime warnings", () => { // pg-pool would close every connection as soon as it opens for a negative // lifetime, and never recycle for NaN (it turns NaN into 0), so these must - // not reach pg.Pool. + // not reach pg.Pool. Above 2^31 - 1 ms, Node's setTimeout fires after 1 ms, + // so pg-pool would close every connection after one use. test.each([ ["negative", -5], ["NaN", Number.NaN], ["infinite", Number.POSITIVE_INFINITY], + ["past Node's timer limit", 2_147_484], + ["just past Node's timer limit, fractionally", 2_147_483.7], + ["30 days", 30 * 24 * 60 * 60], ])( "uses the default lifetime and warns when the lifetime is %s", (_name, value) => { @@ -113,6 +117,16 @@ describe("connection lifetime warnings", () => { }, ); + // The longest whole-second lifetime Node's timer can hold is still valid: + // it's kept (with the too-long warning), not replaced by the default. + test("keeps the longest lifetime Node's timer can hold", () => { + const logger = makeLogger(); + const pool = makePool({ maxLifetimeSeconds: 2_147_483, logger }); + + expect(pool.options.maxLifetimeSeconds).toBe(2_147_483); + expect(logger.warn).toHaveBeenCalled(); + }); + test("warns when the lifetime is too long to recycle connections before claims lapse", () => { const logger = makeLogger(); makePool({ maxLifetimeSeconds: MAX_SAFE_LIFETIME_MS / 1000 + 1, logger }); diff --git a/packages/lakebase/src/pool-config.ts b/packages/lakebase/src/pool-config.ts index 429419502..400b85a2c 100644 --- a/packages/lakebase/src/pool-config.ts +++ b/packages/lakebase/src/pool-config.ts @@ -19,13 +19,21 @@ const DEFAULT_MAX_LIFETIME_SECONDS = 30 * 60; */ const MAX_SAFE_LIFETIME_SECONDS = 4 * 3600 - 3600 - 30 * 60; +/** + * Longest delay Node's `setTimeout` accepts (about 24.8 days). pg-pool passes + * `maxLifetimeSeconds * 1000` to it, and a longer delay fires after 1 ms, so + * every connection would close after a single use. + */ +const MAX_TIMER_MS = 2 ** 31 - 1; + /** * Resolve the connection lifetime, keeping an explicit 0 (never recycle), and * warn when connections could outlive a token's Unity Catalog claims. * * pg-pool closes a connection right after it opens when the lifetime is - * negative, and never recycles it when the lifetime is NaN, so an invalid - * value falls back to the default instead of reaching pg-pool. + * negative or overflows Node's timer, and never recycles it when the lifetime + * is NaN, so an invalid value falls back to the default instead of reaching + * pg-pool. */ function resolveMaxLifetimeSeconds( userConfig: Partial, @@ -34,9 +42,14 @@ function resolveMaxLifetimeSeconds( const maxLifetimeSeconds = userConfig.maxLifetimeSeconds ?? DEFAULT_MAX_LIFETIME_SECONDS; - if (!Number.isFinite(maxLifetimeSeconds) || maxLifetimeSeconds < 0) { + if ( + !Number.isFinite(maxLifetimeSeconds) || + maxLifetimeSeconds < 0 || + maxLifetimeSeconds * 1000 > MAX_TIMER_MS + ) { logger?.warn( - "maxLifetimeSeconds must be a number of seconds, 0 or more, but is %s. Using the default of %d seconds.", + "maxLifetimeSeconds must be a number of seconds from 0 to %d, but is %s. Using the default of %d seconds.", + Math.floor(MAX_TIMER_MS / 1000), String(maxLifetimeSeconds), DEFAULT_MAX_LIFETIME_SECONDS, ); diff --git a/packages/lakebase/src/types.ts b/packages/lakebase/src/types.ts index 4c668c44f..eb7fb9ce4 100644 --- a/packages/lakebase/src/types.ts +++ b/packages/lakebase/src/types.ts @@ -88,8 +88,9 @@ export interface LakebasePoolConfig extends PoolConfig { * granted through Unity Catalog claims in the token can stop working on a * connection that stays open for hours. Recycling connections keeps that * from happening. Values up to 9000 (2.5 hours) are safe; larger values, and - * `0` (never recycle), log a warning. A negative or non-numeric value logs a - * warning and the default is used instead. + * `0` (never recycle), log a warning. A negative or non-numeric value, or + * one above 2147483 (about 24.8 days, the longest timer Node supports), logs + * a warning and the default is used instead. * * pg-pool recycles only idle connections: a client you check out and hold * isn't closed until you release it. From 2e44157b5824597c14f7672ea8fcd6120dc98f29 Mon Sep 17 00:00:00 2001 From: Bill McMillin Date: Thu, 12 Nov 2026 11:23:29 -0500 Subject: [PATCH 3/3] fix(lakebase): deliver getLakebaseOrmConfig() pool settings to TypeORM and Sequelize getLakebaseOrmConfig() returned pool and client settings at the top level, where neither ORM reads them, so TypeORM and Sequelize users silently lost connection recycling, the pool size and timeouts, and (Sequelize) SSL. - TypeORM: pool size, connect timeout and application name go under its own poolSize, connectTimeoutMS and applicationName, so caller values still win; lifetime, idle and statement timeouts go under extra. - Sequelize: SSL and client timeouts go under dialectOptions, pool size and idle under pool. Its pool has no lifetime option, so connections are recycled by age with pool.validate and an afterConnect hook. - Log only when the caller passes a logger, so nothing new is printed by default. Fixes #730 Co-authored-by: Isaac Signed-off-by: Bill McMillin --- .../appkit/Function.getLakebaseOrmConfig.md | 181 +++++++- packages/lakebase/README.md | 32 +- packages/lakebase/package.json | 4 +- .../lakebase/src/__tests__/orm-config.test.ts | 404 ++++++++++++++++++ packages/lakebase/src/pool-config.ts | 169 +++++++- pnpm-lock.yaml | 6 + 6 files changed, 785 insertions(+), 11 deletions(-) create mode 100644 packages/lakebase/src/__tests__/orm-config.test.ts diff --git a/docs/docs/api/appkit/Function.getLakebaseOrmConfig.md b/docs/docs/api/appkit/Function.getLakebaseOrmConfig.md index 4e98e74c9..c2aa79006 100644 --- a/docs/docs/api/appkit/Function.getLakebaseOrmConfig.md +++ b/docs/docs/api/appkit/Function.getLakebaseOrmConfig.md @@ -2,7 +2,40 @@ ```ts function getLakebaseOrmConfig(config?: Partial): { + applicationName: string | undefined; + connectTimeoutMS: number | undefined; + dialectOptions: DefinedOnly<{ + application_name: string | undefined; + connectionTimeoutMillis: number | undefined; + idle_in_transaction_session_timeout?: number; + keepAlive?: boolean; + options?: string; + query_timeout?: number; + ssl: | boolean + | { + rejectUnauthorized: boolean | undefined; + }; + statement_timeout?: number | false; + }>; + extra: DefinedOnly<{ + idle_in_transaction_session_timeout?: number; + idleTimeoutMillis: number | undefined; + keepAlive?: boolean; + maxLifetimeSeconds: number | undefined; + options?: string; + query_timeout?: number; + statement_timeout?: number | false; + }>; + hooks: { + afterConnect: (connection: unknown) => void; + }; password: string | () => string | () => Promise | undefined; + pool: DefinedOnly<{ + idle: number | undefined; + max: number | undefined; + validate: (connection?: unknown) => boolean; + }>; + poolSize: number | undefined; ssl: | boolean | { rejectUnauthorized: boolean | undefined; @@ -21,6 +54,27 @@ Returns connection config with field names compatible with common ORMs: - Simplified SSL config - Password callback support for OAuth token refresh +Pool and client settings (connection lifetime, pool size, timeouts, SSL) +are also returned where each ORM reads them, so recycling works with both: +- TypeORM: pool size, connect timeout and application name under its own + `poolSize`, `connectTimeoutMS` and `applicationName`, so values you set + after the spread still win. Everything else (lifetime, idle timeout, + statement and query timeouts) goes under `extra`, which TypeORM merges + into the `pg.Pool` options last. +- Sequelize: `dialectOptions` (passed to each `pg.Client`), `pool`, and + `hooks`. Sequelize's pool has no lifetime option, so connections are + recycled by age through `pool.validate` and an `afterConnect` hook. + +The result carries keys for both ORMs, and each ignores the other's: +Sequelize ignores `poolSize`, `connectTimeoutMS`, `applicationName` and +`extra`, and TypeORM ignores `dialectOptions`, `pool` and `hooks`. + +If you set your own `extra`, `pool`, +`dialectOptions` or `hooks`, merge them with these instead of replacing +them, or those settings are lost. Replacing `hooks` makes lifetimes count +from a connection's first reuse instead of when it opened, and logs a +warning if you passed a `logger`. Nothing is logged without one. + ## Parameters | Parameter | Type | Description | @@ -31,7 +85,40 @@ Returns connection config with field names compatible with common ORMs: ```ts { + applicationName: string | undefined; + connectTimeoutMS: number | undefined; + dialectOptions: DefinedOnly<{ + application_name: string | undefined; + connectionTimeoutMillis: number | undefined; + idle_in_transaction_session_timeout?: number; + keepAlive?: boolean; + options?: string; + query_timeout?: number; + ssl: | boolean + | { + rejectUnauthorized: boolean | undefined; + }; + statement_timeout?: number | false; + }>; + extra: DefinedOnly<{ + idle_in_transaction_session_timeout?: number; + idleTimeoutMillis: number | undefined; + keepAlive?: boolean; + maxLifetimeSeconds: number | undefined; + options?: string; + query_timeout?: number; + statement_timeout?: number | false; + }>; + hooks: { + afterConnect: (connection: unknown) => void; + }; password: string | () => string | () => Promise | undefined; + pool: DefinedOnly<{ + idle: number | undefined; + max: number | undefined; + validate: (connection?: unknown) => boolean; + }>; + poolSize: number | undefined; ssl: | boolean | { rejectUnauthorized: boolean | undefined; @@ -42,12 +129,96 @@ Returns connection config with field names compatible with common ORMs: ORM-compatible connection configuration +### applicationName + +```ts +applicationName: string | undefined = pgConfig.application_name; +``` + +### connectTimeoutMS + +```ts +connectTimeoutMS: number | undefined = pgConfig.connectionTimeoutMillis; +``` + +### dialectOptions + +```ts +dialectOptions: DefinedOnly<{ + application_name: string | undefined; + connectionTimeoutMillis: number | undefined; + idle_in_transaction_session_timeout?: number; + keepAlive?: boolean; + options?: string; + query_timeout?: number; + ssl: | boolean + | { + rejectUnauthorized: boolean | undefined; + }; + statement_timeout?: number | false; +}>; +``` + +### extra + +```ts +extra: DefinedOnly<{ + idle_in_transaction_session_timeout?: number; + idleTimeoutMillis: number | undefined; + keepAlive?: boolean; + maxLifetimeSeconds: number | undefined; + options?: string; + query_timeout?: number; + statement_timeout?: number | false; +}>; +``` + +### hooks + +```ts +hooks: { + afterConnect: (connection: unknown) => void; +}; +``` + +#### hooks.afterConnect() + +```ts +afterConnect: (connection: unknown) => void = lifetimeCheck.afterConnect; +``` + +##### Parameters + +| Parameter | Type | +| ------ | ------ | +| `connection` | `unknown` | + +##### Returns + +`void` + ### password ```ts password: string | () => string | () => Promise | undefined; ``` +### pool + +```ts +pool: DefinedOnly<{ + idle: number | undefined; + max: number | undefined; + validate: (connection?: unknown) => boolean; +}>; +``` + +### poolSize + +```ts +poolSize: number | undefined = pgConfig.max; +``` + ### ssl ```ts @@ -55,7 +226,7 @@ ssl: | boolean | { rejectUnauthorized: boolean | undefined; -}; +} = ormSsl; ``` ### username @@ -81,4 +252,12 @@ const sequelize = new Sequelize({ ...getLakebaseOrmConfig(), logging: false, }); + +// Adding your own settings: merge, don't replace +const ormConfig = getLakebaseOrmConfig(); +const sequelize = new Sequelize({ + dialect: 'postgres', + ...ormConfig, + pool: { ...ormConfig.pool, min: 2 }, +}); ``` diff --git a/packages/lakebase/README.md b/packages/lakebase/README.md index 99f3beb60..32cb772c8 100644 --- a/packages/lakebase/README.md +++ b/packages/lakebase/README.md @@ -214,9 +214,7 @@ const users = await prisma.user.findMany(); ```typescript import { DataSource } from "typeorm"; -import { createLakebasePool } from "@databricks/lakebase"; - -const pool = createLakebasePool(); +import { getLakebaseOrmConfig } from "@databricks/lakebase"; const dataSource = new DataSource({ type: "postgres", @@ -242,6 +240,34 @@ const sequelize = new Sequelize({ }); ``` +`getLakebaseOrmConfig()` returns the pool and client settings (connection lifetime, pool size, timeouts, SSL) where each ORM reads them: + +- **TypeORM:** pool size, connect timeout and application name go under TypeORM's own `poolSize`, `connectTimeoutMS` and `applicationName`, so values you set after the spread still win. The rest (connection lifetime, idle timeout, statement and query timeouts) goes under `extra`, which TypeORM merges into the `pg.Pool` options. +- **Sequelize:** `dialectOptions` (passed to each `pg.Client`, including SSL), `pool` and `hooks`. Sequelize's pool has no lifetime option, so connections are recycled by age through `pool.validate` and an `afterConnect` hook. + +The result carries keys for both ORMs, and each ignores the other's: Sequelize ignores `poolSize`, `connectTimeoutMS`, `applicationName` and `extra`, and TypeORM ignores `dialectOptions`, `pool` and `hooks`. + +If you set `extra`, `pool`, `dialectOptions` or `hooks` yourself, merge with the returned values instead of replacing them, or those settings are lost. Replacing `hooks` makes connection lifetimes count from a connection's first reuse instead of when it opened, and logs a warning if you passed a `logger` (nothing is logged without one). + +```typescript +const ormConfig = getLakebaseOrmConfig(); + +// TypeORM +const dataSource = new DataSource({ + type: "postgres", + ...ormConfig, + extra: { ...ormConfig.extra, allowExitOnIdle: true }, +}); + +// Sequelize +const sequelize = new Sequelize({ + dialect: "postgres", + ...ormConfig, + pool: { ...ormConfig.pool, min: 2 }, + hooks: { ...ormConfig.hooks, beforeConnect: async () => {} }, +}); +``` + ## OpenTelemetry Integration The driver automatically uses OpenTelemetry's global registry when available. If your application initializes OpenTelemetry providers, the driver will automatically instrument queries and metrics with no additional configuration needed. diff --git a/packages/lakebase/package.json b/packages/lakebase/package.json index 47b77854b..506ee25da 100644 --- a/packages/lakebase/package.json +++ b/packages/lakebase/package.json @@ -62,7 +62,9 @@ "@opentelemetry/api": "1.9.0" }, "devDependencies": { - "@types/pg": "8.16.0" + "@types/pg": "8.16.0", + "sequelize": "6.37.7", + "typeorm": "0.3.28" }, "publishConfig": { "exports": { diff --git a/packages/lakebase/src/__tests__/orm-config.test.ts b/packages/lakebase/src/__tests__/orm-config.test.ts new file mode 100644 index 000000000..b9c470ff9 --- /dev/null +++ b/packages/lakebase/src/__tests__/orm-config.test.ts @@ -0,0 +1,404 @@ +import { EventEmitter } from "node:events"; + +import pg from "pg"; +import { type Options, type PoolOptions, Sequelize } from "sequelize"; +import { DataSource } from "typeorm"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { getLakebaseOrmConfig } from "../pool-config"; +import type { LakebasePoolConfig, Logger } from "../types"; + +// These tests run the real ORMs against a fake `pg` module, so they show what +// TypeORM and Sequelize actually pass to pg, without a database. + +// config.ts falls back to these, so clear them to keep a developer's shell +// from changing the result. +const ENV_KEYS = [ + "PGHOST", + "PGPORT", + "PGDATABASE", + "PGUSER", + "PGSSLMODE", + "LAKEBASE_ENDPOINT", + "DATABRICKS_CLIENT_ID", +] as const; + +beforeEach(() => { + for (const key of ENV_KEYS) { + vi.stubEnv(key, undefined); + } +}); + +afterEach(() => { + vi.unstubAllEnvs(); + vi.useRealTimers(); +}); + +function ormConfig(config: Partial = {}) { + return getLakebaseOrmConfig({ + host: "ep-test.database.us-east-1.databricks.com", + database: "databricks_postgres", + user: "test-user@example.com", + password: "test-only", + ...config, + }); +} + +describe("TypeORM", () => { + // The options TypeORM passes to pg.Pool, given a Lakebase config and any + // TypeORM options the caller sets after the spread. + async function poolOptionsFor( + config: Partial, + typeormOptions: Record = {}, + ): Promise { + let captured: pg.PoolConfig | undefined; + + // Real pg.Pool, so `options` is what pg-pool resolved, but it stops before + // a real connection opens. TypeORM calls `new Pool(...)`; a constructor + // that returns an object gives back that object. + function CapturingPool(options: pg.PoolConfig): pg.Pool { + const pool = new pg.Pool(options); + captured = pool.options; + return Object.assign(pool, { + connect: (callback: (err: Error) => void) => + callback(new Error("test: no database")), + }); + } + + const dataSource = new DataSource({ + type: "postgres", + ...ormConfig(config), + ...typeormOptions, + driver: { ...pg, Pool: CapturingPool as unknown as typeof pg.Pool }, + }); + await expect(dataSource.initialize()).rejects.toThrow("test: no database"); + + if (!captured) { + throw new Error("TypeORM didn't create a pg.Pool"); + } + return captured; + } + + test("passes the connection lifetime to pg.Pool", async () => { + const options = await poolOptionsFor({}); + expect(options.maxLifetimeSeconds).toBe(30 * 60); + }); + + test("passes an explicit lifetime, pool size and timeouts to pg.Pool", async () => { + const options = await poolOptionsFor({ + maxLifetimeSeconds: 600, + max: 7, + idleTimeoutMillis: 5_000, + connectionTimeoutMillis: 2_000, + statement_timeout: 3_000, + }); + + expect(options).toMatchObject({ + maxLifetimeSeconds: 600, + max: 7, + idleTimeoutMillis: 5_000, + connectionTimeoutMillis: 2_000, + statement_timeout: 3_000, + }); + }); + + // TypeORM merges `extra` over its own options, so these must not go in + // `extra`, or they'd override what the caller sets. + test("lets the caller's poolSize, connectTimeoutMS and applicationName win", async () => { + const options = await poolOptionsFor( + { max: 7, connectionTimeoutMillis: 2_000, application_name: "lakebase" }, + { poolSize: 20, connectTimeoutMS: 4_000, applicationName: "mine" }, + ); + + expect(options).toMatchObject({ + max: 20, + connectionTimeoutMillis: 4_000, + application_name: "mine", + }); + }); + + test("passes the Lakebase pool size, connect timeout and application name when the caller doesn't set them", async () => { + const options = await poolOptionsFor({ + max: 7, + connectionTimeoutMillis: 2_000, + application_name: "lakebase", + }); + + expect(options).toMatchObject({ + max: 7, + connectionTimeoutMillis: 2_000, + application_name: "lakebase", + }); + }); + + test("keeps the caller's own extra settings when they merge as documented", async () => { + const config = ormConfig({ maxLifetimeSeconds: 600 }); + const options = await poolOptionsFor( + { maxLifetimeSeconds: 600 }, + { extra: { ...config.extra, allowExitOnIdle: true } }, + ); + + expect(options).toMatchObject({ + maxLifetimeSeconds: 600, + allowExitOnIdle: true, + }); + }); +}); + +describe("Sequelize", () => { + // A pg.Client stand-in: connects at once and answers every query with no + // rows. + class FakeClient extends EventEmitter { + static instances: FakeClient[] = []; + + readonly connection = new EventEmitter(); + // Saves Sequelize a "SET standard_conforming_strings" round trip. + readonly standard_conforming_strings = "on"; + ended = false; + + constructor(readonly config: pg.ClientConfig) { + super(); + FakeClient.instances.push(this); + } + + connect(callback: (err: Error | null) => void): void { + setImmediate(() => callback(null)); + } + + async query(): Promise<{ rows: never[]; rowCount: number }> { + return { rows: [], rowCount: 0 }; + } + + end(callback?: () => void): void { + this.ended = true; + this.emit("end"); + callback?.(); + } + } + + beforeEach(() => { + FakeClient.instances = []; + }); + + function makeSequelize( + config: Partial = {}, + sequelizeOptions: Record = {}, + ) { + // Cast: Sequelize's types don't list `databaseVersion`, or a function for + // `password` (both work at runtime). + return new Sequelize({ + dialect: "postgres", + dialectModule: { types: pg.types, Client: FakeClient }, + // Skips the server-version query. + databaseVersion: "16.0.0", + logging: false, + ...ormConfig(config), + ...sequelizeOptions, + } as Options); + } + + test("uses the pool size from the Lakebase config", () => { + const sequelize = makeSequelize({ max: 7 }); + // `pool` is sequelize-pool's Pool; Sequelize's types leave it out. + const { pool } = sequelize.connectionManager as unknown as { + pool: { maxSize: number }; + }; + expect(pool.maxSize).toBe(7); + }); + + test("replaces an idle connection once it's older than the lifetime", async () => { + vi.useFakeTimers({ toFake: ["Date"] }); + const sequelize = makeSequelize({ maxLifetimeSeconds: 600 }); + const { connectionManager } = sequelize; + + const first = await connectionManager.getConnection({ type: "write" }); + connectionManager.releaseConnection(first); + + vi.advanceTimersByTime(599_000); + const reused = await connectionManager.getConnection({ type: "write" }); + expect(reused).toBe(first); + connectionManager.releaseConnection(reused); + + vi.advanceTimersByTime(2_000); + const replacement = await connectionManager.getConnection({ + type: "write", + }); + expect(replacement).not.toBe(first); + expect((first as unknown as FakeClient).ended).toBe(true); + + connectionManager.releaseConnection(replacement); + await sequelize.close(); + }); + + test("never replaces a connection for age when the lifetime is 0", async () => { + vi.useFakeTimers({ toFake: ["Date"] }); + const sequelize = makeSequelize({ maxLifetimeSeconds: 0 }); + const { connectionManager } = sequelize; + + const first = await connectionManager.getConnection({ type: "write" }); + connectionManager.releaseConnection(first); + + vi.advanceTimersByTime(24 * 60 * 60_000); + const reused = await connectionManager.getConnection({ type: "write" }); + expect(reused).toBe(first); + + connectionManager.releaseConnection(reused); + await sequelize.close(); + }); + + test("still drops a connection Sequelize marked invalid", async () => { + const sequelize = makeSequelize(); + const { connectionManager } = sequelize; + + const first = await connectionManager.getConnection({ type: "write" }); + (first as unknown as { _invalid: boolean })._invalid = true; + connectionManager.releaseConnection(first); + + const next = await connectionManager.getConnection({ type: "write" }); + expect(next).not.toBe(first); + + connectionManager.releaseConnection(next); + await sequelize.close(); + }); + async function clientConfigFor(config: Partial) { + const sequelize = makeSequelize(config); + const connection = await sequelize.connectionManager.getConnection({ + type: "write", + }); + sequelize.connectionManager.releaseConnection(connection); + await sequelize.close(); + return FakeClient.instances[0].config; + } + + // Before, Sequelize never received `ssl` or the client timeouts: pg fell + // back to PGSSLMODE and its own defaults. Each sslMode now reaches + // pg.Client the same way it reaches pg.Pool, along with the timeouts. + test.each([ + ["require", { rejectUnauthorized: true }], + ["prefer", { rejectUnauthorized: false }], + ["disable", false], + ] as const)( + "passes sslMode %s and the client timeouts to pg.Client", + async (sslMode, expected) => { + const config = await clientConfigFor({ + sslMode, + statement_timeout: 3_000, + connectionTimeoutMillis: 2_000, + }); + expect(config.ssl).toEqual(expected); + expect(config).toMatchObject({ + statement_timeout: 3_000, + connectionTimeoutMillis: 2_000, + }); + }, + ); + + test("keeps recycling, from first reuse, and warns once when the caller replaces hooks", async () => { + vi.useFakeTimers({ toFake: ["Date"] }); + const logger: Logger = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + const sequelize = makeSequelize( + { maxLifetimeSeconds: 600, logger }, + { hooks: { beforeConnect: () => {} } }, + ); + const { connectionManager } = sequelize; + + const first = await connectionManager.getConnection({ type: "write" }); + connectionManager.releaseConnection(first); + + // First reuse: no open time was recorded, so the clock starts here. + const reused = await connectionManager.getConnection({ type: "write" }); + expect(reused).toBe(first); + connectionManager.releaseConnection(reused); + expect(logger.warn).toHaveBeenCalledTimes(1); + + vi.advanceTimersByTime(601_000); + const replacement = await connectionManager.getConnection({ + type: "write", + }); + expect(replacement).not.toBe(first); + // The replacement isn't stamped either, but the warning isn't repeated. + connectionManager.releaseConnection(replacement); + const again = await connectionManager.getConnection({ type: "write" }); + expect(again).toBe(replacement); + expect(logger.warn).toHaveBeenCalledTimes(1); + + connectionManager.releaseConnection(again); + await sequelize.close(); + }); + + // Compile-time check: the documented merge pattern typechecks against + // Sequelize's own option types (typecheck covers this file). + test("the documented merge pattern fits Sequelize's types", () => { + const config = ormConfig(); + const pool: PoolOptions = { ...config.pool, min: 2 }; + const hooks: Options["hooks"] = { + ...config.hooks, + beforeConnect: async () => {}, + }; + + expect(pool.validate).toBe(config.pool.validate); + expect(hooks?.afterConnect).toBe(config.hooks.afterConnect); + }); +}); + +describe("logging", () => { + // Fails every token fetch: any use of the client throws. + const failingClient = new Proxy( + {}, + { + get() { + throw new Error("test: token fetch failed"); + }, + }, + ) as LakebasePoolConfig["workspaceClient"]; + + // An invalid lifetime (which warns) and an OAuth password callback whose + // token fetch fails (which logs an error). + function noisyConfig(logger?: Logger) { + return ormConfig({ + password: undefined, + endpoint: "projects/p/branches/b/endpoints/e", + workspaceClient: failingClient, + maxLifetimeSeconds: -5, + logger, + }); + } + + test("doesn't log anything when no logger is configured", async () => { + const consoleSpies = ( + ["error", "warn", "info", "log", "debug"] as const + ).map((method) => vi.spyOn(console, method).mockImplementation(() => {})); + + const config = noisyConfig(); + await expect((config.password as () => Promise)()).rejects.toThrow( + "test: token fetch failed", + ); + + for (const spy of consoleSpies) { + expect(spy).not.toHaveBeenCalled(); + spy.mockRestore(); + } + }); + + test("logs warnings and token-fetch errors to a configured logger", async () => { + const logger: Logger = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + + const config = noisyConfig(logger); + await expect((config.password as () => Promise)()).rejects.toThrow( + "test: token fetch failed", + ); + + expect(logger.warn).toHaveBeenCalled(); + expect(logger.error).toHaveBeenCalled(); + }); +}); diff --git a/packages/lakebase/src/pool-config.ts b/packages/lakebase/src/pool-config.ts index 400b85a2c..1fa0dd5b0 100644 --- a/packages/lakebase/src/pool-config.ts +++ b/packages/lakebase/src/pool-config.ts @@ -1,6 +1,7 @@ import type { PoolConfig } from "pg"; import { getUsernameSync, parsePoolConfig } from "./config"; +import { resolveLogger } from "./logger"; import { type DriverTelemetry, initTelemetry } from "./telemetry"; import { createTokenRefreshCallback } from "./token-refresh"; import type { LakebasePoolConfig, Logger } from "./types"; @@ -152,6 +153,82 @@ export function getLakebasePgConfig( }; } +/** + * `T` after `definedOnly`: keys whose value can be undefined become optional + * (and no longer include undefined); keys that are always set stay required. + */ +type DefinedOnly = { + [K in keyof T as undefined extends T[K] ? never : K]: T[K]; +} & { + [K in keyof T as undefined extends T[K] ? K : never]?: Exclude< + T[K], + undefined + >; +}; + +/** Drop undefined values, so they don't override an ORM's own defaults. */ +function definedOnly(values: T): DefinedOnly { + return Object.fromEntries( + Object.entries(values).filter(([, value]) => value !== undefined), + ) as DefinedOnly; +} + +/** + * Age-based connection recycling for Sequelize, whose pool (sequelize-pool) + * has no lifetime option. + * + * `afterConnect` records when each connection opens. sequelize-pool calls + * `validate` whenever it reuses an idle connection and destroys the + * connection when it returns false, so a connection older than the lifetime + * is replaced by a new one. As with pg-pool, a connection that's checked out + * isn't closed until it's released. + */ +function createSequelizeLifetimeCheck( + maxLifetimeSeconds: number, + logger?: Logger, +) { + const openedAt = new WeakMap(); + const lifetimeMs = maxLifetimeSeconds * 1000; + let warnedMissingHook = false; + + return { + afterConnect(connection: unknown): void { + if (typeof connection === "object" && connection !== null) { + openedAt.set(connection, Date.now()); + } + }, + validate(connection?: unknown): boolean { + if (typeof connection !== "object" || connection === null) { + return false; + } + // A custom validate replaces Sequelize's own check, so repeat it. + const client = connection as { _invalid?: boolean; _ending?: boolean }; + if (client._invalid || client._ending) { + return false; + } + if (lifetimeMs === 0) { + return true; + } + const opened = openedAt.get(connection); + if (opened === undefined) { + // afterConnect didn't run, so the caller replaced `hooks` with their + // own. Start the clock at this first reuse instead. Rejecting the + // connection would be worse: no connection would ever have an open + // time, so every reuse would open a new one. + if (!warnedMissingHook) { + warnedMissingHook = true; + logger?.warn( + "The afterConnect hook from getLakebaseOrmConfig() isn't registered, so connection lifetimes are measured from first reuse rather than from when connections open. Merge your hooks with the returned `hooks` instead of replacing them.", + ); + } + openedAt.set(connection, Date.now()); + return true; + } + return Date.now() - opened < lifetimeMs; + }, + }; +} + /** * Get Lakebase connection configuration for ORMs that don't accept pg.Pool directly. * @@ -163,6 +240,27 @@ export function getLakebasePgConfig( * - Simplified SSL config * - Password callback support for OAuth token refresh * + * Pool and client settings (connection lifetime, pool size, timeouts, SSL) + * are also returned where each ORM reads them, so recycling works with both: + * - TypeORM: pool size, connect timeout and application name under its own + * `poolSize`, `connectTimeoutMS` and `applicationName`, so values you set + * after the spread still win. Everything else (lifetime, idle timeout, + * statement and query timeouts) goes under `extra`, which TypeORM merges + * into the `pg.Pool` options last. + * - Sequelize: `dialectOptions` (passed to each `pg.Client`), `pool`, and + * `hooks`. Sequelize's pool has no lifetime option, so connections are + * recycled by age through `pool.validate` and an `afterConnect` hook. + * + * The result carries keys for both ORMs, and each ignores the other's: + * Sequelize ignores `poolSize`, `connectTimeoutMS`, `applicationName` and + * `extra`, and TypeORM ignores `dialectOptions`, `pool` and `hooks`. + * + * If you set your own `extra`, `pool`, + * `dialectOptions` or `hooks`, merge them with these instead of replacing + * them, or those settings are lost. Replacing `hooks` makes lifetimes count + * from a connection's first reuse instead of when it opened, and logs a + * warning if you passed a `logger`. Nothing is logged without one. + * * @param config - Optional configuration (reads from environment if not provided) * @returns ORM-compatible connection configuration * @@ -182,10 +280,48 @@ export function getLakebasePgConfig( * ...getLakebaseOrmConfig(), * logging: false, * }); + * + * // Adding your own settings: merge, don't replace + * const ormConfig = getLakebaseOrmConfig(); + * const sequelize = new Sequelize({ + * dialect: 'postgres', + * ...ormConfig, + * pool: { ...ormConfig.pool, min: 2 }, + * }); * ``` */ export function getLakebaseOrmConfig(config?: Partial) { - const { user, password, ssl, ...pgConfig } = getLakebasePgConfig(config); + // Log only when the caller configured a logger. resolveLogger(undefined) + // would return a console logger for errors, which would print token-fetch + // failures that this function has never printed. + const logger = config?.logger ? resolveLogger(config.logger) : undefined; + const { user, password, ssl, ...pgConfig } = getLakebasePgConfig( + config, + undefined, + logger, + ); + + const ormSsl = ssl + ? typeof ssl === "boolean" + ? ssl + : { rejectUnauthorized: ssl.rejectUnauthorized } + : false; + + // Settings pg.Client reads that TypeORM has no option of its own for. + // Neither ORM forwards these from the top level. + const clientOptions = definedOnly({ + keepAlive: pgConfig.keepAlive, + statement_timeout: pgConfig.statement_timeout, + query_timeout: pgConfig.query_timeout, + idle_in_transaction_session_timeout: + pgConfig.idle_in_transaction_session_timeout, + options: pgConfig.options, + }); + + const lifetimeCheck = createSequelizeLifetimeCheck( + pgConfig.maxLifetimeSeconds ?? 0, + logger, + ); return { ...pgConfig, @@ -195,10 +331,31 @@ export function getLakebaseOrmConfig(config?: Partial) { | (() => string) | (() => Promise) | undefined, - ssl: ssl - ? typeof ssl === "boolean" - ? ssl - : { rejectUnauthorized: ssl.rejectUnauthorized } - : false, + ssl: ormSsl, + // TypeORM. Settings it has its own option for go under that option, so + // a value the caller sets after the spread still wins: TypeORM applies + // `extra` last, over its own options. + poolSize: pgConfig.max, + connectTimeoutMS: pgConfig.connectionTimeoutMillis, + applicationName: pgConfig.application_name, + extra: definedOnly({ + ...clientOptions, + idleTimeoutMillis: pgConfig.idleTimeoutMillis ?? undefined, + maxLifetimeSeconds: pgConfig.maxLifetimeSeconds, + }), + // Sequelize + dialectOptions: definedOnly({ + ...clientOptions, + application_name: pgConfig.application_name, + connectionTimeoutMillis: pgConfig.connectionTimeoutMillis, + ssl: ormSsl, + }), + pool: definedOnly({ + max: pgConfig.max, + // @types/pg allows null here; Sequelize's types don't. + idle: pgConfig.idleTimeoutMillis ?? undefined, + validate: lifetimeCheck.validate, + }), + hooks: { afterConnect: lifetimeCheck.afterConnect }, }; } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 12413586b..c4388175e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -644,6 +644,12 @@ importers: '@types/pg': specifier: 8.16.0 version: 8.16.0 + sequelize: + specifier: 6.37.7 + version: 6.37.7(pg@8.18.0) + typeorm: + specifier: 0.3.28 + version: 0.3.28(pg@8.18.0) packages/shared: dependencies: