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

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<LakebasePoolConfig> = {}): 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<typeof vi.fn> } {
+    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();
+  });
+});

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();
+    }
+  });
+});

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<LakebasePoolConfig>,
+  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,
   };
 }
 

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
    *