diff --git a/.github/workflows/prepare-release-lakebase-auth.yml b/.github/workflows/prepare-release-lakebase-auth.yml new file mode 100644 index 000000000..9acad53be --- /dev/null +++ b/.github/workflows/prepare-release-lakebase-auth.yml @@ -0,0 +1,110 @@ +name: Prepare Release Lakebase Auth + +on: + push: + branches: + - main + paths: + - 'packages/lakebase-auth/**' + +concurrency: + group: prepare-release-lakebase-auth + cancel-in-progress: false + +permissions: + contents: read + id-token: write + +jobs: + prepare: + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + + steps: + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + fetch-depth: 0 + + - name: Setup JFrog npm + uses: ./.github/actions/setup-jfrog-npm + + - name: Setup pnpm + uses: pnpm/action-setup@fc06bc1257f339d1d5d8b3a19a8cae5388b55320 # v5.0.0 + + - name: Setup Node.js + uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: 24 + cache: "pnpm" + + - name: Install dependencies + run: pnpm install --frozen-lockfile + + - name: Check for releasable commits + id: version + working-directory: packages/lakebase-auth + run: | + VERSION=$(pnpm exec release-it --release-version --ci) || true + if [[ "$VERSION" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then + echo "version=$VERSION" >> "$GITHUB_OUTPUT" + echo "Next version: $VERSION" + else + echo "No releasable commits — skipping release preparation" + echo "version=" >> "$GITHUB_OUTPUT" + fi + + - name: Generate changelog + if: steps.version.outputs.version != '' + working-directory: packages/lakebase-auth + run: | + pnpm exec release-it ${{ steps.version.outputs.version }} --ci + + - name: Sync version + if: steps.version.outputs.version != '' + run: pnpm exec tsx tools/sync-lakebase-auth-version.ts "${{ steps.version.outputs.version }}" + + - name: Build + if: steps.version.outputs.version != '' + run: pnpm --filter=@databricks/lakebase-auth build:package + + - name: Dist + if: steps.version.outputs.version != '' + run: pnpm --filter=@databricks/lakebase-auth dist + + - name: SBOM + if: steps.version.outputs.version != '' + run: pnpm --filter=@databricks/lakebase-auth release:sbom + + - name: Pack + if: steps.version.outputs.version != '' + run: npm pack packages/lakebase-auth/tmp + + - name: Generate SHA256 + if: steps.version.outputs.version != '' + run: sha256sum *.tgz > SHA256SUMS + + - name: Write version file + if: steps.version.outputs.version != '' + run: echo "${{ steps.version.outputs.version }}" > VERSION + + - name: Upload release metadata + if: steps.version.outputs.version != '' + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: lakebase-auth-release-meta-${{ github.run_number }} + retention-days: 7 + path: VERSION + + - name: Upload release artifacts + if: steps.version.outputs.version != '' + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: lakebase-auth-release-${{ github.run_number }} + retention-days: 7 + path: | + *.tgz + packages/lakebase-auth/changelog-diff.md + VERSION + SHA256SUMS diff --git a/.gitignore b/.gitignore index 645f5cf52..1ec112f60 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,5 @@ coverage .databricks .claude/scheduled_tasks.lock + +.cursor/ diff --git a/CLAUDE.md b/CLAUDE.md index 8e400e492..870dbda75 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -28,7 +28,8 @@ Examples: /packages/ /appkit/ - Core SDK with plugin architecture /appkit-ui/ - React components and JS utilities - /lakebase/ - Standalone Lakebase (PostgreSQL) connector package + /lakebase/ - Standalone Lakebase (PostgreSQL) connector package (pg.Pool + OTel) + /lakebase-auth/ - Driver-agnostic Lakebase OAuth credential/token-refresh package (SDK-only deps) /shared/ - Shared TypeScript types across packages /apps/ @@ -144,7 +145,7 @@ pnpm clean:full # Remove build artifacts + node_modules ### Releasing -This project uses a two-stage release pipeline. Both packages (`appkit` and `appkit-ui`) are always released together with the same version. `@databricks/lakebase` is released independently. +This project uses a two-stage release pipeline. Both packages (`appkit` and `appkit-ui`) are always released together with the same version. `@databricks/lakebase` and `@databricks/lakebase-auth` are each released independently. #### Stage 1: Prepare (this repo) @@ -154,7 +155,7 @@ The `prepare-release` workflow runs automatically on push to `main`: 3. Builds, packs, and uploads artifacts (`.tgz`, changelog, SHA256 digests) 4. **Does NOT** commit, tag, push, or publish — only uploads artifacts -Lakebase has a separate `prepare-release-lakebase` workflow triggered by changes to `packages/lakebase/**`. +Lakebase has a separate `prepare-release-lakebase` workflow triggered by changes to `packages/lakebase/**`, and `@databricks/lakebase-auth` has a `prepare-release-lakebase-auth` workflow triggered by changes to `packages/lakebase-auth/**`. #### Stage 2: Publish (secure repo) @@ -248,10 +249,11 @@ The AnalyticsPlugin provides SQL query execution: ### Lakebase Connector -Lakebase support is split into two layers: +Lakebase support is split into three layers: -1. **`@databricks/lakebase` package** (`packages/lakebase/`) - Standalone connector with OAuth token refresh, ORM helpers, and full API. See the [`@databricks/lakebase` README](https://github.com/databricks/appkit/blob/main/packages/lakebase/README.md). -2. **AppKit integration** (`packages/appkit/src/connectors/lakebase/`) - Thin wrapper that adds AppKit logger integration and re-exports the standalone package. +1. **`@databricks/lakebase-auth` package** (`packages/lakebase-auth/`) - Driver-agnostic OAuth credential generation and token refresh (eager by default, plus lazy; retries transient failures). Dependency-light (Databricks SDK only, no `pg`/OTel). Exposes `getPgConfig()` and the low-level `createPasswordProvider()` for use with `pg`, `postgres.js`, `Bun.SQL`, etc. See the [`@databricks/lakebase-auth` README](https://github.com/databricks/appkit/blob/main/packages/lakebase-auth/README.md). +2. **`@databricks/lakebase` package** (`packages/lakebase/`) - Builds on `@databricks/lakebase-auth` to provide a ready-to-use `pg.Pool` (`createLakebasePool`) with OpenTelemetry instrumentation and logger integration, plus ORM helpers. See the [`@databricks/lakebase` README](https://github.com/databricks/appkit/blob/main/packages/lakebase/README.md). +3. **AppKit integration** (`packages/appkit/src/connectors/lakebase/`) - Thin wrapper that adds AppKit logger integration and re-exports the standalone package. **Quick Example:** ```typescript diff --git a/docs/docs/api/appkit/Function.getUsernameWithApiLookup.md b/docs/docs/api/appkit/Function.getUsernameWithApiLookup.md index 29b92f2ec..64504a827 100644 --- a/docs/docs/api/appkit/Function.getUsernameWithApiLookup.md +++ b/docs/docs/api/appkit/Function.getUsernameWithApiLookup.md @@ -1,7 +1,7 @@ # Function: getUsernameWithApiLookup() ```ts -function getUsernameWithApiLookup(config?: Partial): Promise; +function getUsernameWithApiLookup(config?: Partial): Promise; ``` Resolves the PostgreSQL username for a Lakebase connection. @@ -24,7 +24,7 @@ caller can decide whether to proceed or surface an error. | Parameter | Type | | ------ | ------ | -| `config?` | `Partial`\<[`LakebasePoolConfig`](Interface.LakebasePoolConfig.md)\> | +| `config?` | `Partial`\<`LakebaseAuthConfig`\> | ## Returns diff --git a/docs/docs/api/appkit/Function.getWorkspaceClient.md b/docs/docs/api/appkit/Function.getWorkspaceClient.md index 9511b944d..00bb45373 100644 --- a/docs/docs/api/appkit/Function.getWorkspaceClient.md +++ b/docs/docs/api/appkit/Function.getWorkspaceClient.md @@ -1,7 +1,7 @@ # Function: getWorkspaceClient() ```ts -function getWorkspaceClient(config: Partial): WorkspaceClient; +function getWorkspaceClient(config: Partial): WorkspaceClient; ``` Get workspace client from config or SDK default auth chain @@ -10,7 +10,7 @@ Get workspace client from config or SDK default auth chain | Parameter | Type | | ------ | ------ | -| `config` | `Partial`\<[`LakebasePoolConfig`](Interface.LakebasePoolConfig.md)\> | +| `config` | `Partial`\<`LakebaseAuthConfig`\> | ## Returns diff --git a/docs/docs/api/appkit/Interface.DatabaseCredential.md b/docs/docs/api/appkit/Interface.DatabaseCredential.md index b2a0b255e..aeb12358a 100644 --- a/docs/docs/api/appkit/Interface.DatabaseCredential.md +++ b/docs/docs/api/appkit/Interface.DatabaseCredential.md @@ -1,6 +1,6 @@ # Interface: DatabaseCredential -Database credentials with OAuth token for Postgres connection +Database credentials with OAuth token for Postgres connection. ## Properties diff --git a/docs/docs/api/appkit/Interface.LakebasePoolConfig.md b/docs/docs/api/appkit/Interface.LakebasePoolConfig.md index 3f40f3021..ea95de865 100644 --- a/docs/docs/api/appkit/Interface.LakebasePoolConfig.md +++ b/docs/docs/api/appkit/Interface.LakebasePoolConfig.md @@ -18,6 +18,33 @@ https://docs.databricks.com/aws/en/oltp/projects/authentication ## Properties +### claims? + +```ts +optional claims: RequestedClaims[]; +``` + +Optional UC claims for fine-grained Unity Catalog table permissions on the +generated Postgres token. + +*** + +### earlyRefreshMs? + +```ts +optional earlyRefreshMs: number; +``` + +How long before token expiry to refresh, in milliseconds. + +#### Default + +```ts +120000 (2 minutes) +``` + +*** + ### endpoint? ```ts @@ -73,19 +100,57 @@ const pool = createLakebasePool({ *** +### refresh? + +```ts +optional refresh: RefreshMode; +``` + +Token refresh strategy. + +- `"eager"` (default): fetch a token immediately and refresh it in the + background before it expires. Best for time-sensitive, user-facing apps. +- `"lazy"`: fetch a token on first use and refresh it on demand. + +#### Default + +```ts +"eager" +``` + +*** + +### retry? + +```ts +optional retry: RetryOptions; +``` + +Retry options for transient credential-fetch failures (e.g. the OAuth +server being briefly unreachable). + +#### Default + +```ts +{ schedule: [50, 500, 5000] } +``` + +*** + ### sslMode? ```ts -optional sslMode: "require" | "disable" | "prefer"; +optional sslMode: "verify-full" | "verify-ca" | "require" | "prefer" | "disable"; ``` -SSL mode for the connection (convenience helper) -Can also be set via PGSSLMODE environment variable +SSL mode for the connection (convenience helper). Can also be set via +PGSSLMODE. All values other than "disable" are treated as "verify-full" +with system root certs. #### Default ```ts -"require" +"verify-full" ``` *** diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index 4d9950512..0849194f9 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -42,7 +42,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [AutoInheritToolsConfig](Interface.AutoInheritToolsConfig.md) | Auto-inherit configuration. When enabled for a given agent origin, agents with no explicit `tools:` declaration receive every registered ToolProvider plugin tool whose author marked `autoInheritable: true`. Tools without that flag — destructive, state-mutating, or privilege-sensitive — never spread automatically and must be wired via `tools:` (object or function form in code, `plugin:NAME` entries in markdown frontmatter). | | [BasePluginConfig](Interface.BasePluginConfig.md) | Base configuration interface for AppKit plugins | | [CacheConfig](Interface.CacheConfig.md) | Configuration for the CacheInterceptor. Controls TTL, size limits, storage backend, and probabilistic cleanup. | -| [DatabaseCredential](Interface.DatabaseCredential.md) | Database credentials with OAuth token for Postgres connection | +| [DatabaseCredential](Interface.DatabaseCredential.md) | Database credentials with OAuth token for Postgres connection. | | [EndpointConfig](Interface.EndpointConfig.md) | - | | [FilePolicyUser](Interface.FilePolicyUser.md) | Minimal user identity passed to the policy function. | | [FileResource](Interface.FileResource.md) | Describes the file or directory being acted upon. | diff --git a/docs/docs/plugins/lakebase.md b/docs/docs/plugins/lakebase.md index d3f77e78d..d617078c9 100644 --- a/docs/docs/plugins/lakebase.md +++ b/docs/docs/plugins/lakebase.md @@ -8,11 +8,14 @@ Provides a PostgreSQL connection pool for Databricks Lakebase Autoscaling with a **Key features:** - Standard `pg.Pool` compatible with any PostgreSQL library or ORM -- Automatic OAuth token refresh (1-hour tokens, 2-minute refresh buffer) +- Automatic OAuth token refresh (1-hour tokens, 2-minute refresh buffer) — eager (background) by default, or lazy (on-demand) +- Retries transient credential-fetch failures - Token caching to minimize API calls - Built-in OpenTelemetry instrumentation (query duration, pool connections, token refresh) - AppKit logger configured by default for query and connection events +The underlying OAuth credential generation and token-refresh logic lives in the lightweight, driver-agnostic [`@databricks/lakebase-auth`](https://www.npmjs.com/package/@databricks/lakebase-auth) package, which can also be used standalone with `pg`, `postgres.js`, or `Bun.SQL`. + ## Getting started with the Lakebase The easiest way to get started with the Lakebase plugin is to use the Databricks CLI to create a new Databricks app with AppKit installed and the Lakebase plugin. diff --git a/packages/appkit/src/plugins/lakebase/manifest.json b/packages/appkit/src/plugins/lakebase/manifest.json index 17798d33f..69919f237 100644 --- a/packages/appkit/src/plugins/lakebase/manifest.json +++ b/packages/appkit/src/plugins/lakebase/manifest.json @@ -74,7 +74,7 @@ "sslmode": { "env": "PGSSLMODE", "localOnly": true, - "value": "require", + "value": "verify-full", "description": "Postgres SSL mode. Auto-injected by the platform at deploy time." } } diff --git a/packages/lakebase-auth/.release-it.json b/packages/lakebase-auth/.release-it.json new file mode 100644 index 000000000..db7a363f0 --- /dev/null +++ b/packages/lakebase-auth/.release-it.json @@ -0,0 +1,32 @@ +{ + "$schema": "https://unpkg.com/release-it@19/schema/release-it.json", + "git": { + "commit": false, + "tag": false, + "push": false, + "requireBranch": false, + "requireCleanWorkingDir": false, + "requireCommits": true, + "requireCommitsFail": false, + "tagMatch": "lakebase-auth-v*", + "tagName": "lakebase-auth-v${version}", + "getLatestTagFromAllRefs": true, + "commitsPath": "." + }, + "github": { + "release": false + }, + "npm": false, + "hooks": {}, + "plugins": { + "@release-it/conventional-changelog": { + "preset": { + "name": "conventionalcommits", + "bumpStrict": true + }, + "infile": "changelog-diff.md", + "gitRawCommitsOpts": { "path": "." }, + "commitsOpts": { "path": "." } + } + } +} diff --git a/packages/lakebase-auth/CHANGELOG.md b/packages/lakebase-auth/CHANGELOG.md new file mode 100644 index 000000000..7726c9e52 --- /dev/null +++ b/packages/lakebase-auth/CHANGELOG.md @@ -0,0 +1,3 @@ +# Changelog + +All notable changes to @databricks/lakebase-auth will be documented in this file. diff --git a/packages/lakebase-auth/README.md b/packages/lakebase-auth/README.md new file mode 100644 index 000000000..526ceb0a0 --- /dev/null +++ b/packages/lakebase-auth/README.md @@ -0,0 +1,158 @@ +# @databricks/lakebase-auth + +OAuth credential generation and token refresh for Databricks Lakebase Autoscaling, usable with any PostgreSQL driver. + +## Overview + +`@databricks/lakebase-auth` produces a Postgres connection config (including an auto-refreshing OAuth password callback) for Databricks Lakebase Autoscaling (OLTP) databases. It is dependency-light (only the Databricks SDK) and driver-agnostic, so it works with [`pg`](https://node-postgres.com/), [`postgres.js`](https://github.com/porsager/postgres), [`Bun.SQL`](https://bun.sh/docs/api/sql), and any ORM built on top of them. + +It: + +- Generates time-limited Lakebase OAuth tokens via the Databricks SDK (full auth chain: `.databrickscfg`, PAT, OAuth M2M, env vars) +- Refreshes tokens automatically — **eagerly** in the background by default, or **lazily** on demand +- Retries transient credential-fetch failures +- Has no `pg` or OpenTelemetry dependency + +For a batteries-included `pg.Pool` with OpenTelemetry instrumentation and AppKit integration, use [`@databricks/lakebase`](https://www.npmjs.com/package/@databricks/lakebase), which builds on this package. + +**NOTE:** This package is NOT compatible with Databricks Lakebase Provisioned. + +## Installation + +```bash +npm install @databricks/lakebase-auth +``` + +## Quick Start + +Ensure Databricks credentials are available, for example in `.databrickscfg` or by setting `DATABRICKS_HOST`, `DATABRICKS_CLIENT_ID`, and `DATABRICKS_CLIENT_SECRET`. + +Set the following environment variables: + +```bash +export PGHOST=your-lakebase-host.databricks.com +export PGDATABASE=your_database_name +export LAKEBASE_ENDPOINT=projects/6bef4151-4b5d-4147-b4d0-c2f4fd5b40db/branches/br-broad-pine-y12n6gnv/endpoints/ep-summer-frost-y131l3vx +export PGUSER=your_user # optional: defaults to DATABRICKS_CLIENT_ID +``` + +Your `LAKEBASE_ENDPOINT` has the structure `projects/${project}/branches/${branch}/endpoints/${endpoint}`. To find it, run the Databricks CLI and use the `name` field: + +```bash +databricks postgres list-endpoints projects/{project-id}/branches/{branch-id} +``` + +You can obtain the Project ID and Branch ID from the Lakebase Autoscaling UI, like the "Branch Overview" page (Project list -> Project dashboard -> Branch overview). + +Then use with node-postgres: + +```typescript +import pg from "pg"; +import { getPgConfig } from "@databricks/lakebase-auth"; + +const { dispose, ...config } = getPgConfig(); +const pool = new pg.Pool(config); + +const result = await pool.query("SELECT * FROM users"); + +// on shutdown, stop the background token refresh: +await pool.end(); +dispose(); +``` + +## Usage with other drivers + +`getPgConfig()` returns `host`, `port`, `user`, `database`, `password` (a function returning a current OAuth token), `ssl`, and `dispose`. These are accepted by postgres.js and Bun.sql as well as pg: + +```typescript +// postgres.js +import postgres from "postgres"; +const { dispose, ...config } = getPgConfig(); +const sql = postgres({ + ...config, + // any custom options or overrides here +}); +const result = await sql`SELECT now()` +await sql.end(); +dispose(); // stop background token refresh + +// Bun.SQL +const { dispose, ...config } = getPgConfig(); +const sql = new Bun.SQL({ + ...config, + // any custom options or overrides here +}); +const result = await sql`SELECT now()`; +await sql.end(); +dispose(); // stop background token refresh +``` + +The emitted `ssl` object carries a `serverName` for `Bun.SQL`, which [fails to derive SNI from the host](https://github.com/oven-sh/bun/issues/26369) when TLS is passed as an object. Lakebase requires SNI, so this makes Bun connections work; `pg` and `postgres.js` set SNI themselves and ignore the key. + +### Low-level password provider + +If you only need the password callback (and manage the rest of the connection yourself), use `createPasswordProvider`: + +```typescript +import { createPasswordProvider } from "@databricks/lakebase-auth"; + +const { password, dispose } = createPasswordProvider({ + endpoint: process.env.LAKEBASE_ENDPOINT, +}); + +const pool = new pg.Pool({ host, user, database, password }); +// on shutdown: await pool.end(); dispose(); +``` + +## Token refresh strategies + +| Mode | When the token is fetched/refreshed | Best for | +| ---------------- | -------------------------------------------------------------------- | ----------------------------------------- | +| `"eager"` (default) | Immediately on creation, then in the background before each expiry | Time-sensitive, user-facing apps and APIs | +| `"lazy"` | On first use, then on demand when the token nears expiry | Background jobs, infrequent connections | + +```typescript +const config = getPgConfig({ refresh: "lazy" }); +``` + +Eager refresh uses an `unref`'d timer, so it never keeps the process alive on its own. Call `dispose()` to cancel it during graceful shutdown. + +## Retries + +Transient credential-fetch failures (e.g. the OAuth server being briefly unreachable) are retried automatically. The default schedule is `[50, 500, 5000]` ms (i.e. an initial attempt plus three retries with backoff). Customize or disable it: + +```typescript +// custom backoff +getPgConfig({ retry: { schedule: [100, 1000] } }); + +// disable retries +getPgConfig({ retry: { schedule: [] } }); +``` + +## Logging + +The package emits log events through an optional `onLog` callback (no logging dependency): + +```typescript +getPgConfig({ + onLog: (level, message, ...args) => console[level](message, ...args), +}); +``` + +## Configuration + +| Option | Environment Variable | Description | Default | +| ---------------- | ---------------------------------- | ------------------------------------ | ------------------ | +| `host` | `PGHOST` | Lakebase host | _Required_ | +| `database` | `PGDATABASE` | Database name | _Required_ | +| `endpoint` | `LAKEBASE_ENDPOINT` | Endpoint resource path | _Required_ | +| `user` | `PGUSER` or `DATABRICKS_CLIENT_ID` | Username or service principal ID | _Required_ | +| `port` | `PGPORT` | Port number | `5432` | +| `sslMode` | `PGSSLMODE` | SSL mode | `require` | +| `refresh` | - | `"eager"` or `"lazy"` | `"eager"` | +| `earlyRefreshMs` | - | How long before expiry to refresh | `120000` | +| `retry` | - | Retry schedule for credential fetch | `[50, 500, 5000]` | + +## Learn more + +For Lakebase Autoscaling documentation, see [docs.databricks.com/aws/en/oltp/projects](https://docs.databricks.com/aws/en/oltp/projects/). diff --git a/packages/lakebase-auth/package.json b/packages/lakebase-auth/package.json new file mode 100644 index 000000000..d227f3a64 --- /dev/null +++ b/packages/lakebase-auth/package.json @@ -0,0 +1,62 @@ +{ + "name": "@databricks/lakebase-auth", + "type": "module", + "version": "0.1.0", + "description": "Databricks Lakebase OAuth credential generation and token refresh for any PostgreSQL driver (pg, postgres.js, Bun.SQL)", + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "packageManager": "pnpm@10.21.0", + "repository": { + "type": "git", + "url": "git+https://github.com/databricks/appkit.git", + "directory": "packages/lakebase-auth" + }, + "keywords": [ + "databricks", + "lakebase", + "postgres", + "postgresql", + "oauth", + "auth", + "credentials", + "oltp" + ], + "license": "Apache-2.0", + "files": [ + "dist", + "README.md", + "LICENSE", + "DCO", + "sbom.cdx.json" + ], + "exports": { + ".": { + "types": "./dist/index.d.ts", + "development": "./src/index.ts", + "default": "./dist/index.js" + }, + "./package.json": "./package.json" + }, + "scripts": { + "build:package": "tsdown --config tsdown.config.ts", + "build:watch": "tsdown --config tsdown.config.ts --watch", + "clean:full": "rm -rf dist node_modules tmp", + "clean": "rm -rf dist tmp", + "dist": "tsx ../../tools/dist-lakebase-auth.ts", + "tarball": "rm -rf tmp && pnpm dist && npm pack ./tmp --pack-destination ./tmp", + "tarball:prerelease": "rm -rf tmp && SHORTSHA=$(git rev-parse --short HEAD) && pnpm dist --prerelease $SHORTSHA && npm pack ./tmp --pack-destination ./tmp", + "typecheck": "tsc --noEmit", + "release:dry": "release-it --dry-run", + "release:sbom": "pnpm exec cdxgen -t js --no-recurse --required-only -o tmp/sbom.cdx.json ." + }, + "dependencies": { + "@databricks/sdk-experimental": "0.17.0" + }, + "module": "./dist/index.js", + "publishConfig": { + "exports": { + ".": "./dist/index.js", + "./package.json": "./package.json" + } + } +} diff --git a/packages/lakebase-auth/src/__tests__/caching.test.ts b/packages/lakebase-auth/src/__tests__/caching.test.ts new file mode 100644 index 000000000..98af1f5b9 --- /dev/null +++ b/packages/lakebase-auth/src/__tests__/caching.test.ts @@ -0,0 +1,159 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { cachedWithOnDemandRefresh, cachedWithTimedRefresh } from "../caching"; +import type { Credential } from "../types"; + +const HOUR = 3_600_000; +const EARLY = 120_000; + +describe("cachedWithTimedRefresh (eager)", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + test("fetches immediately on creation and serves the cached token", async () => { + const fetch = vi.fn( + async (): Promise => ({ + token: "t1", + expiresAt: Date.now() + HOUR, + }), + ); + + const provider = cachedWithTimedRefresh(fetch, EARLY); + await vi.advanceTimersByTimeAsync(0); // flush the eager initial fetch + + expect(fetch).toHaveBeenCalledTimes(1); + await expect(provider.getToken()).resolves.toBe("t1"); + expect(fetch).toHaveBeenCalledTimes(1); + + provider.dispose(); + }); + + test("refreshes in the background before the token expires", async () => { + let n = 0; + const fetch = vi.fn(async (): Promise => { + n += 1; + return { token: `t${n}`, expiresAt: Date.now() + 200_000 }; + }); + + const provider = cachedWithTimedRefresh(fetch, EARLY); + await vi.advanceTimersByTimeAsync(0); + expect(fetch).toHaveBeenCalledTimes(1); + + // Scheduled refresh fires 200s - 120s = 80s after the fetch. + await vi.advanceTimersByTimeAsync(80_000); + expect(fetch).toHaveBeenCalledTimes(2); + await expect(provider.getToken()).resolves.toBe("t2"); + + provider.dispose(); + }); + + test("retries a failed background refresh", async () => { + const fetch = vi + .fn<() => Promise>() + .mockRejectedValueOnce(new Error("boom")) + .mockResolvedValue({ token: "t-ok", expiresAt: Date.now() + HOUR }); + const onLog = vi.fn(); + + const provider = cachedWithTimedRefresh(fetch, EARLY, onLog); + await vi.advanceTimersByTimeAsync(0); // initial attempt fails + expect(fetch).toHaveBeenCalledTimes(1); + expect(onLog).toHaveBeenCalled(); + + // BACKGROUND_RETRY_MS is 30000ms. + await vi.advanceTimersByTimeAsync(30000); + expect(fetch).toHaveBeenCalledTimes(2); + await expect(provider.getToken()).resolves.toBe("t-ok"); + + provider.dispose(); + }); + + test("dispose stops further background refreshes", async () => { + const fetch = vi.fn( + async (): Promise => ({ + token: "t", + expiresAt: Date.now() + 200_000, + }), + ); + + const provider = cachedWithTimedRefresh(fetch, EARLY); + await vi.advanceTimersByTimeAsync(0); + expect(fetch).toHaveBeenCalledTimes(1); + + provider.dispose(); + await vi.advanceTimersByTimeAsync(200_000); + // No further fetches after disposal. + expect(fetch).toHaveBeenCalledTimes(1); + }); +}); + +describe("cachedWithOnDemandRefresh (lazy)", () => { + test("does not fetch until first use, then caches", async () => { + const fetch = vi.fn( + async (): Promise => ({ + token: "t1", + expiresAt: Date.now() + HOUR, + }), + ); + + const provider = cachedWithOnDemandRefresh(fetch, EARLY); + expect(fetch).not.toHaveBeenCalled(); + + await expect(provider.getToken()).resolves.toBe("t1"); + await expect(provider.getToken()).resolves.toBe("t1"); + expect(fetch).toHaveBeenCalledTimes(1); + }); + + test("deduplicates concurrent refreshes", async () => { + let resolveFetch: (c: Credential) => void = () => {}; + const fetch = vi.fn( + (): Promise => + new Promise((resolve) => { + resolveFetch = resolve; + }), + ); + + const provider = cachedWithOnDemandRefresh(fetch, EARLY); + const p1 = provider.getToken(); + const p2 = provider.getToken(); + const p3 = provider.getToken(); + + resolveFetch({ token: "t1", expiresAt: Date.now() + HOUR }); + + await expect(Promise.all([p1, p2, p3])).resolves.toEqual([ + "t1", + "t1", + "t1", + ]); + expect(fetch).toHaveBeenCalledTimes(1); + }); + + test("refreshes on demand when the token nears expiry", async () => { + let n = 0; + const fetch = vi.fn(async (): Promise => { + n += 1; + // Expires within the early-refresh buffer, so the next call must refresh. + return { token: `t${n}`, expiresAt: Date.now() + 100 }; + }); + + const provider = cachedWithOnDemandRefresh(fetch, EARLY); + await expect(provider.getToken()).resolves.toBe("t1"); + await expect(provider.getToken()).resolves.toBe("t2"); + expect(fetch).toHaveBeenCalledTimes(2); + }); + + test("retries after a failed refresh instead of caching the rejection", async () => { + const fetch = vi + .fn<() => Promise>() + .mockRejectedValueOnce(new Error("boom")) + .mockResolvedValue({ token: "t-ok", expiresAt: Date.now() + HOUR }); + + const provider = cachedWithOnDemandRefresh(fetch, EARLY); + await expect(provider.getToken()).rejects.toThrow("boom"); + await expect(provider.getToken()).resolves.toBe("t-ok"); + expect(fetch).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/lakebase-auth/src/__tests__/config.test.ts b/packages/lakebase-auth/src/__tests__/config.test.ts new file mode 100644 index 000000000..3134cd21a --- /dev/null +++ b/packages/lakebase-auth/src/__tests__/config.test.ts @@ -0,0 +1,222 @@ +import type { WorkspaceClient } from "@databricks/sdk-experimental"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +const mockMe = vi.fn(); + +// Mock the SDK so getWorkspaceClient() can construct a client and +// getUsernameWithApiLookup() can exercise the currentUser.me() fallback. +vi.mock("@databricks/sdk-experimental", () => ({ + WorkspaceClient: vi + .fn() + .mockImplementation(() => ({ currentUser: { me: mockMe } })), +})); + +import { WorkspaceClient as MockWorkspaceClient } from "@databricks/sdk-experimental"; +import { + getUsernameSync, + getUsernameWithApiLookup, + getWorkspaceClient, + mapSslConfig, + parseConfig, +} from "../config"; + +const ENV_KEYS = [ + "PGHOST", + "PGDATABASE", + "LAKEBASE_ENDPOINT", + "PGUSER", + "PGPORT", + "PGSSLMODE", + "DATABRICKS_CLIENT_ID", +] as const; + +const original: Record = {}; + +beforeEach(() => { + vi.clearAllMocks(); + for (const key of ENV_KEYS) { + original[key] = process.env[key]; + delete process.env[key]; + } +}); + +afterEach(() => { + for (const key of ENV_KEYS) { + if (original[key] === undefined) delete process.env[key]; + else process.env[key] = original[key]; + } +}); + +describe("mapSslConfig", () => { + test.each(["verify-full", "verify-ca", "require", "prefer"] as const)( + "maps %s to certificate verification", + (mode) => { + expect(mapSslConfig(mode)).toEqual({ rejectUnauthorized: true }); + }, + ); + + test("maps disable to false", () => { + expect(mapSslConfig("disable")).toBe(false); + }); + + test("populates the Bun SNI server name from a DNS host", () => { + expect(mapSslConfig("verify-full", "ep-test.databricks.com")).toEqual({ + rejectUnauthorized: true, + serverName: "ep-test.databricks.com", + }); + }); + + test("omits SNI server name for IP-literal hosts", () => { + expect(mapSslConfig("require", "10.0.0.5")).toEqual({ + rejectUnauthorized: true, + }); + expect(mapSslConfig("require", "::1")).toEqual({ + rejectUnauthorized: true, + }); + }); + + test("never sets SNI when SSL is disabled", () => { + expect(mapSslConfig("disable", "ep-test.databricks.com")).toBe(false); + }); +}); + +describe("parseConfig", () => { + const base = { + host: "ep.databricks.com", + database: "databricks_postgres", + endpoint: "projects/p/branches/b/endpoints/e", + }; + + test("reads connection essentials from explicit config", () => { + const parsed = parseConfig(base); + expect(parsed).toMatchObject({ + host: "ep.databricks.com", + database: "databricks_postgres", + endpoint: "projects/p/branches/b/endpoints/e", + port: 5432, + sslMode: "verify-full", + }); + }); + + test("falls back to environment variables", () => { + process.env.PGHOST = "env-host"; + process.env.PGDATABASE = "env-db"; + process.env.LAKEBASE_ENDPOINT = "env-endpoint"; + process.env.PGPORT = "6543"; + process.env.PGSSLMODE = "require"; + + const parsed = parseConfig(); + expect(parsed.host).toBe("env-host"); + expect(parsed.database).toBe("env-db"); + expect(parsed.endpoint).toBe("env-endpoint"); + expect(parsed.port).toBe(6543); + expect(parsed.sslMode).toBe("require"); + }); + + test("throws when neither endpoint nor password is provided", () => { + expect(() => + parseConfig({ host: base.host, database: base.database }), + ).toThrow("LAKEBASE_ENDPOINT or config.endpoint"); + }); + + test("allows a missing endpoint when a native password is provided", () => { + const parsed = parseConfig({ + host: base.host, + database: base.database, + password: "secret", + }); + expect(parsed.endpoint).toBeUndefined(); + }); + + test("throws when the host is missing", () => { + expect(() => + parseConfig({ database: base.database, endpoint: base.endpoint }), + ).toThrow("PGHOST or config.host"); + }); + + test("throws when the database is missing", () => { + expect(() => + parseConfig({ host: base.host, endpoint: base.endpoint }), + ).toThrow("PGDATABASE or config.database"); + }); + + test("throws when PGPORT is not a number", () => { + process.env.PGPORT = "not-a-number"; + expect(() => parseConfig(base)).toThrow("port"); + }); + + test("throws on an invalid sslMode", () => { + expect(() => + parseConfig({ ...base, sslMode: "wide-open" as never }), + ).toThrow("one of: verify-full, verify-ca, require, prefer, disable"); + }); + + test("passes through an explicit ssl config", () => { + const ssl = { rejectUnauthorized: false }; + expect(parseConfig({ ...base, ssl }).ssl).toBe(ssl); + }); +}); + +describe("getWorkspaceClient", () => { + test("returns an explicitly provided workspace client", () => { + const explicit = { sentinel: true } as unknown as WorkspaceClient; + expect(getWorkspaceClient({ workspaceClient: explicit })).toBe(explicit); + expect(MockWorkspaceClient).not.toHaveBeenCalled(); + }); + + test("constructs a client from the SDK default auth chain otherwise", () => { + const client = getWorkspaceClient({}); + expect(MockWorkspaceClient).toHaveBeenCalledWith({}); + expect(client).toBeDefined(); + }); +}); + +describe("getUsernameSync", () => { + test("prefers config.user", () => { + process.env.PGUSER = "env-user"; + expect(getUsernameSync({ user: "config-user" })).toBe("config-user"); + }); + + test("falls back to PGUSER", () => { + process.env.PGUSER = "env-user"; + expect(getUsernameSync({})).toBe("env-user"); + }); + + test("falls back to DATABRICKS_CLIENT_ID", () => { + process.env.DATABRICKS_CLIENT_ID = "sp-123"; + expect(getUsernameSync({})).toBe("sp-123"); + }); + + test("throws when nothing resolves", () => { + expect(() => getUsernameSync({})).toThrow( + "config.user, PGUSER or DATABRICKS_CLIENT_ID", + ); + }); +}); + +describe("getUsernameWithApiLookup", () => { + test("returns the synchronously resolved username without an API call", async () => { + await expect(getUsernameWithApiLookup({ user: "sync-user" })).resolves.toBe( + "sync-user", + ); + expect(mockMe).not.toHaveBeenCalled(); + }); + + test("falls back to the workspace API when sync resolution fails", async () => { + mockMe.mockResolvedValue({ userName: "api-user@example.com" }); + await expect(getUsernameWithApiLookup({})).resolves.toBe( + "api-user@example.com", + ); + expect(mockMe).toHaveBeenCalledTimes(1); + }); + + test("returns undefined when the API lookup throws", async () => { + mockMe.mockRejectedValue(new Error("network down")); + await expect(getUsernameWithApiLookup({})).resolves.toBeUndefined(); + }); + + test("returns undefined when the API has no userName", async () => { + mockMe.mockResolvedValue({ userName: null }); + await expect(getUsernameWithApiLookup({})).resolves.toBeUndefined(); + }); +}); diff --git a/packages/lakebase/src/__tests__/credentials.test.ts b/packages/lakebase-auth/src/__tests__/credentials.test.ts similarity index 85% rename from packages/lakebase/src/__tests__/credentials.test.ts rename to packages/lakebase-auth/src/__tests__/credentials.test.ts index 837163b11..3218b5f63 100644 --- a/packages/lakebase/src/__tests__/credentials.test.ts +++ b/packages/lakebase-auth/src/__tests__/credentials.test.ts @@ -145,6 +145,34 @@ describe("Lakebase Authentication", () => { ).rejects.toThrow("API request failed"); }); + it.each([ + ["null", null], + ["a non-object", "not-an-object"], + ["an object missing token", { expire_time: "2026-02-06T18:00:00Z" }], + ["an object missing expire_time", { token: "abc" }], + ])("should reject when the response is %s", async (_label, response) => { + vi.mocked(mockApiClient.request).mockResolvedValue(response); + + await expect( + generateDatabaseCredential(mockWorkspaceClient, { + endpoint: "projects/test/branches/main/endpoints/primary", + }), + ).rejects.toThrow("Invalid value for credential response"); + }); + + it("should reject when token or expire_time are not strings", async () => { + vi.mocked(mockApiClient.request).mockResolvedValue({ + token: 123, + expire_time: 456, + }); + + await expect( + generateDatabaseCredential(mockWorkspaceClient, { + endpoint: "projects/test/branches/main/endpoints/primary", + }), + ).rejects.toThrow("Invalid value for credential response fields"); + }); + it("should use correct workspace host for API calls", async () => { const customHost = "https://custom-workspace.databricks.com"; diff --git a/packages/lakebase-auth/src/__tests__/errors.test.ts b/packages/lakebase-auth/src/__tests__/errors.test.ts new file mode 100644 index 000000000..a239825aa --- /dev/null +++ b/packages/lakebase-auth/src/__tests__/errors.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, test } from "vitest"; +import { ConfigurationError, LakebaseError, ValidationError } from "../errors"; + +describe("ConfigurationError", () => { + test("is a LakebaseError with a stable code", () => { + const err = new ConfigurationError("nope"); + expect(err).toBeInstanceOf(LakebaseError); + expect(err).toBeInstanceOf(Error); + expect(err.code).toBe("CONFIGURATION_ERROR"); + expect(err.name).toBe("ConfigurationError"); + }); + + test("missingEnvVar builds a descriptive message and context", () => { + const err = ConfigurationError.missingEnvVar("PGHOST"); + expect(err.message).toBe("PGHOST environment variable is required"); + expect(err.context).toEqual({ envVar: "PGHOST" }); + }); + + test("retains an optional cause", () => { + const cause = new Error("root"); + const err = new ConfigurationError("wrapped", { cause }); + expect(err.cause).toBe(cause); + }); +}); + +describe("ValidationError", () => { + test("has its own code and name", () => { + const err = new ValidationError("bad"); + expect(err).toBeInstanceOf(LakebaseError); + expect(err.code).toBe("VALIDATION_ERROR"); + expect(err.name).toBe("ValidationError"); + }); + + test("invalidValue includes the expected description and field context", () => { + const err = ValidationError.invalidValue("port", "abc", "a number"); + expect(err.message).toBe("Invalid value for port: expected a number"); + expect(err.context).toEqual({ + field: "port", + valueType: "string", + expected: "a number", + }); + }); + + test("invalidValue omits the expectation when not provided", () => { + const err = ValidationError.invalidValue("field", 42); + expect(err.message).toBe("Invalid value for field"); + expect(err.context).toMatchObject({ field: "field", valueType: "number" }); + }); + + test("reports null values distinctly from objects", () => { + const err = ValidationError.invalidValue("response", null, "an object"); + expect(err.context).toMatchObject({ valueType: "null" }); + }); +}); diff --git a/packages/lakebase-auth/src/__tests__/password-provider.test.ts b/packages/lakebase-auth/src/__tests__/password-provider.test.ts new file mode 100644 index 000000000..4d5c5067c --- /dev/null +++ b/packages/lakebase-auth/src/__tests__/password-provider.test.ts @@ -0,0 +1,208 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { + createPasswordProvider, + DEFAULT_EARLY_REFRESH_MS, +} from "../password-provider"; +import type { Credential } from "../types"; + +// Spy on the credential generation and workspace-client creation so the default +// fetcher can be exercised without any network or SDK auth. +vi.mock("../credentials", () => ({ + generateDatabaseCredential: vi.fn(), +})); +vi.mock("../config", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, getWorkspaceClient: vi.fn() }; +}); + +import { getWorkspaceClient } from "../config"; +import { generateDatabaseCredential } from "../credentials"; + +const HOUR = 3_600_000; + +const mockGenerate = vi.mocked(generateDatabaseCredential); +const mockGetWorkspaceClient = vi.mocked(getWorkspaceClient); + +describe("createPasswordProvider", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + test("exposes the documented default early-refresh buffer", () => { + expect(DEFAULT_EARLY_REFRESH_MS).toBe(120_000); + }); + + test("eager mode (default) fetches a token immediately", async () => { + const fetchCredential = vi.fn( + async (): Promise => ({ + token: "eager-token", + expiresAt: Date.now() + HOUR, + }), + ); + + const provider = createPasswordProvider({ fetchCredential }); + await vi.advanceTimersByTimeAsync(0); // flush the eager initial fetch + + expect(fetchCredential).toHaveBeenCalledTimes(1); + await expect(provider.password()).resolves.toBe("eager-token"); + expect(fetchCredential).toHaveBeenCalledTimes(1); // served from cache + + provider.dispose(); + }); + + test("lazy mode defers fetching until the password is requested", async () => { + const fetchCredential = vi.fn( + async (): Promise => ({ + token: "lazy-token", + expiresAt: Date.now() + HOUR, + }), + ); + + const provider = createPasswordProvider({ fetchCredential, mode: "lazy" }); + await vi.advanceTimersByTimeAsync(0); + expect(fetchCredential).not.toHaveBeenCalled(); + + await expect(provider.password()).resolves.toBe("lazy-token"); + expect(fetchCredential).toHaveBeenCalledTimes(1); + + provider.dispose(); + }); + + test("dispose stops eager background refreshes", async () => { + const fetchCredential = vi.fn( + async (): Promise => ({ + token: "t", + // Short-lived so a background refresh would be scheduled soon. + expiresAt: Date.now() + 200_000, + }), + ); + + const provider = createPasswordProvider({ + fetchCredential, + earlyRefreshMs: 120_000, + }); + await vi.advanceTimersByTimeAsync(0); + expect(fetchCredential).toHaveBeenCalledTimes(1); + + provider.dispose(); + await vi.advanceTimersByTimeAsync(200_000); + expect(fetchCredential).toHaveBeenCalledTimes(1); // no further refresh + + expect(() => provider.dispose()).not.toThrow(); // idempotent + }); + + test("applies the retry schedule and logs retries on failure", async () => { + const fetchCredential = vi + .fn<() => Promise>() + .mockRejectedValueOnce(new Error("transient")) + .mockResolvedValue({ + token: "after-retry", + expiresAt: Date.now() + HOUR, + }); + const onLog = vi.fn(); + + const provider = createPasswordProvider({ + fetchCredential, + mode: "lazy", + retry: { schedule: [10] }, + onLog, + }); + + const pending = provider.password(); + await vi.runAllTimersAsync(); + await expect(pending).resolves.toBe("after-retry"); + + expect(fetchCredential).toHaveBeenCalledTimes(2); + expect(onLog).toHaveBeenCalledWith( + "warn", + expect.stringContaining("Retrying"), + 10, + expect.stringContaining("transient"), + ); + + provider.dispose(); + }); + + test("an empty retry schedule disables retries", async () => { + const fetchCredential = vi + .fn<() => Promise>() + .mockRejectedValue(new Error("boom")); + + const provider = createPasswordProvider({ + fetchCredential, + mode: "lazy", + retry: { schedule: [] }, + }); + + await expect(provider.password()).rejects.toThrow("boom"); + expect(fetchCredential).toHaveBeenCalledTimes(1); + + provider.dispose(); + }); + + describe("default SDK-based fetcher", () => { + test("throws when no endpoint is provided", () => { + expect(() => createPasswordProvider({ mode: "lazy" })).toThrow( + "config.endpoint", + ); + }); + + test("uses a provided workspace client, forwarding claims and mapping expiry", async () => { + const expireTime = "2099-01-01T00:00:00Z"; + mockGenerate.mockResolvedValue({ + token: "sdk-token", + expire_time: expireTime, + }); + const workspaceClient = { config: { host: "test" } } as never; + const claims = [{ permission_set: undefined, resources: [] }]; + + const provider = createPasswordProvider({ + mode: "lazy", + endpoint: "projects/p/branches/b/endpoints/e", + workspaceClient, + claims, + }); + + await expect(provider.password()).resolves.toBe("sdk-token"); + expect(mockGetWorkspaceClient).not.toHaveBeenCalled(); + expect(mockGenerate).toHaveBeenCalledWith(workspaceClient, { + endpoint: "projects/p/branches/b/endpoints/e", + claims, + }); + + provider.dispose(); + }); + + test("lazily creates a workspace client when none is provided", async () => { + const createdClient = { config: { host: "auto" } }; + mockGetWorkspaceClient.mockReturnValue(createdClient as never); + mockGenerate.mockResolvedValue({ + token: "auto-token", + expire_time: "2099-01-01T00:00:00Z", + }); + + const provider = createPasswordProvider({ + mode: "lazy", + endpoint: "projects/p/branches/b/endpoints/e", + }); + + await expect(provider.password()).resolves.toBe("auto-token"); + expect(mockGetWorkspaceClient).toHaveBeenCalledTimes(1); + // No claims provided -> payload omits the claims key entirely. + expect(mockGenerate).toHaveBeenCalledWith(createdClient, { + endpoint: "projects/p/branches/b/endpoints/e", + }); + + // The client is created once and reused across fetches. + await provider.password(); + expect(mockGetWorkspaceClient).toHaveBeenCalledTimes(1); + + provider.dispose(); + }); + }); +}); diff --git a/packages/lakebase-auth/src/__tests__/pg-config.test.ts b/packages/lakebase-auth/src/__tests__/pg-config.test.ts new file mode 100644 index 000000000..22ee6b3f0 --- /dev/null +++ b/packages/lakebase-auth/src/__tests__/pg-config.test.ts @@ -0,0 +1,102 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { getPgConfig } from "../pg-config"; +import type { Credential } from "../types"; + +const ENV_KEYS = [ + "PGHOST", + "PGDATABASE", + "LAKEBASE_ENDPOINT", + "PGUSER", + "PGPORT", + "PGSSLMODE", + "DATABRICKS_CLIENT_ID", +] as const; + +const original: Record = {}; + +function fakeFetch(token = "tok"): () => Promise { + return vi.fn(async () => ({ token, expiresAt: Date.now() + 3_600_000 })); +} + +describe("getPgConfig", () => { + beforeEach(() => { + for (const key of ENV_KEYS) original[key] = process.env[key]; + process.env.PGHOST = "ep-test.databricks.com"; + process.env.PGDATABASE = "databricks_postgres"; + process.env.LAKEBASE_ENDPOINT = + "projects/test/branches/main/endpoints/primary"; + process.env.PGUSER = "user@example.com"; + }); + + afterEach(() => { + for (const key of ENV_KEYS) { + if (original[key] === undefined) delete process.env[key]; + else process.env[key] = original[key]; + } + }); + + test("builds a pg config from environment variables", () => { + const cfg = getPgConfig({ fetchCredential: fakeFetch(), refresh: "lazy" }); + + expect(cfg.host).toBe("ep-test.databricks.com"); + expect(cfg.database).toBe("databricks_postgres"); + expect(cfg.user).toBe("user@example.com"); + expect(cfg.port).toBe(5432); + expect(cfg.ssl).toEqual({ + rejectUnauthorized: true, + serverName: "ep-test.databricks.com", + }); + expect(typeof cfg.password).toBe("function"); + + cfg.dispose(); + }); + + test("password callback resolves to the fetched token", async () => { + const fetchCredential = fakeFetch("the-token"); + const cfg = getPgConfig({ fetchCredential, refresh: "lazy" }); + + const password = cfg.password as () => Promise; + await expect(password()).resolves.toBe("the-token"); + expect(fetchCredential).toHaveBeenCalledTimes(1); + + cfg.dispose(); + }); + + test("uses a native password without OAuth, with a no-op dispose", () => { + const cfg = getPgConfig({ + password: "static-password", + host: "ep-test.databricks.com", + database: "databricks_postgres", + user: "user@example.com", + }); + + expect(cfg.password).toBe("static-password"); + expect(() => cfg.dispose()).not.toThrow(); + }); + + test("maps sslMode 'disable' to false", () => { + const cfg = getPgConfig({ + password: "x", + sslMode: "disable", + }); + + expect(cfg.ssl).toBe(false); + }); + + test("throws when neither endpoint nor password is provided", () => { + delete process.env.LAKEBASE_ENDPOINT; + + expect(() => getPgConfig({ fetchCredential: fakeFetch() })).toThrow( + "LAKEBASE_ENDPOINT or config.endpoint", + ); + }); + + test("throws when the username cannot be resolved", () => { + delete process.env.PGUSER; + delete process.env.DATABRICKS_CLIENT_ID; + + expect(() => + getPgConfig({ fetchCredential: fakeFetch(), refresh: "lazy" }), + ).toThrow("PGUSER or DATABRICKS_CLIENT_ID"); + }); +}); diff --git a/packages/lakebase-auth/src/__tests__/retry.test.ts b/packages/lakebase-auth/src/__tests__/retry.test.ts new file mode 100644 index 000000000..e09b831f4 --- /dev/null +++ b/packages/lakebase-auth/src/__tests__/retry.test.ts @@ -0,0 +1,73 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { DEFAULT_RETRY_SCHEDULE, withRetries } from "../retry"; + +describe("withRetries", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + test("returns the result on first success without retrying", async () => { + const fn = vi.fn(async () => "ok"); + const wrapped = withRetries(fn, [50, 500]); + + await expect(wrapped()).resolves.toBe("ok"); + expect(fn).toHaveBeenCalledTimes(1); + }); + + test("retries after each delay then succeeds", async () => { + const fn = vi + .fn<() => Promise>() + .mockRejectedValueOnce(new Error("fail 1")) + .mockRejectedValueOnce(new Error("fail 2")) + .mockResolvedValueOnce("ok"); + const onLog = vi.fn(); + const wrapped = withRetries(fn, [50, 500], onLog); + + const promise = wrapped(); + // Run all pending timers (the inter-attempt delays) to completion. + await vi.runAllTimersAsync(); + + await expect(promise).resolves.toBe("ok"); + expect(fn).toHaveBeenCalledTimes(3); + expect(onLog).toHaveBeenCalledTimes(2); + expect(onLog).toHaveBeenCalledWith( + "warn", + expect.stringContaining("Retrying"), + 50, + expect.stringContaining("fail 1"), + ); + }); + + test("throws the final error after exhausting the schedule", async () => { + const fn = vi.fn(async () => { + throw new Error("always fails"); + }); + const wrapped = withRetries(fn, [10, 20]); + + const promise = wrapped(); + const assertion = expect(promise).rejects.toThrow("always fails"); + await vi.runAllTimersAsync(); + await assertion; + + // 2 retries (one per delay) + 1 final attempt + expect(fn).toHaveBeenCalledTimes(3); + }); + + test("makes a single attempt when schedule is empty", async () => { + const fn = vi.fn(async () => { + throw new Error("nope"); + }); + const wrapped = withRetries(fn, []); + + await expect(wrapped()).rejects.toThrow("nope"); + expect(fn).toHaveBeenCalledTimes(1); + }); + + test("exposes a sensible default schedule", () => { + expect(DEFAULT_RETRY_SCHEDULE).toEqual([50, 500, 5000]); + }); +}); diff --git a/packages/lakebase-auth/src/caching.ts b/packages/lakebase-auth/src/caching.ts new file mode 100644 index 000000000..c90f6f5b4 --- /dev/null +++ b/packages/lakebase-auth/src/caching.ts @@ -0,0 +1,142 @@ +import type { FetchCredential, LogFn } from "./types"; + +/** + * Delay before retrying a background (eager) refresh that failed even after + * the inner {@link withRetries} schedule was exhausted. + */ +const BACKGROUND_RETRY_MS = 30000; + +/** A cached token source with a disposer to release any background timers. */ +export interface CachedTokenProvider { + /** Resolve to a valid token, refreshing if necessary. */ + getToken: () => Promise; + /** Stop any scheduled background refresh (idempotent). */ + dispose: () => void; +} + +/** + * Eagerly cache a token, refreshing it in the background before it expires. + * + * A first fetch is kicked off immediately (not awaited) and subsequent + * refreshes are scheduled via an `unref`'d timer, so the token is kept warm + * without holding the event loop open. If a background refresh fails it is + * retried after a short delay. {@link CachedTokenProvider.getToken} returns the + * cached token while valid and otherwise falls back to an on-demand refresh. + * + * @param fetchCredential Fetches a fresh credential (typically retry-wrapped) + * @param earlyRefreshMs How long before expiry to refresh + * @param onLog Optional structured logging callback + */ +export function cachedWithTimedRefresh( + fetchCredential: FetchCredential, + earlyRefreshMs: number, + onLog?: LogFn, +): CachedTokenProvider { + let cachedToken = ""; + let expiresAt = 0; + let inFlight: Promise | null = null; + let timer: ReturnType | null = null; + let disposed = false; + + function scheduleNext(delayMs: number): void { + if (disposed) return; + if (timer) clearTimeout(timer); + timer = setTimeout( + () => { + // Background refresh: swallow rejection here (getToken surfaces errors + // to callers); failures reschedule themselves below. + refresh().catch(() => {}); + }, + Math.max(0, delayMs), + ); + timer.unref?.(); + } + + function refresh(): Promise { + if (inFlight) return inFlight; + inFlight = (async () => { + try { + const { token, expiresAt: exp } = await fetchCredential(); + cachedToken = token; + expiresAt = exp; + scheduleNext(exp - Date.now() - earlyRefreshMs); + return token; + } catch (err) { + onLog?.( + "warn", + "Background token refresh failed, retrying in %dms: %s", + BACKGROUND_RETRY_MS, + err instanceof Error ? err.message : String(err), + ); + scheduleNext(BACKGROUND_RETRY_MS); + throw err; + } finally { + inFlight = null; + } + })(); + return inFlight; + } + + // Kick off the eager fetch immediately. Swallow the rejection so an initial + // failure doesn't surface as an unhandled rejection; getToken re-triggers. + refresh().catch(() => {}); + + return { + getToken() { + if (cachedToken && Date.now() < expiresAt - earlyRefreshMs) { + return Promise.resolve(cachedToken); + } + return refresh(); + }, + dispose() { + disposed = true; + if (timer) { + clearTimeout(timer); + timer = null; + } + }, + }; +} + +const PAST_EPOCH = Number.NEGATIVE_INFINITY; +const FAR_FUTURE_EPOCH = Number.POSITIVE_INFINITY; + +/** + * Lazily cache a token, refreshing it on demand when it nears expiry. + * + * The first call triggers a fetch; concurrent callers share the same in-flight + * promise (deduplication). A failed refresh resets the cache so the next call + * retries rather than serving a rejected promise indefinitely. + * + * @param fetchCredential Fetches a fresh credential (typically retry-wrapped) + * @param earlyRefreshMs How long before expiry to refresh + */ +export function cachedWithOnDemandRefresh( + fetchCredential: FetchCredential, + earlyRefreshMs: number, +): CachedTokenProvider { + let refreshAfter = PAST_EPOCH; // first call must refresh + let cachedToken: Promise = Promise.resolve(""); + + async function refresh(): Promise { + refreshAfter = FAR_FUTURE_EPOCH; // dedupe: no more refreshes until done + try { + const { token, expiresAt } = await fetchCredential(); + refreshAfter = expiresAt - earlyRefreshMs; + return token; + } catch (err) { + refreshAfter = PAST_EPOCH; // allow the next call to retry + throw err; + } + } + + return { + getToken() { + if (Date.now() > refreshAfter) cachedToken = refresh(); + return cachedToken; + }, + dispose() { + // No background timers to clean up in lazy mode. + }, + }; +} diff --git a/packages/lakebase/src/config.ts b/packages/lakebase-auth/src/config.ts similarity index 63% rename from packages/lakebase/src/config.ts rename to packages/lakebase-auth/src/config.ts index ef055edc1..038d23517 100644 --- a/packages/lakebase/src/config.ts +++ b/packages/lakebase-auth/src/config.ts @@ -1,36 +1,72 @@ import { WorkspaceClient } from "@databricks/sdk-experimental"; -import type pg from "pg"; import { ConfigurationError, ValidationError } from "./errors"; -import type { LakebasePoolConfig } from "./types"; +import type { DriverSslConfig, LakebaseAuthConfig, SslConfig } from "./types"; -/** Default configuration values for the Lakebase connector */ +/** Default connection values for Lakebase auth */ const defaults = { port: 5432, - sslMode: "require" as const, - max: 10, - idleTimeoutMillis: 30_000, - connectionTimeoutMillis: 10_000, + sslMode: "verify-full" as const, }; -const VALID_SSL_MODES = ["require", "disable", "prefer"] as const; +const VALID_SSL_MODES = [ + "verify-full", + "verify-ca", + "require", + "prefer", + "disable", +] as const; type SslMode = (typeof VALID_SSL_MODES)[number]; -export interface ParsedPoolConfig { +/** Connection essentials parsed from config and environment variables. */ +export interface ParsedAuthConfig { endpoint?: string; host: string; database: string; port: number; sslMode: SslMode; - ssl?: pg.PoolConfig["ssl"]; - max: number; - idleTimeoutMillis: number; - connectionTimeoutMillis: number; + ssl?: SslConfig; } -/** Parse pool configuration from provided config and environment variables */ -export function parsePoolConfig( - userConfig?: Partial, -): ParsedPoolConfig { +/** + * Map an SSL mode string to the corresponding SSL/TLS configuration. + * - `"verify-full"` -- SSL enabled with certificate verification against system certs + * - `"verify-ca"` -- ditto (upgraded) + * - `"require"` -- ditto (upgraded) + * - `"prefer"` -- ditto (upgraded) + * - `"disable"` -- SSL disabled (not compatible with Lakebase) + * + * When `host` is a DNS name (not an IP literal), the `Bun.SQL` SNI server name + * is set on the returned object — see {@link DriverSslConfig.serverName}. + */ +export function mapSslConfig(sslMode: SslMode, host?: string): DriverSslConfig { + switch (sslMode) { + case "verify-full": // since JS drivers check root certs, there's an implied sslrootcert=system + case "verify-ca": // upgraded to equivalent of verify-full, sslrootcert=system + case "require": // upgraded to equivalent of verify-full, sslrootcert=system + case "prefer": { + // upgraded to equivalent of verify-full, sslrootcert=system + const ssl: DriverSslConfig = { rejectUnauthorized: true }; + // `Bun.SQL` won't derive the SNI server name from the host when TLS is an + // object (https://github.com/oven-sh/bun/issues/26369), so set it. `pg` + // and `postgres.js` already do this themselves. SNI must not be an IP. + if (host && !isIpAddress(host)) ssl.serverName = host; + return ssl; + } + case "disable": + return false; + } +} + +/** Whether a host is an IP literal, which must not be sent as a TLS SNI name. */ +function isIpAddress(host: string): boolean { + // IPv4 dotted-quad, or anything containing a colon (IPv6). + return /^\d{1,3}(?:\.\d{1,3}){3}$/.test(host) || host.includes(":"); +} + +/** Parse connection configuration from provided config and environment variables */ +export function parseConfig( + userConfig?: Partial, +): ParsedAuthConfig { // Get endpoint (required only for OAuth auth) const endpoint = userConfig?.endpoint ?? process.env.LAKEBASE_ENDPOINT; @@ -65,16 +101,8 @@ export function parsePoolConfig( // Get SSL mode (optional, default from defaults) const rawSslMode = userConfig?.sslMode ?? process.env.PGSSLMODE ?? undefined; - const sslMode = validateSslMode(rawSslMode) ?? defaults.sslMode; - // Pool options (with defaults) - const max = userConfig?.max ?? defaults.max; - const idleTimeoutMillis = - userConfig?.idleTimeoutMillis ?? defaults.idleTimeoutMillis; - const connectionTimeoutMillis = - userConfig?.connectionTimeoutMillis ?? defaults.connectionTimeoutMillis; - return { endpoint, host, @@ -82,9 +110,6 @@ export function parsePoolConfig( port, sslMode, ssl: userConfig?.ssl, - max, - idleTimeoutMillis, - connectionTimeoutMillis, }; } @@ -98,7 +123,7 @@ function validateSslMode(value: string | undefined): SslMode | undefined { throw ValidationError.invalidValue( "sslMode (PGSSLMODE)", value, - `one of: ${VALID_SSL_MODES.join(", ")}`, + `one of: ${VALID_SSL_MODES.join(", ")} (all except "disable" are treated as "verify-full")`, ); } @@ -107,7 +132,7 @@ function validateSslMode(value: string | undefined): SslMode | undefined { /** Get workspace client from config or SDK default auth chain */ export function getWorkspaceClient( - config: Partial, + config: Partial, ): WorkspaceClient { // Priority 1: Explicit workspaceClient in config if (config.workspaceClient) { @@ -121,7 +146,7 @@ export function getWorkspaceClient( } /** Get username synchronously from config or environment */ -export function getUsernameSync(config: Partial): string { +export function getUsernameSync(config: Partial): string { // Priority 1: Explicit user in config if (config.user) { return config.user; @@ -162,7 +187,7 @@ export function getUsernameSync(config: Partial): string { * caller can decide whether to proceed or surface an error. */ export async function getUsernameWithApiLookup( - config?: Partial, + config?: Partial, ): Promise { try { return getUsernameSync(config ?? {}); diff --git a/packages/lakebase/src/credentials.ts b/packages/lakebase-auth/src/credentials.ts similarity index 100% rename from packages/lakebase/src/credentials.ts rename to packages/lakebase-auth/src/credentials.ts diff --git a/packages/lakebase/src/errors.ts b/packages/lakebase-auth/src/errors.ts similarity index 97% rename from packages/lakebase/src/errors.ts rename to packages/lakebase-auth/src/errors.ts index 86c3ecc1f..1edf3c711 100644 --- a/packages/lakebase/src/errors.ts +++ b/packages/lakebase-auth/src/errors.ts @@ -1,5 +1,5 @@ /** - * Base error class for Lakebase driver errors. + * Base error class for Lakebase auth errors. */ export abstract class LakebaseError extends Error { abstract readonly code: string; diff --git a/packages/lakebase-auth/src/index.ts b/packages/lakebase-auth/src/index.ts new file mode 100644 index 000000000..355a43391 --- /dev/null +++ b/packages/lakebase-auth/src/index.ts @@ -0,0 +1,42 @@ +export { + getUsernameSync, + getUsernameWithApiLookup, + getWorkspaceClient, + mapSslConfig, + type ParsedAuthConfig, + parseConfig, +} from "./config"; +export { generateDatabaseCredential } from "./credentials"; +export { + ConfigurationError, + LakebaseError, + ValidationError, +} from "./errors"; +export { + type CreatePasswordProviderOptions, + createPasswordProvider, + DEFAULT_EARLY_REFRESH_MS, + type PasswordProvider, +} from "./password-provider"; +export { + type GetPgConfigOptions, + getPgConfig, + type PgConfig, +} from "./pg-config"; +export { DEFAULT_RETRY_SCHEDULE, withRetries } from "./retry"; +export type { + Credential, + DatabaseCredential, + DriverSslConfig, + FetchCredential, + GenerateDatabaseCredentialRequest, + LakebaseAuthConfig, + LogFn, + LogLevel, + RefreshMode, + RequestedClaims, + RequestedResource, + RetryOptions, + SslConfig, +} from "./types"; +export { RequestedClaimsPermissionSet } from "./types"; diff --git a/packages/lakebase-auth/src/password-provider.ts b/packages/lakebase-auth/src/password-provider.ts new file mode 100644 index 000000000..ec44ba3b9 --- /dev/null +++ b/packages/lakebase-auth/src/password-provider.ts @@ -0,0 +1,152 @@ +import type { WorkspaceClient } from "@databricks/sdk-experimental"; +import { + type CachedTokenProvider, + cachedWithOnDemandRefresh, + cachedWithTimedRefresh, +} from "./caching"; +import { getWorkspaceClient } from "./config"; +import { generateDatabaseCredential } from "./credentials"; +import { ConfigurationError } from "./errors"; +import { withRetries } from "./retry"; +import type { + Credential, + FetchCredential, + LakebaseAuthConfig, + LogFn, + RefreshMode, + RequestedClaims, + RetryOptions, +} from "./types"; + +/** Default early-refresh buffer: refresh ~2 minutes before the 1-hour expiry. */ +export const DEFAULT_EARLY_REFRESH_MS = 2 * 60 * 1000; + +/** + * An async password callback paired with a disposer. + * + * `password` is the function to hand to a Postgres driver (`pg`, `postgres.js`, + * `Bun.SQL`, ...) as the connection password. `dispose` stops any background + * refresh timer and should be called when the connection/pool is closed. + */ +export interface PasswordProvider { + /** Async callback returning a valid OAuth token to use as the password. */ + password: () => Promise; + /** Release background refresh timers (idempotent). Call on shutdown. */ + dispose: () => void; +} + +/** + * Options for {@link createPasswordProvider}. + */ +export interface CreatePasswordProviderOptions { + /** + * Custom credential fetcher. When omitted, a default fetcher is built from + * `endpoint` + `workspaceClient`/`userConfig` using the Databricks SDK. This + * is the seam consumers use to wrap credential generation with telemetry. + */ + fetchCredential?: FetchCredential; + + /** Endpoint resource path (required when `fetchCredential` is not provided). */ + endpoint?: string; + + /** Workspace client used by the default fetcher (created lazily if omitted). */ + workspaceClient?: WorkspaceClient; + + /** Config passed to {@link getWorkspaceClient} by the default fetcher. */ + userConfig?: Partial; + + /** Optional UC claims for the default fetcher. */ + claims?: RequestedClaims[]; + + /** Refresh strategy. @default "eager" */ + mode?: RefreshMode; + + /** How long before expiry to refresh (ms). @default 120000 */ + earlyRefreshMs?: number; + + /** Retry options for credential fetches. */ + retry?: RetryOptions; + + /** Optional structured logging callback. */ + onLog?: LogFn; +} + +/** Build the default SDK-based credential fetcher. */ +function defaultFetchCredential( + opts: CreatePasswordProviderOptions, +): FetchCredential { + const { endpoint } = opts; + if (!endpoint) { + throw ConfigurationError.missingEnvVar( + "config.endpoint (required to generate OAuth credentials)", + ); + } + + // Lazily initialize the workspace client on first fetch. + let workspaceClient: WorkspaceClient | null = opts.workspaceClient ?? null; + + return async (): Promise => { + if (!workspaceClient) { + workspaceClient = getWorkspaceClient(opts.userConfig ?? {}); + } + const credential = await generateDatabaseCredential(workspaceClient, { + endpoint, + ...(opts.claims ? { claims: opts.claims } : {}), + }); + return { + token: credential.token, + expiresAt: new Date(credential.expire_time).getTime(), + }; + }; +} + +/** + * Create an OAuth password provider for Lakebase Postgres connections. + * + * The returned `password` callback fetches, caches, and refreshes a Lakebase + * OAuth token, ready to use as the password for any Postgres driver. Refresh + * is eager by default (token kept warm in the background) and retries + * transient failures. + * + * @example pg + * ```typescript + * const { password, dispose } = createPasswordProvider({ endpoint }); + * const pool = new pg.Pool({ host, user, database, password }); + * // on shutdown: await pool.end(); dispose(); + * ``` + * + * @example postgres.js + * ```typescript + * const { password } = createPasswordProvider({ endpoint }); + * const sql = postgres({ host, username, database, password }); + * ``` + * + * @example Bun.SQL + * ```typescript + * const { password } = createPasswordProvider({ endpoint }); + * const sql = new Bun.SQL({ hostname: host, username, database, password }); + * ``` + */ +export function createPasswordProvider( + opts: CreatePasswordProviderOptions, +): PasswordProvider { + const earlyRefreshMs = opts.earlyRefreshMs ?? DEFAULT_EARLY_REFRESH_MS; + const mode: RefreshMode = opts.mode ?? "eager"; + + const baseFetch = opts.fetchCredential ?? defaultFetchCredential(opts); + const fetchWithRetries = withRetries( + baseFetch, + opts.retry?.schedule, + opts.onLog, + ); + + const cache: CachedTokenProvider = + mode === "eager" + ? cachedWithTimedRefresh(fetchWithRetries, earlyRefreshMs, opts.onLog) + : cachedWithOnDemandRefresh(fetchWithRetries, earlyRefreshMs); + + return { + password: () => cache.getToken(), + dispose: () => cache.dispose(), + }; +} diff --git a/packages/lakebase-auth/src/pg-config.ts b/packages/lakebase-auth/src/pg-config.ts new file mode 100644 index 000000000..2bfb90784 --- /dev/null +++ b/packages/lakebase-auth/src/pg-config.ts @@ -0,0 +1,120 @@ +import { getUsernameSync, mapSslConfig, parseConfig } from "./config"; +import { createPasswordProvider } from "./password-provider"; +import type { + DriverSslConfig, + FetchCredential, + LakebaseAuthConfig, +} from "./types"; + +/** + * A driver-agnostic Postgres connection config with an OAuth password callback. + * + * The connection fields are compatible with `pg` (and, by spreading, with + * `postgres.js` and `Bun.SQL` — note those use `username`/`hostname`/`tls`, + * so map the fields accordingly). `dispose` stops any background token refresh + * and should be called when the connection/pool is closed; it is a no-op for + * native-password and lazy-refresh configs. + */ +export interface PgConfig { + host: string; + port: number; + user: string; + database: string; + password: string | (() => string | Promise); + /** + * SSL config, typed narrowly ({@link DriverSslConfig}) so it can be passed + * directly to `pg`, `postgres.js`, and `Bun.SQL` without a cast. A wider + * `tls.ConnectionOptions` supplied via `config.ssl` is still emitted at + * runtime — only the static type is narrowed. + */ + ssl: DriverSslConfig; + /** Stop background token refresh (idempotent). Call on shutdown. */ + dispose: () => void; +} + +/** Options accepted by {@link getPgConfig}. */ +export interface GetPgConfigOptions extends Partial { + /** + * Custom credential fetcher (e.g. telemetry-wrapped). When omitted, the + * default SDK-based fetcher is used with `endpoint` + `workspaceClient`. + */ + fetchCredential?: FetchCredential; +} + +/** + * Build a Postgres connection config for Lakebase with OAuth token refresh. + * + * Reads from the provided config and falls back to environment variables + * (`PGHOST`, `PGDATABASE`, `LAKEBASE_ENDPOINT`, `PGUSER`/`DATABRICKS_CLIENT_ID`, + * `PGPORT`, `PGSSLMODE`). When a native `password` is supplied, OAuth is skipped. + * + * @example pg + * ```typescript + * import pg from "pg"; + * import { getPgConfig } from "@databricks/lakebase-auth"; + * + * const { dispose, ...config } = getPgConfig(); + * const pool = new pg.Pool(config); + * // on shutdown: await pool.end(); dispose(); + * ``` + * + * @example postgres.js + * ```typescript + * import postgres from "postgres"; + * import { getPgConfig } from "@databricks/lakebase-auth"; + * + * const { host, port, user, database, password, ssl } = getPgConfig(); + * const sql = postgres({ host, port, username: user, database, password, ssl }); + * ``` + * + * @example Bun.SQL + * ```typescript + * import { getPgConfig } from "@databricks/lakebase-auth"; + * + * const { host, port, user, database, password, ssl } = getPgConfig(); + * const sql = new Bun.SQL({ hostname: host, port, username: user, database, password, tls: ssl }); + * ``` + */ +export function getPgConfig(config?: GetPgConfigOptions): PgConfig { + const userConfig = config ?? {}; + const parsed = parseConfig(userConfig); + const user = getUsernameSync(userConfig); + const ssl = parsed.ssl ?? mapSslConfig(parsed.sslMode, parsed.host); + + // Native password authentication: no OAuth provider needed. + if (userConfig.password !== undefined) { + return { + host: parsed.host, + port: parsed.port, + user, + database: parsed.database, + password: userConfig.password, + ssl, + dispose: () => {}, + }; + } + + const provider = createPasswordProvider({ + fetchCredential: userConfig.fetchCredential, + // endpoint is guaranteed here -- parseConfig() throws if neither + // endpoint nor password is provided. + endpoint: parsed.endpoint, + workspaceClient: userConfig.workspaceClient, + userConfig, + claims: userConfig.claims, + mode: userConfig.refresh, + earlyRefreshMs: userConfig.earlyRefreshMs, + retry: userConfig.retry, + onLog: userConfig.onLog, + }); + + return { + host: parsed.host, + port: parsed.port, + user, + database: parsed.database, + password: provider.password, + ssl, + dispose: provider.dispose, + }; +} diff --git a/packages/lakebase-auth/src/retry.ts b/packages/lakebase-auth/src/retry.ts new file mode 100644 index 000000000..14d8ce234 --- /dev/null +++ b/packages/lakebase-auth/src/retry.ts @@ -0,0 +1,42 @@ +import type { LogFn } from "./types"; + +/** Default retry schedule: one retry after each delay, then a final attempt. */ +export const DEFAULT_RETRY_SCHEDULE = [50, 500, 5000]; + +/** + * Wrap an async function with basic retry logic. + * + * On each failure the function is retried after the next delay in `schedule`. + * Once the schedule is exhausted, one final attempt is made and its result + * (or error) is returned — rethrowing the original rather than a wrapped error + * preserves the stack trace. An empty `schedule` disables retries. + * + * @param fn The async function to execute + * @param schedule Delays (ms) between retry attempts + * @param onLog Optional structured logging callback for retry warnings + * @returns A function with the same signature that applies the retry policy + */ +export function withRetries( + fn: () => Promise, + schedule: number[] = DEFAULT_RETRY_SCHEDULE, + onLog?: LogFn, +): () => Promise { + return async function retrying(): Promise { + for (const retryDelay of schedule) { + try { + return await fn(); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + onLog?.( + "warn", + "Retrying credential fetch in %dms after error: %s", + retryDelay, + message, + ); + await new Promise((resolve) => setTimeout(resolve, retryDelay)); + } + } + // Final attempt: call again rather than rethrowing to preserve stack trace. + return await fn(); + }; +} diff --git a/packages/lakebase-auth/src/types.ts b/packages/lakebase-auth/src/types.ts new file mode 100644 index 000000000..b05d61f1f --- /dev/null +++ b/packages/lakebase-auth/src/types.ts @@ -0,0 +1,279 @@ +import type { ConnectionOptions } from "node:tls"; +import type { WorkspaceClient } from "@databricks/sdk-experimental"; + +/** + * SSL configuration *accepted* as input (e.g. `config.ssl`). + * + * Intentionally wide — structurally identical to `pg`'s `ssl` option + * (`boolean | tls.ConnectionOptions`) — so callers can supply a full TLS + * config (custom CA, client certs, etc.) without depending on `pg` here. + * + * Note: what {@link getPgConfig} *emits* is typed as the narrower + * {@link DriverSslConfig} (see below) so the result is assignable to every + * driver. Any wider value you pass in is still emitted unchanged at runtime; + * only the static type is narrowed. + */ +export type SslConfig = boolean | ConnectionOptions; + +/** + * SSL configuration *emitted* by {@link getPgConfig} (the `ssl` field). + * + * Deliberately narrower than {@link SslConfig}: this structural type is + * assignable to the `ssl`/`tls` option of `pg`, `postgres.js`, and `Bun.SQL` + * alike (Bun's `TLSOptions` rejects `node:tls.ConnectionOptions`'s `key`/`cert` + * shapes, so the wide type can't be used directly). The default `sslMode` + * mapping only ever produces `false` or `{ rejectUnauthorized: true }`. + * + * This is a type-level narrowing only: if you pass a wider + * `tls.ConnectionOptions` via `config.ssl`, that exact object is still emitted + * at runtime. + */ +export type DriverSslConfig = + | boolean + | { + rejectUnauthorized?: boolean; + /** + * SNI server name, in the `Bun.SQL` spelling. Only emitted for `Bun.SQL`'s + * benefit: it works around + * {@link https://github.com/oven-sh/bun/issues/26369}, where `Bun.SQL` + * fails to derive SNI from the host when TLS is given as an object. + * + * `pg` and `postgres.js` set `servername` from the host themselves (and + * harmlessly ignore this camelCase key), so it isn't emitted for them. + */ + serverName?: string; + }; + +/** Log severity levels emitted by the auth package. */ +export type LogLevel = "debug" | "info" | "warn" | "error"; + +/** + * Optional structured logging callback. + * + * The auth package has no logging dependency; instead it emits log events + * through this callback when provided. Arguments follow `util.format`/`console` + * semantics (a format string plus interpolation args). + */ +export type LogFn = ( + level: LogLevel, + message: string, + ...args: unknown[] +) => void; + +/** + * Token refresh strategy. + * + * - `"eager"` (default): fetch a token immediately and refresh it in the + * background before it expires, regardless of whether it is used. Best for + * time-sensitive, user-facing apps where the first query must be warm. + * - `"lazy"`: fetch a token on first use and refresh it on demand when it + * nears expiry. Best for background jobs and infrequently-used connections. + */ +export type RefreshMode = "eager" | "lazy"; + +/** + * Retry configuration for credential fetches. + */ +export interface RetryOptions { + /** + * Delays (in ms) between retry attempts; one retry per array entry. An empty + * array disables retries. + * + * @default [50, 500, 5000] + */ + schedule?: number[]; +} + +/** + * Database credentials with OAuth token for Postgres connection. + */ +export interface DatabaseCredential { + /** OAuth token to use as the password when connecting to Postgres */ + token: string; + + /** + * Token expiration time in UTC (ISO 8601 format) + * Tokens expire after 1 hour from generation + * @example "2026-02-06T17:07:00Z" + */ + expire_time: string; +} + +/** + * Normalized credential with the expiry parsed to epoch milliseconds. + * This is the shape returned by a {@link FetchCredential} function. + */ +export interface Credential { + /** OAuth token to use as the Postgres password */ + token: string; + /** Token expiration time as epoch milliseconds */ + expiresAt: number; +} + +/** + * A function that fetches a fresh credential. Used as the injectable seam that + * lets consumers (e.g. `@databricks/lakebase`) wrap credential generation with + * telemetry while keeping this package dependency-light. + */ +export type FetchCredential = () => Promise; + +/** + * Permission set for Unity Catalog table access + */ +export enum RequestedClaimsPermissionSet { + /** + * Read-only access to specified UC tables + */ + READ_ONLY = "READ_ONLY", +} + +/** + * Resource to request permissions for in Unity Catalog + */ +export interface RequestedResource { + /** + * Unity Catalog table name to request access to + * @example "catalog.schema.table" + */ + table_name?: string; + + /** + * Generic resource name for non-table resources + */ + unspecified_resource_name?: string; +} + +/** + * Optional claims for fine-grained Unity Catalog table permissions + * When specified, the returned token will be scoped to only the requested tables + */ +export interface RequestedClaims { + /** + * Permission level to request + */ + permission_set?: RequestedClaimsPermissionSet; + + /** + * List of UC resources to request access to + */ + resources?: RequestedResource[]; +} + +/** + * Request parameters for generating database OAuth credentials + */ +export interface GenerateDatabaseCredentialRequest { + /** + * Endpoint resource path. Retrieve using the Databricks CLI: + * ``` + * databricks postgres list-endpoints projects/{project-id}/branches/{branch-id} + * ``` + * Use the `name` field from the output. + * + * @example "projects/{project-id}/branches/{branch-id}/endpoints/{endpoint-identifier}" + */ + endpoint: string; + + /** + * Optional claims for fine-grained UC table permissions. + * When specified, the token will only grant access to the specified tables. + * + * @example + * ```typescript + * { + * claims: [{ + * permission_set: RequestedClaimsPermissionSet.READ_ONLY, + * resources: [{ table_name: "catalog.schema.users" }] + * }] + * } + * ``` + */ + claims?: RequestedClaims[]; +} + +/** + * Driver-agnostic configuration for Lakebase OAuth authentication. + * + * Supports two authentication methods: + * 1. OAuth token authentication - Provide workspaceClient + endpoint (automatic token rotation) + * 2. Native Postgres password authentication - Provide password string or function + * + * @see https://docs.databricks.com/aws/en/oltp/projects/authentication + */ +export interface LakebaseAuthConfig { + /** + * Databricks workspace client for OAuth authentication. + * If not provided along with endpoint, the SDK default auth chain is used. + * + * Note: If password is provided, OAuth auth is not used. + */ + workspaceClient?: WorkspaceClient; + + /** + * Endpoint resource path for OAuth token generation. + * + * Retrieve the value using the Databricks CLI: + * ``` + * databricks postgres list-endpoints projects/{project-id}/branches/{branch-id} + * ``` + * Use the `name` field from the output. + * + * Required for OAuth authentication (unless password is provided). + * Can also be set via the LAKEBASE_ENDPOINT environment variable. + * + * @example "projects/{project-id}/branches/{branch-id}/endpoints/{endpoint-identifier}" + */ + endpoint?: string; + + /** PostgreSQL host. Can also be set via PGHOST. */ + host?: string; + + /** Database name. Can also be set via PGDATABASE. */ + database?: string; + + /** PostgreSQL username. Can also be set via PGUSER or DATABRICKS_CLIENT_ID. */ + user?: string; + + /** Port number. Can also be set via PGPORT. @default 5432 */ + port?: number; + + /** + * SSL mode (convenience helper). Can also be set via PGSSLMODE. All values + * other than "disable" are treated as "verify-full" with system root certs. + * @default "verify-full" + */ + sslMode?: "verify-full" | "verify-ca" | "require" | "prefer" | "disable"; + + /** Explicit SSL configuration; overrides {@link sslMode} when provided. */ + ssl?: SslConfig; + + /** + * Native Postgres password (skips OAuth). A string or an (optionally async) + * callback returning a password. + */ + password?: string | (() => string | Promise); + + /** + * Optional UC claims for fine-grained table permissions on the generated + * Postgres token. + */ + claims?: RequestedClaims[]; + + /** + * Token refresh strategy. + * @default "eager" + */ + refresh?: RefreshMode; + + /** + * How long before token expiry to refresh, in milliseconds. + * @default 120000 (2 minutes) + */ + earlyRefreshMs?: number; + + /** Retry options for transient credential-fetch failures. */ + retry?: RetryOptions; + + /** Optional structured logging callback. */ + onLog?: LogFn; +} diff --git a/packages/lakebase-auth/tsconfig.json b/packages/lakebase-auth/tsconfig.json new file mode 100644 index 000000000..4a6e68b3e --- /dev/null +++ b/packages/lakebase-auth/tsconfig.json @@ -0,0 +1,12 @@ +{ + "extends": "../../tsconfig.json", + "compilerOptions": { + "outDir": "dist", + "baseUrl": ".", + "paths": { + "@/*": ["src/*"] + } + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/lakebase-auth/tsdown.config.ts b/packages/lakebase-auth/tsdown.config.ts new file mode 100644 index 000000000..37cf9b9fa --- /dev/null +++ b/packages/lakebase-auth/tsdown.config.ts @@ -0,0 +1,31 @@ +import { defineConfig } from "tsdown"; + +export default defineConfig([ + { + publint: true, + name: "@databricks/lakebase-auth", + entry: "src/index.ts", + outDir: "dist", + hash: false, + format: "esm", + platform: "node", + minify: false, + dts: { + resolver: "oxc", + }, + sourcemap: false, + clean: false, + unbundle: true, + noExternal: [], + outExtensions: () => ({ + js: ".js", + }), + external: (id) => { + // Bundle all internal modules + if (id.startsWith("@/")) return false; + // Externalize all npm packages + return /^[^./]/.test(id) || id.includes("/node_modules/"); + }, + tsconfig: "./tsconfig.json", + }, +]); diff --git a/packages/lakebase/README.md b/packages/lakebase/README.md index 99f3beb60..7816a072d 100644 --- a/packages/lakebase/README.md +++ b/packages/lakebase/README.md @@ -8,13 +8,16 @@ PostgreSQL driver for Databricks Lakebase Autoscaling with automatic OAuth token It: -- Returns a standard `pg.Pool` - works with any PostgreSQL library or ORM -- Automatically refreshes OAuth tokens (1-hour lifetime, with 2-minute buffer) +- Returns a standard `pg.Pool` - works with many PostgreSQL libraries and ORMs +- Automatically refreshes OAuth tokens (1-hour lifetime, with 2-minute buffer) — **eagerly**, in the background (by default), or **lazily**, on demand +- Retries transient credential-fetch failures (e.g. a briefly unreachable OAuth server) - Caches tokens to minimize API calls - Zero configuration with environment variables - Optional OpenTelemetry instrumentation -**NOTE:** This package is NOT compatible with the Databricks Lakebase Provisioned. +OAuth credential generation and token-refresh logic lives in the smaller, driver-agnostic [`@databricks/lakebase-auth`](https://www.npmjs.com/package/@databricks/lakebase-auth) package. Use that package directly if you want a Postgres config for `postgres.js`, `Bun.SQL`, or `pg` without the bundled `pg.Pool` and OpenTelemetry instrumentation. + +**NOTE:** This package is NOT compatible with Databricks Lakebase Provisioned. ## Installation @@ -26,23 +29,24 @@ npm install @databricks/lakebase ### Using Environment Variables +Ensure Databricks credentials are available, for example in `.databrickscfg` or by setting `DATABRICKS_HOST`, `DATABRICKS_CLIENT_ID`, and `DATABRICKS_CLIENT_SECRET`. + Set the following environment variables: ```bash export PGHOST=your-lakebase-host.databricks.com export PGDATABASE=your_database_name export LAKEBASE_ENDPOINT=projects/6bef4151-4b5d-4147-b4d0-c2f4fd5b40db/branches/br-broad-pine-y12n6gnv/endpoints/ep-summer-frost-y131l3vx -export PGUSER=your_user # optionally, defaults to DATABRICKS_CLIENT_ID -export PGSSLMODE=require +export PGUSER=your_user # optional: defaults to DATABRICKS_CLIENT_ID ``` -To find your `LAKEBASE_ENDPOINT`, run the Databricks CLI and use the `name` field from the output: +Your `LAKEBASE_ENDPOINT` has the structure `projects/${project}/branches/${branch}/endpoints/${endpoint}`. To find it, run the Databricks CLI and use the `name` field: ```bash databricks postgres list-endpoints projects/{project-id}/branches/{branch-id} ``` -You can obtain the Project ID and Branch ID from the Lakebase Autoscaling UI, like the "Branch Overview" page. (Project list -> Project dashboard -> Branch overview). +You can obtain the Project ID and Branch ID from the Lakebase Autoscaling UI, like the "Branch Overview" page (Project list -> Project dashboard -> Branch overview). Then use the driver: @@ -79,6 +83,28 @@ The driver supports Databricks authentication via: See [Databricks authentication docs](https://docs.databricks.com/en/dev-tools/auth/index.html) or [Lakebase Autoscaling authentication docs](https://docs.databricks.com/aws/en/oltp/projects/authentication#overview) for more information. +## Token Refresh & Retries + +OAuth tokens (1-hour lifetime) are refreshed automatically. Two strategies are available via the `refresh` option: + +| Mode | Behavior | Best for | +| ------------------- | ----------------------------------------------------------------------- | ----------------------------------------- | +| `"eager"` (default) | Fetches a token at pool creation and refreshes in the background before expiry | Time-sensitive, user-facing apps and APIs | +| `"lazy"` | Fetches on first connection and refreshes on demand when nearing expiry | Background jobs, infrequent connections | + +```typescript +const pool = createLakebasePool({ refresh: "lazy" }); +``` + +The eager refresh uses an `unref`'d timer (never keeps the process alive on its own) and is cancelled automatically when you call `pool.end()`. + +Transient credential-fetch failures are retried automatically (default schedule `[50, 500, 5000]` ms). Customize or disable with the `retry` option: + +```typescript +createLakebasePool({ retry: { schedule: [100, 1000] } }); // custom backoff +createLakebasePool({ retry: { schedule: [] } }); // disable retries +``` + ## PostgreSQL Username Resolution The driver resolves the PostgreSQL username (`user` configuration option) using the following priority order: @@ -89,7 +115,7 @@ The driver resolves the PostgreSQL username (`user` configuration option) using If none of these are set, the driver throws a `ConfigurationError`. -### Automatic resolution via Workspace API +### Automatic Resolution via Workspace API For human users authenticating with a PAT token or browser OAuth via `~/.databrickscfg`, none of the above are typically set. Use `getUsernameWithApiLookup` to automatically fetch the username from the Databricks workspace before creating the pool: @@ -121,6 +147,9 @@ const pool = createLakebasePool({ user }); | `max` | - | Max pool connections | `10` | | `idleTimeoutMillis` | - | Idle connection timeout | `30000` | | `connectionTimeoutMillis` | - | Connection timeout | `10000` | +| `refresh` | - | Token refresh strategy (`eager`/`lazy`) | `eager` | +| `earlyRefreshMs` | - | How long before expiry to refresh | `120000` | +| `retry` | - | Retry schedule for credential fetches | `[50, 500, 5000]` | | `logger` | - | Logger instance or config | `{ error: true }` | ## Logging diff --git a/packages/lakebase/package.json b/packages/lakebase/package.json index 3fdbfecc1..5c6c06914 100644 --- a/packages/lakebase/package.json +++ b/packages/lakebase/package.json @@ -49,6 +49,7 @@ "release:sbom": "pnpm exec cdxgen -t js --no-recurse --required-only -o tmp/sbom.cdx.json ." }, "dependencies": { + "@databricks/lakebase-auth": "workspace:*", "@databricks/sdk-experimental": "0.17.0", "pg": "8.18.0", "@opentelemetry/api": "1.9.0" diff --git a/packages/lakebase/src/__tests__/pool-config.test.ts b/packages/lakebase/src/__tests__/pool-config.test.ts new file mode 100644 index 000000000..0f1269db3 --- /dev/null +++ b/packages/lakebase/src/__tests__/pool-config.test.ts @@ -0,0 +1,236 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { + createTelemetryFetchCredential, + getLakebaseOrmConfig, + loggerToOnLog, +} from "../pool-config"; +import type { DriverTelemetry } from "../telemetry"; +import { createTokenRefreshCallback } from "../token-refresh"; + +// Keep the real auth logic, but stub the network-bound credential generation +// and workspace-client creation. +vi.mock("@databricks/lakebase-auth", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + generateDatabaseCredential: vi.fn(), + getWorkspaceClient: vi.fn(() => ({ config: { host: "test" } })), + }; +}); + +import { + generateDatabaseCredential, + getWorkspaceClient, +} from "@databricks/lakebase-auth"; + +const mockGenerate = vi.mocked(generateDatabaseCredential); +const mockGetWorkspaceClient = vi.mocked(getWorkspaceClient); + +/** A telemetry double whose tracer simply runs the active-span callback. */ +function fakeTelemetry(): DriverTelemetry { + return { + tracer: { + startActiveSpan: (_n: string, _o: unknown, fn: (s: never) => T): T => + fn({ + setAttribute: vi.fn(), + setStatus: vi.fn(), + end: vi.fn(), + recordException: vi.fn(), + } as never), + }, + meter: {}, + tokenRefreshDuration: { record: vi.fn() }, + queryDuration: { record: vi.fn() }, + poolErrors: { add: vi.fn() }, + } as unknown as DriverTelemetry; +} + +const ENV_KEYS = [ + "PGHOST", + "PGDATABASE", + "LAKEBASE_ENDPOINT", + "PGUSER", + "PGPORT", + "PGSSLMODE", + "DATABRICKS_CLIENT_ID", +] as const; +const original: Record = {}; + +beforeEach(() => { + vi.clearAllMocks(); + for (const key of ENV_KEYS) original[key] = process.env[key]; + process.env.PGHOST = "ep-test.databricks.com"; + process.env.PGDATABASE = "databricks_postgres"; + process.env.LAKEBASE_ENDPOINT = "projects/p/branches/b/endpoints/e"; + process.env.PGUSER = "user@example.com"; + mockGenerate.mockResolvedValue({ + token: "oauth-token", + expire_time: new Date(Date.now() + 3_600_000).toISOString(), + }); +}); + +afterEach(() => { + for (const key of ENV_KEYS) { + if (original[key] === undefined) delete process.env[key]; + else process.env[key] = original[key]; + } +}); + +describe("loggerToOnLog", () => { + test("returns undefined without a logger", () => { + expect(loggerToOnLog(undefined)).toBeUndefined(); + }); + + test("bridges structured log calls to the matching logger level", () => { + const logger = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + const onLog = loggerToOnLog(logger); + + onLog?.("warn", "msg %s", "arg"); + expect(logger.warn).toHaveBeenCalledWith("msg %s", "arg"); + onLog?.("error", "boom"); + expect(logger.error).toHaveBeenCalledWith("boom"); + }); +}); + +describe("getLakebaseOrmConfig", () => { + test("renames user to username and normalizes a boolean ssl (disable)", () => { + const cfg = getLakebaseOrmConfig({ + password: "static", + sslMode: "disable", + }); + + expect(cfg.username).toBe("user@example.com"); + expect("user" in cfg).toBe(false); + expect(cfg.password).toBe("static"); + expect(cfg.ssl).toBe(false); + }); + + test("normalizes an object ssl to just rejectUnauthorized", () => { + const cfg = getLakebaseOrmConfig({ + password: "static", + sslMode: "verify-full", + }); + + expect(cfg.ssl).toEqual({ rejectUnauthorized: true }); + }); + + test("exposes a password callback for the OAuth path", () => { + const cfg = getLakebaseOrmConfig({ + workspaceClient: {} as never, + refresh: "lazy", + }); + + expect(cfg.username).toBe("user@example.com"); + expect(typeof cfg.password).toBe("function"); + }); +}); + +describe("createTelemetryFetchCredential", () => { + test("fetches, traces, and maps the credential", async () => { + const telemetry = fakeTelemetry(); + const fetchCredential = createTelemetryFetchCredential({ + userConfig: { workspaceClient: {} as never }, + endpoint: "projects/p/branches/b/endpoints/e", + telemetry, + }); + + const credential = await fetchCredential(); + expect(credential.token).toBe("oauth-token"); + expect(typeof credential.expiresAt).toBe("number"); + expect(telemetry.tokenRefreshDuration.record).toHaveBeenCalledWith( + expect.any(Number), + ); + }); + + test("logs and rethrows when the workspace client cannot be created", async () => { + mockGetWorkspaceClient.mockImplementationOnce(() => { + throw new Error("no auth"); + }); + const logger = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + + const fetchCredential = createTelemetryFetchCredential({ + userConfig: {}, // no workspaceClient -> getWorkspaceClient is invoked + endpoint: "projects/p/branches/b/endpoints/e", + telemetry: fakeTelemetry(), + logger, + }); + + await expect(fetchCredential()).rejects.toThrow("no auth"); + expect(logger.error).toHaveBeenCalledWith( + expect.stringContaining("workspace client"), + expect.any(Error), + ); + }); + + test("logs and rethrows when credential generation fails", async () => { + mockGenerate.mockRejectedValueOnce(new Error("api down")); + const logger = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + + const fetchCredential = createTelemetryFetchCredential({ + userConfig: { workspaceClient: {} as never }, + endpoint: "projects/p/branches/b/endpoints/e", + telemetry: fakeTelemetry(), + logger, + }); + + await expect(fetchCredential()).rejects.toThrow("api down"); + expect(logger.error).toHaveBeenCalled(); + }); +}); + +describe("createTokenRefreshCallback (deprecated)", () => { + test("returns a lazy password callback that resolves to the token", async () => { + const password = createTokenRefreshCallback({ + userConfig: { workspaceClient: {} as never }, + endpoint: "projects/p/branches/b/endpoints/e", + telemetry: fakeTelemetry(), + }); + + // Lazy: nothing fetched until the callback is invoked. + expect(mockGenerate).not.toHaveBeenCalled(); + await expect(password()).resolves.toBe("oauth-token"); + expect(mockGenerate).toHaveBeenCalledTimes(1); + }); + + test("deduplicates concurrent callback invocations", async () => { + mockGenerate.mockImplementation( + () => + new Promise((resolve) => + setTimeout( + () => + resolve({ + token: "deduped", + expire_time: new Date(Date.now() + 3_600_000).toISOString(), + }), + 10, + ), + ), + ); + + const password = createTokenRefreshCallback({ + userConfig: { workspaceClient: {} as never }, + endpoint: "projects/p/branches/b/endpoints/e", + telemetry: fakeTelemetry(), + }); + + const [a, b, c] = await Promise.all([password(), password(), password()]); + expect([a, b, c]).toEqual(["deduped", "deduped", "deduped"]); + expect(mockGenerate).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/lakebase/src/__tests__/pool.test.ts b/packages/lakebase/src/__tests__/pool.test.ts index 97fb862b6..08cc388dc 100644 --- a/packages/lakebase/src/__tests__/pool.test.ts +++ b/packages/lakebase/src/__tests__/pool.test.ts @@ -32,12 +32,15 @@ vi.mock("pg", () => { }; }); -// Mock generateDatabaseCredential -vi.mock("../credentials", async (importOriginal) => { - const actual = await importOriginal(); +// Mock the auth package: keep the real config/provider/caching logic, but stub +// out the network-bound credential generation and workspace-client creation. +vi.mock("@databricks/lakebase-auth", async (importOriginal) => { + const actual = + await importOriginal(); return { ...actual, generateDatabaseCredential: vi.fn(), + getWorkspaceClient: vi.fn(() => ({ config: { host: "test" } })), }; }); @@ -117,8 +120,9 @@ describe("createLakebasePool", () => { process.env.PGUSER = "test-user@example.com"; // Setup mock for generateDatabaseCredential - const utils = await import("../credentials"); - mockGenerateCredential = utils.generateDatabaseCredential as any; + const auth = await import("@databricks/lakebase-auth"); + mockGenerateCredential = + auth.generateDatabaseCredential as unknown as ReturnType; mockGenerateCredential.mockResolvedValue({ token: "test-oauth-token-12345", expire_time: new Date(Date.now() + 3600000).toISOString(), // 1 hour from now @@ -237,10 +241,13 @@ describe("createLakebasePool", () => { test("should configure SSL based on sslMode", () => { const pool = createLakebasePool({ workspaceClient: {} as any, - sslMode: "require", + sslMode: "verify-full", }); - expect(pool.options.ssl).toEqual({ rejectUnauthorized: true }); + expect(pool.options.ssl).toEqual({ + rejectUnauthorized: true, + serverName: "ep-test.database.us-east-1.databricks.com", + }); }); test("should allow custom SSL configuration", () => { @@ -254,17 +261,23 @@ describe("createLakebasePool", () => { }); test("should throw on invalid PGSSLMODE", () => { - process.env.PGSSLMODE = "verify-full"; + process.env.PGSSLMODE = "wide-open"; expect(() => createLakebasePool({ workspaceClient: {} as any, }), - ).toThrow("one of: require, disable, prefer"); + ).toThrow("one of: verify-full, verify-ca, require, prefer, disable"); }); test("should accept valid PGSSLMODE values", () => { - for (const mode of ["require", "disable", "prefer"]) { + for (const mode of [ + "verify-full", + "verify-ca", + "require", + "prefer", + "disable", + ]) { process.env.PGSSLMODE = mode; expect(() => @@ -285,7 +298,17 @@ describe("createLakebasePool", () => { expect(typeof pool.options.password).toBe("function"); }); - test("should fetch OAuth token when password callback is invoked", async () => { + test("should eagerly fetch a token on pool creation by default", () => { + const workspaceClient = { config: { host: "test" } } as any; + + createLakebasePool({ workspaceClient }); + + // Eager mode (the default) fetches a token at construction time, + // without the password callback being invoked. + expect(mockGenerateCredential).toHaveBeenCalledTimes(1); + }); + + test("should fetch OAuth token when password callback is invoked (lazy)", async () => { const workspaceClient = { test: "client", config: { host: "test" }, @@ -293,9 +316,12 @@ describe("createLakebasePool", () => { const pool = createLakebasePool({ workspaceClient, endpoint: "projects/test/branches/main/endpoints/primary", + refresh: "lazy", }); - // Invoke the password callback + // Lazy mode does not fetch until the callback is invoked. + expect(mockGenerateCredential).not.toHaveBeenCalled(); + const passwordFn = pool.options.password as () => Promise; const password = await passwordFn(); @@ -309,6 +335,7 @@ describe("createLakebasePool", () => { const workspaceClient = { config: { host: "test" } } as any; const pool = createLakebasePool({ workspaceClient, + refresh: "lazy", }); const passwordFn = pool.options.password as () => Promise; @@ -340,6 +367,7 @@ describe("createLakebasePool", () => { const pool = createLakebasePool({ workspaceClient, + refresh: "lazy", }); const passwordFn = pool.options.password as () => Promise; @@ -362,6 +390,8 @@ describe("createLakebasePool", () => { const pool = createLakebasePool({ workspaceClient, + refresh: "lazy", + retry: { schedule: [] }, // disable retries for a fast, deterministic failure }); const passwordFn = pool.options.password as () => Promise; @@ -388,6 +418,7 @@ describe("createLakebasePool", () => { const pool = createLakebasePool({ workspaceClient, + refresh: "lazy", }); const passwordFn = pool.options.password as () => Promise; @@ -405,6 +436,15 @@ describe("createLakebasePool", () => { expect(p2).toBe("deduped-token"); expect(p3).toBe("deduped-token"); }); + + test("should clear background refresh when the pool is closed", async () => { + const workspaceClient = { config: { host: "test" } } as any; + const pool = createLakebasePool({ workspaceClient }); + + // pool.end is wrapped to dispose the eager refresh timer. + expect(pool.end.name).toBe("endWithDispose"); + await expect(pool.end()).resolves.toBeUndefined(); + }); }); describe("workspace client", () => { @@ -578,6 +618,7 @@ describe("createLakebasePool", () => { const workspaceClient = { config: { host: "test" } } as any; const pool = createLakebasePool({ workspaceClient, + refresh: "lazy", }); const passwordFn = pool.options.password as () => Promise; @@ -591,6 +632,7 @@ describe("createLakebasePool", () => { const workspaceClient = { config: { host: "test" } } as any; const pool = createLakebasePool({ workspaceClient, + refresh: "lazy", }); const passwordFn = pool.options.password as () => Promise; diff --git a/packages/lakebase/src/index.ts b/packages/lakebase/src/index.ts index 452173173..0eb5959a3 100644 --- a/packages/lakebase/src/index.ts +++ b/packages/lakebase/src/index.ts @@ -1,5 +1,8 @@ -export { getUsernameWithApiLookup, getWorkspaceClient } from "./config"; -export { generateDatabaseCredential } from "./credentials"; +export { + generateDatabaseCredential, + getUsernameWithApiLookup, + getWorkspaceClient, +} from "@databricks/lakebase-auth"; export { createLakebasePool } from "./pool"; export { getLakebaseOrmConfig, @@ -14,7 +17,9 @@ export type { LakebasePoolConfig, Logger, LoggerConfig, + RefreshMode, RequestedClaims, RequestedResource, + RetryOptions, } from "./types"; export { RequestedClaimsPermissionSet } from "./types"; diff --git a/packages/lakebase/src/pool-config.ts b/packages/lakebase/src/pool-config.ts index a843bc173..f0dc65d73 100644 --- a/packages/lakebase/src/pool-config.ts +++ b/packages/lakebase/src/pool-config.ts @@ -1,27 +1,146 @@ +import { + type Credential, + type FetchCredential, + generateDatabaseCredential, + getPgConfig, + getWorkspaceClient, + type LogFn, +} from "@databricks/lakebase-auth"; +import type { WorkspaceClient } from "@databricks/sdk-experimental"; import type { PoolConfig } from "pg"; -import { getUsernameSync, parsePoolConfig } from "./config"; -import { type DriverTelemetry, initTelemetry } from "./telemetry"; -import { createTokenRefreshCallback } from "./token-refresh"; +import { + type DriverTelemetry, + initTelemetry, + SpanStatusCode, +} from "./telemetry"; import type { LakebasePoolConfig, Logger } from "./types"; +/** Default pool sizing values for the Lakebase connector */ +const poolDefaults = { + max: 10, + idleTimeoutMillis: 30_000, + connectionTimeoutMillis: 10_000, +}; + +/** Bridge a {@link Logger} to the auth package's structured `onLog` callback. */ +export function loggerToOnLog(logger?: Logger): LogFn | undefined { + if (!logger) return undefined; + return (level, message, ...args) => logger[level](message, ...args); +} + /** - * Map an SSL mode string to the corresponding `pg` SSL configuration. - * - * - `"require"` -- SSL enabled with certificate verification - * - `"prefer"` -- SSL enabled without certificate verification (try SSL, accept any cert) - * - `"disable"` -- SSL disabled + * Build a credential fetcher that wraps {@link generateDatabaseCredential} with + * OpenTelemetry tracing/metrics and logger integration. Injected into the auth + * package's password provider so observability stays in `@databricks/lakebase`. */ -function mapSslConfig( - sslMode: "require" | "prefer" | "disable", -): PoolConfig["ssl"] { - switch (sslMode) { - case "require": - return { rejectUnauthorized: true }; - case "prefer": - return { rejectUnauthorized: false }; - case "disable": - return false; - } +export function createTelemetryFetchCredential(deps: { + userConfig: Partial; + endpoint: string; + telemetry: DriverTelemetry; + logger?: Logger; +}): FetchCredential { + // Lazily initialize the workspace client on first password fetch. + let workspaceClient: WorkspaceClient | null = + deps.userConfig.workspaceClient ?? null; + + return async (): Promise => { + if (!workspaceClient) { + try { + workspaceClient = getWorkspaceClient(deps.userConfig); + } catch (error) { + deps.logger?.error("Failed to initialize workspace client: %O", error); + throw error; + } + } + const client = workspaceClient; + + const startTime = Date.now(); + try { + return await deps.telemetry.tracer.startActiveSpan( + "lakebase.token.refresh", + { attributes: { "lakebase.endpoint": deps.endpoint } }, + async (span) => { + const credential = await generateDatabaseCredential(client, { + endpoint: deps.endpoint, + ...(deps.userConfig.claims + ? { claims: deps.userConfig.claims } + : {}), + }); + const expiresAt = new Date(credential.expire_time).getTime(); + span.setAttribute( + "lakebase.token.expires_at", + new Date(expiresAt).toISOString(), + ); + span.setStatus({ code: SpanStatusCode.OK }); + span.end(); + return { token: credential.token, expiresAt }; + }, + ); + } catch (error) { + deps.logger?.error("Failed to fetch OAuth token: %O", { + error, + message: error instanceof Error ? error.message : String(error), + endpoint: deps.endpoint, + }); + throw error; + } finally { + deps.telemetry.tokenRefreshDuration.record(Date.now() - startTime); + } + }; +} + +/** + * Build the Lakebase `pg.PoolConfig` along with a disposer that stops the + * background token refresh (used by {@link createLakebasePool} to clean up on + * `pool.end()`). + */ +export function buildLakebasePgConfig( + config?: Partial, + telemetry?: DriverTelemetry, + logger?: Logger, +): { poolConfig: PoolConfig; dispose: () => void } { + const userConfig = config ?? {}; + const onLog = loggerToOnLog(logger); + + // Resolve the endpoint (with env fallback) so the telemetry-wrapped fetcher + // is used for OAuth auth regardless of whether the endpoint came from config + // or LAKEBASE_ENDPOINT. When neither endpoint nor password is set, getPgConfig + // below throws the appropriate configuration error. + const endpoint = userConfig.endpoint ?? process.env.LAKEBASE_ENDPOINT; + + // Only wrap with telemetry when using OAuth (no native password provided). + const fetchCredential = + userConfig.password === undefined && endpoint !== undefined + ? createTelemetryFetchCredential({ + userConfig, + endpoint, + telemetry: telemetry ?? initTelemetry(), + logger, + }) + : undefined; + + const { dispose, ...pg } = getPgConfig({ + ...userConfig, + fetchCredential, + onLog, + }); + + const poolConfig: PoolConfig = { + host: pg.host, + port: pg.port, + user: pg.user, + database: pg.database, + password: pg.password, + ssl: pg.ssl, + max: userConfig.max ?? poolDefaults.max, + idleTimeoutMillis: + userConfig.idleTimeoutMillis ?? poolDefaults.idleTimeoutMillis, + connectionTimeoutMillis: + userConfig.connectionTimeoutMillis ?? + poolDefaults.connectionTimeoutMillis, + }; + + return { poolConfig, dispose }; } /** @@ -44,36 +163,7 @@ export function getLakebasePgConfig( telemetry?: DriverTelemetry, logger?: Logger, ): PoolConfig { - const userConfig = config ?? {}; - const poolConfig = parsePoolConfig(userConfig); - const username = getUsernameSync(userConfig); - - let passwordConfig: string | (() => string | Promise) | undefined; - - if (userConfig.password !== undefined) { - passwordConfig = userConfig.password; - } else if (poolConfig.endpoint) { - // endpoint is guaranteed here -- parsePoolConfig() throws if - // neither endpoint nor password is provided - passwordConfig = createTokenRefreshCallback({ - userConfig, - endpoint: poolConfig.endpoint, - telemetry: telemetry ?? initTelemetry(), - logger, - }); - } - - return { - host: poolConfig.host, - port: poolConfig.port, - user: username, - database: poolConfig.database, - password: passwordConfig, - ssl: poolConfig.ssl ?? mapSslConfig(poolConfig.sslMode), - max: poolConfig.max, - idleTimeoutMillis: poolConfig.idleTimeoutMillis, - connectionTimeoutMillis: poolConfig.connectionTimeoutMillis, - }; + return buildLakebasePgConfig(config, telemetry, logger).poolConfig; } /** diff --git a/packages/lakebase/src/pool.ts b/packages/lakebase/src/pool.ts index c07c114cf..20921ec64 100644 --- a/packages/lakebase/src/pool.ts +++ b/packages/lakebase/src/pool.ts @@ -1,6 +1,6 @@ import pg from "pg"; import { resolveLogger } from "./logger"; -import { getLakebasePgConfig } from "./pool-config"; +import { buildLakebasePgConfig } from "./pool-config"; import { attachPoolMetrics, initTelemetry, @@ -63,12 +63,27 @@ export function createLakebasePool( const telemetry = initTelemetry(); - const poolConfig = getLakebasePgConfig(userConfig, telemetry, logger); + const { poolConfig, dispose } = buildLakebasePgConfig( + userConfig, + telemetry, + logger, + ); const pool = new pg.Pool(poolConfig); attachPoolMetrics(pool, telemetry, logger); + // Stop the background token refresh (eager mode) when the pool is closed. + const origEnd = pool.end.bind(pool); + pool.end = function endWithDispose( + ...args: unknown[] + ): ReturnType { + dispose(); + return (origEnd as (...a: unknown[]) => unknown)(...args) as ReturnType< + typeof pool.end + >; + } as typeof pool.end; + // Wrap pool.query to track query duration and create trace spans. // pg.Pool.query has 15+ overloads that are difficult to type-preserve, // so we use a loosely-typed wrapper and cast back. diff --git a/packages/lakebase/src/token-refresh.ts b/packages/lakebase/src/token-refresh.ts index 1da75ba38..27bd9bc48 100644 --- a/packages/lakebase/src/token-refresh.ts +++ b/packages/lakebase/src/token-refresh.ts @@ -1,13 +1,8 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; -import { getWorkspaceClient } from "./config"; -import { generateDatabaseCredential } from "./credentials"; -import { type DriverTelemetry, SpanStatusCode } from "./telemetry"; +import { createPasswordProvider } from "@databricks/lakebase-auth"; +import { createTelemetryFetchCredential, loggerToOnLog } from "./pool-config"; +import type { DriverTelemetry } from "./telemetry"; import type { LakebasePoolConfig, Logger } from "./types"; -// 2-minute buffer before token expiration to prevent race conditions -// Lakebase tokens expire after 1 hour, so we refresh when ~58 minutes remain -const CACHE_BUFFER_MS = 2 * 60 * 1000; - export interface TokenRefreshDeps { userConfig: Partial; endpoint: string; @@ -15,23 +10,12 @@ export interface TokenRefreshDeps { logger?: Logger; } -/** Fetch a fresh OAuth token from Databricks */ -async function refreshToken( - workspaceClient: WorkspaceClient, - endpoint: string, -): Promise<{ token: string; expiresAt: number }> { - const credential = await generateDatabaseCredential(workspaceClient, { - endpoint, - }); - - return { - token: credential.token, - expiresAt: new Date(credential.expire_time).getTime(), - }; -} - /** - * Build the password callback with token caching, deduplication, and telemetry. + * Build a password callback with token caching, deduplication, and telemetry. + * + * @deprecated Prefer `createPasswordProvider` / `getPgConfig` from + * `@databricks/lakebase-auth`. Retained for backwards compatibility; uses lazy + * (on-demand) refresh to match the previous behavior. * * The returned async function is called by `pg.Pool` each time a new connection * is established. It caches OAuth tokens and deduplicates concurrent refresh @@ -40,76 +24,11 @@ async function refreshToken( export function createTokenRefreshCallback( deps: TokenRefreshDeps, ): () => Promise { - let cachedToken: string | undefined; - let tokenExpiresAt = 0; - let workspaceClient: WorkspaceClient | null = null; - let refreshPromise: Promise | null = null; - - return async (): Promise => { - // Lazily initialize workspace client on first password fetch - if (!workspaceClient) { - try { - workspaceClient = getWorkspaceClient(deps.userConfig); - } catch (error) { - deps.logger?.error("Failed to initialize workspace client: %O", error); - throw error; - } - } - - const now = Date.now(); - const hasValidToken = cachedToken && now < tokenExpiresAt - CACHE_BUFFER_MS; - if (hasValidToken) { - // Return cached token if still valid (with buffer) - const expiresIn = Math.round((tokenExpiresAt - now) / 1000 / 60); - deps.logger?.debug( - "Using cached OAuth token (expires in %d minutes at %s)", - expiresIn, - new Date(tokenExpiresAt).toISOString(), - ); - return cachedToken as string; - } - - const client = workspaceClient; - - // Deduplicate concurrent refresh requests - if (!refreshPromise) { - refreshPromise = (async () => { - const startTime = Date.now(); - try { - const result = await deps.telemetry.tracer.startActiveSpan( - "lakebase.token.refresh", - { - attributes: { "lakebase.endpoint": deps.endpoint }, - }, - async (span) => { - const tokenResult = await refreshToken(client, deps.endpoint); - span.setAttribute( - "lakebase.token.expires_at", - new Date(tokenResult.expiresAt).toISOString(), - ); - span.setStatus({ code: SpanStatusCode.OK }); - span.end(); - return tokenResult; - }, - ); - - cachedToken = result.token; - tokenExpiresAt = result.expiresAt; - return cachedToken; - } catch (error) { - deps.logger?.error("Failed to fetch OAuth token: %O", { - error, - message: error instanceof Error ? error.message : String(error), - endpoint: deps.endpoint, - }); - throw error; - } finally { - deps.telemetry.tokenRefreshDuration.record(Date.now() - startTime); - refreshPromise = null; - } - })(); - } - - return refreshPromise; - }; + const fetchCredential = createTelemetryFetchCredential(deps); + const { password } = createPasswordProvider({ + fetchCredential, + mode: "lazy", + onLog: loggerToOnLog(deps.logger), + }); + return password; } diff --git a/packages/lakebase/src/types.ts b/packages/lakebase/src/types.ts index 18e51e15e..c04a54d48 100644 --- a/packages/lakebase/src/types.ts +++ b/packages/lakebase/src/types.ts @@ -1,6 +1,24 @@ +import type { + RefreshMode, + RequestedClaims, + RetryOptions, +} from "@databricks/lakebase-auth"; import type { WorkspaceClient } from "@databricks/sdk-experimental"; import type { PoolConfig } from "pg"; +// Re-export the auth/credential types so existing `@databricks/lakebase` +// imports keep working after the split into `@databricks/lakebase-auth`. +export type { + Credential, + DatabaseCredential, + GenerateDatabaseCredentialRequest, + RefreshMode, + RequestedClaims, + RequestedResource, + RetryOptions, +} from "@databricks/lakebase-auth"; +export { RequestedClaimsPermissionSet } from "@databricks/lakebase-auth"; + /** * Optional logger interface for the Lakebase driver. * When not provided, the driver operates silently (no logging). @@ -74,12 +92,45 @@ export interface LakebasePoolConfig extends PoolConfig { endpoint?: string; /** - * SSL mode for the connection (convenience helper) - * Can also be set via PGSSLMODE environment variable + * SSL mode for the connection (convenience helper). Can also be set via + * PGSSLMODE. All values other than "disable" are treated as "verify-full" + * with system root certs. + * + * @default "verify-full" + */ + sslMode?: "verify-full" | "verify-ca" | "require" | "prefer" | "disable"; + + /** + * Optional UC claims for fine-grained Unity Catalog table permissions on the + * generated Postgres token. + */ + claims?: RequestedClaims[]; + + /** + * Token refresh strategy. + * + * - `"eager"` (default): fetch a token immediately and refresh it in the + * background before it expires. Best for time-sensitive, user-facing apps. + * - `"lazy"`: fetch a token on first use and refresh it on demand. + * + * @default "eager" + */ + refresh?: RefreshMode; + + /** + * How long before token expiry to refresh, in milliseconds. + * + * @default 120000 (2 minutes) + */ + earlyRefreshMs?: number; + + /** + * Retry options for transient credential-fetch failures (e.g. the OAuth + * server being briefly unreachable). * - * @default "require" + * @default { schedule: [50, 500, 5000] } */ - sslMode?: "require" | "disable" | "prefer"; + retry?: RetryOptions; /** * Telemetry configuration @@ -115,97 +166,3 @@ export interface LakebasePoolConfig extends PoolConfig { */ logger?: Logger | LoggerConfig; } - -// --------------------------------------------------------------------------- -// Authentication types for Lakebase Postgres OAuth token generation -// @see https://docs.databricks.com/aws/en/oltp/projects/authentication -// --------------------------------------------------------------------------- - -/** - * Database credentials with OAuth token for Postgres connection - */ -export interface DatabaseCredential { - /** OAuth token to use as the password when connecting to Postgres */ - token: string; - - /** - * Token expiration time in UTC (ISO 8601 format) - * Tokens expire after 1 hour from generation - * @example "2026-02-06T17:07:00Z" - */ - expire_time: string; -} - -/** - * Permission set for Unity Catalog table access - */ -export enum RequestedClaimsPermissionSet { - /** - * Read-only access to specified UC tables - */ - READ_ONLY = "READ_ONLY", -} - -/** - * Resource to request permissions for in Unity Catalog - */ -export interface RequestedResource { - /** - * Unity Catalog table name to request access to - * @example "catalog.schema.table" - */ - table_name?: string; - - /** - * Generic resource name for non-table resources - */ - unspecified_resource_name?: string; -} - -/** - * Optional claims for fine-grained Unity Catalog table permissions - * When specified, the returned token will be scoped to only the requested tables - */ -export interface RequestedClaims { - /** - * Permission level to request - */ - permission_set?: RequestedClaimsPermissionSet; - - /** - * List of UC resources to request access to - */ - resources?: RequestedResource[]; -} - -/** - * Request parameters for generating database OAuth credentials - */ -export interface GenerateDatabaseCredentialRequest { - /** - * Endpoint resource path. Retrieve using the Databricks CLI: - * ``` - * databricks postgres list-endpoints projects/{project-id}/branches/{branch-id} - * ``` - * Use the `name` field from the output. - * - * @example "projects/{project-id}/branches/{branch-id}/endpoints/{endpoint-identifier}" - */ - endpoint: string; - - /** - * Optional claims for fine-grained UC table permissions. - * When specified, the token will only grant access to the specified tables. - * - * @example - * ```typescript - * { - * claims: [{ - * permission_set: RequestedClaimsPermissionSet.READ_ONLY, - * resources: [{ table_name: "catalog.schema.users" }] - * }] - * } - * ``` - */ - claims?: RequestedClaims[]; -} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 0d08e5808..b3de9bb22 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -529,6 +529,9 @@ importers: packages/lakebase: dependencies: + '@databricks/lakebase-auth': + specifier: workspace:* + version: link:../lakebase-auth '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 @@ -543,6 +546,12 @@ importers: specifier: 8.16.0 version: 8.16.0 + packages/lakebase-auth: + dependencies: + '@databricks/sdk-experimental': + specifier: 0.17.0 + version: 0.17.0 + packages/shared: dependencies: '@ast-grep/napi': diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 3b88e3501..016833852 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -1,4 +1,11 @@ packages: - - "packages/*" - - "apps/*" - - "docs" + - packages/* + - apps/* + - docs + +onlyBuiltDependencies: + - bufferutil + - core-js + - core-js-pure + - esbuild + - protobufjs diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index a4aa1c26e..f8c4a2278 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -256,7 +256,7 @@ "env": "PGSSLMODE", "description": "Postgres SSL mode. Auto-injected by the platform at deploy time.", "localOnly": true, - "value": "require", + "value": "verify-full", "origin": "platform" } } diff --git a/tools/dist-lakebase-auth.ts b/tools/dist-lakebase-auth.ts new file mode 100644 index 000000000..e1bd82c1f --- /dev/null +++ b/tools/dist-lakebase-auth.ts @@ -0,0 +1,35 @@ +import fs from "node:fs"; +import path from "node:path"; +import { parseArgs } from "node:util"; + +const __dirname = path.dirname(new URL(import.meta.url).pathname); +const { values } = parseArgs({ + options: { + prerelease: { type: "string" }, + }, +}); +const prerelease = values.prerelease; + +fs.mkdirSync("tmp", { recursive: true }); + +const pkg = JSON.parse(fs.readFileSync("package.json", "utf-8")); +if (prerelease) { + pkg.version = `${pkg.version}-pr.${prerelease}`; +} + +pkg.exports = pkg.publishConfig.exports; +delete pkg.publishConfig.exports; + +fs.writeFileSync("tmp/package.json", JSON.stringify(pkg, null, 2)); + +fs.cpSync("dist", "tmp/dist", { recursive: true }); + +// Use the package's own README.md if present, otherwise fall back to the root one +const localReadme = "README.md"; +const rootReadme = path.join(__dirname, "../README.md"); +fs.copyFileSync( + fs.existsSync(localReadme) ? localReadme : rootReadme, + "tmp/README.md", +); +fs.copyFileSync(path.join(__dirname, "../LICENSE"), "tmp/LICENSE"); +fs.copyFileSync(path.join(__dirname, "../DCO"), "tmp/DCO"); diff --git a/tools/sync-lakebase-auth-version.ts b/tools/sync-lakebase-auth-version.ts new file mode 100644 index 000000000..6c0b3072f --- /dev/null +++ b/tools/sync-lakebase-auth-version.ts @@ -0,0 +1,26 @@ +#!/usr/bin/env tsx +/** + * Syncs the version to the lakebase-auth package. + * Used by the prepare-release-lakebase-auth workflow after version bump. + */ + +/** + * NOTE: This script is also used by the private secure release repo + * during the finalize step. Changes here affect the release pipeline. + */ + +import { readFileSync, writeFileSync } from "node:fs"; +import { join } from "node:path"; + +const ROOT = process.cwd(); +const version = process.argv[2]; +if (!version) { + console.error("Usage: sync-lakebase-auth-version.ts "); + process.exit(1); +} + +const pkgJsonPath = join(ROOT, "packages/lakebase-auth/package.json"); +const pkgJson = JSON.parse(readFileSync(pkgJsonPath, "utf-8")); +pkgJson.version = version; +writeFileSync(pkgJsonPath, `${JSON.stringify(pkgJson, null, 2)}\n`); +console.log(`✓ packages/lakebase-auth/package.json → ${version}`); diff --git a/vitest.config.ts b/vitest.config.ts index 8c2893b01..9532a5f19 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -50,6 +50,14 @@ export default defineConfig({ environment: "node", }, }, + { + plugins: [tsconfigPaths()], + test: { + name: "lakebase-auth", + root: "./packages/lakebase-auth", + environment: "node", + }, + }, { plugins: [tsconfigPaths()], test: {