diff --git a/CLAUDE.md b/CLAUDE.md index 9a17f870..b77b7322 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -257,7 +257,7 @@ const result = await pool.query('SELECT * FROM users'); ``` **ORM Integration:** -Works with Drizzle, Prisma, TypeORM - see [Lakebase Integration Docs](docs/docs/integrations/lakebase.md) for examples. +Works with Drizzle, Prisma, TypeORM - see the [`@databricks/lakebase` README](https://github.com/databricks/appkit/blob/main/packages/lakebase/README.md) for examples. **Architecture:** - Connector files: `packages/appkit/src/connectors/lakebase/` diff --git a/docs/docs/api/appkit/Function.getUsernameWithApiLookup.md b/docs/docs/api/appkit/Function.getUsernameWithApiLookup.md new file mode 100644 index 00000000..29b92f2e --- /dev/null +++ b/docs/docs/api/appkit/Function.getUsernameWithApiLookup.md @@ -0,0 +1,31 @@ +# Function: getUsernameWithApiLookup() + +```ts +function getUsernameWithApiLookup(config?: Partial): Promise; +``` + +Resolves the PostgreSQL username for a Lakebase connection. + +Extends getUsernameSync with an async fallback that fetches the +current user's identity from the Databricks workspace API. Use this when +you don't have an explicit username configured and want automatic resolution +(e.g. human users authenticating via PAT or browser OAuth in ~/.databrickscfg). + +Resolution priority: +1. `config.user` — explicit config value +2. `PGUSER` env var +3. `DATABRICKS_CLIENT_ID` env var (service principals) +4. `currentUser.me()` — fetched from Databricks workspace API + +Returns `undefined` if all attempts fail rather than throwing, so the +caller can decide whether to proceed or surface an error. + +## Parameters + +| Parameter | Type | +| ------ | ------ | +| `config?` | `Partial`\<[`LakebasePoolConfig`](Interface.LakebasePoolConfig.md)\> | + +## Returns + +`Promise`\<`string` \| `undefined`\> diff --git a/docs/docs/api/appkit/Function.getWorkspaceClient.md b/docs/docs/api/appkit/Function.getWorkspaceClient.md index f8b856b2..9511b944 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): Promise; +function getWorkspaceClient(config: Partial): WorkspaceClient; ``` Get workspace client from config or SDK default auth chain @@ -14,4 +14,4 @@ Get workspace client from config or SDK default auth chain ## Returns -`Promise`\<`WorkspaceClient`\> +`WorkspaceClient` diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index 91f9a66e..f64043b0 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -73,5 +73,6 @@ plugin architecture, and React integration. | [getLakebasePgConfig](Function.getLakebasePgConfig.md) | Get Lakebase connection configuration for PostgreSQL clients. | | [getPluginManifest](Function.getPluginManifest.md) | Loads and validates the manifest from a plugin constructor. Normalizes string type/permission to strict ResourceType/ResourcePermission. | | [getResourceRequirements](Function.getResourceRequirements.md) | Gets the resource requirements from a plugin's manifest. | +| [getUsernameWithApiLookup](Function.getUsernameWithApiLookup.md) | Resolves the PostgreSQL username for a Lakebase connection. | | [getWorkspaceClient](Function.getWorkspaceClient.md) | Get workspace client from config or SDK default auth chain | | [isSQLTypeMarker](Function.isSQLTypeMarker.md) | Type guard to check if a value is a SQL type marker | diff --git a/docs/docs/api/appkit/typedoc-sidebar.ts b/docs/docs/api/appkit/typedoc-sidebar.ts index 3421d7ee..53f15b15 100644 --- a/docs/docs/api/appkit/typedoc-sidebar.ts +++ b/docs/docs/api/appkit/typedoc-sidebar.ts @@ -240,6 +240,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Function.getResourceRequirements", label: "getResourceRequirements" }, + { + type: "doc", + id: "api/appkit/Function.getUsernameWithApiLookup", + label: "getUsernameWithApiLookup" + }, { type: "doc", id: "api/appkit/Function.getWorkspaceClient", diff --git a/packages/appkit/src/connectors/lakebase/index.ts b/packages/appkit/src/connectors/lakebase/index.ts index e33ac077..cc121a28 100644 --- a/packages/appkit/src/connectors/lakebase/index.ts +++ b/packages/appkit/src/connectors/lakebase/index.ts @@ -32,6 +32,7 @@ export { generateDatabaseCredential, getLakebaseOrmConfig, getLakebasePgConfig, + getUsernameWithApiLookup, getWorkspaceClient, type LakebasePoolConfig, type Logger, diff --git a/packages/appkit/src/index.ts b/packages/appkit/src/index.ts index 8ba528ef..84221411 100644 --- a/packages/appkit/src/index.ts +++ b/packages/appkit/src/index.ts @@ -27,6 +27,7 @@ export { generateDatabaseCredential, getLakebaseOrmConfig, getLakebasePgConfig, + getUsernameWithApiLookup, getWorkspaceClient, RequestedClaimsPermissionSet, } from "./connectors/lakebase"; diff --git a/packages/lakebase/README.md b/packages/lakebase/README.md index 1d1d8162..99f3beb6 100644 --- a/packages/lakebase/README.md +++ b/packages/lakebase/README.md @@ -42,13 +42,7 @@ To find your `LAKEBASE_ENDPOINT`, run the Databricks CLI and use the `name` fiel 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). When using the driver as a part of Databricks Apps, the `LAKEBASE_ENDPOINT` is automatically injected using `fromValue`: - -```yaml -env: - - name: LAKEBASE_ENDPOINT - valueFrom: database # Lakebase Autoscaling database resource name -``` +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: @@ -85,6 +79,35 @@ 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. +## PostgreSQL Username Resolution + +The driver resolves the PostgreSQL username (`user` configuration option) using the following priority order: + +1. `config.user` — explicit value passed to `createLakebasePool` +2. `PGUSER` environment variable +3. `DATABRICKS_CLIENT_ID` environment variable (service principals using OAuth M2M) + +If none of these are set, the driver throws a `ConfigurationError`. + +### 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: + +```typescript +import { createLakebasePool, getUsernameWithApiLookup } from "@databricks/lakebase"; + +// Tries config/env vars first, then falls back to currentUser.me() API call +const user = await getUsernameWithApiLookup(); + +const pool = createLakebasePool({ user }); +``` + +`getUsernameWithApiLookup` extends the sync resolution above with a fourth step: + +4. `currentUser.me()` — fetches the current user's identity from the Databricks workspace API (works with PAT tokens and browser OAuth in `~/.databrickscfg`) + +> **Note:** `getUsernameWithApiLookup` makes a network call to the Databricks workspace API when the sync resolution steps (config, `PGUSER`, `DATABRICKS_CLIENT_ID`) all fail. Call it once during initialization, not on every request. + ## Configuration | Option | Environment Variable | Description | Default | @@ -92,7 +115,7 @@ See [Databricks authentication docs](https://docs.databricks.com/en/dev-tools/au | `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 | Auto-detected | +| `user` | `PGUSER` or `DATABRICKS_CLIENT_ID` | Username or service principal ID | See [Username Resolution](#username-resolution)| | `port` | `PGPORT` | Port number | `5432` | | `sslMode` | `PGSSLMODE` | SSL mode | `require` | | `max` | - | Max pool connections | `10` | diff --git a/packages/lakebase/src/__tests__/pool.test.ts b/packages/lakebase/src/__tests__/pool.test.ts index 5cd3db30..97fb862b 100644 --- a/packages/lakebase/src/__tests__/pool.test.ts +++ b/packages/lakebase/src/__tests__/pool.test.ts @@ -209,7 +209,7 @@ describe("createLakebasePool", () => { createLakebasePool({ workspaceClient: {} as any, }), - ).toThrow("PGUSER, DATABRICKS_CLIENT_ID, or config.user"); + ).toThrow("config.user, PGUSER or DATABRICKS_CLIENT_ID"); }); test("should use DATABRICKS_CLIENT_ID as fallback for user", () => { diff --git a/packages/lakebase/src/config.ts b/packages/lakebase/src/config.ts index 43a44051..ef055edc 100644 --- a/packages/lakebase/src/config.ts +++ b/packages/lakebase/src/config.ts @@ -106,9 +106,9 @@ function validateSslMode(value: string | undefined): SslMode | undefined { } /** Get workspace client from config or SDK default auth chain */ -export async function getWorkspaceClient( +export function getWorkspaceClient( config: Partial, -): Promise { +): WorkspaceClient { // Priority 1: Explicit workspaceClient in config if (config.workspaceClient) { return config.workspaceClient; @@ -140,6 +140,41 @@ export function getUsernameSync(config: Partial): string { } throw ConfigurationError.missingEnvVar( - "PGUSER, DATABRICKS_CLIENT_ID, or config.user", + "config.user, PGUSER or DATABRICKS_CLIENT_ID", ); } + +/** + * Resolves the PostgreSQL username for a Lakebase connection. + * + * Extends {@link getUsernameSync} with an async fallback that fetches the + * current user's identity from the Databricks workspace API. Use this when + * you don't have an explicit username configured and want automatic resolution + * (e.g. human users authenticating via PAT or browser OAuth in ~/.databrickscfg). + * + * Resolution priority: + * 1. `config.user` — explicit config value + * 2. `PGUSER` env var + * 3. `DATABRICKS_CLIENT_ID` env var (service principals) + * 4. `currentUser.me()` — fetched from Databricks workspace API + * + * Returns `undefined` if all attempts fail rather than throwing, so the + * caller can decide whether to proceed or surface an error. + */ +export async function getUsernameWithApiLookup( + config?: Partial, +): Promise { + try { + return getUsernameSync(config ?? {}); + } catch { + // sync resolution failed, try workspace API + } + + try { + const workspaceClient = getWorkspaceClient(config ?? {}); + const me = await workspaceClient.currentUser.me(); + return me.userName ?? undefined; + } catch { + return undefined; + } +} diff --git a/packages/lakebase/src/index.ts b/packages/lakebase/src/index.ts index b3c5a288..45217317 100644 --- a/packages/lakebase/src/index.ts +++ b/packages/lakebase/src/index.ts @@ -1,4 +1,4 @@ -export { getWorkspaceClient } from "./config"; +export { getUsernameWithApiLookup, getWorkspaceClient } from "./config"; export { generateDatabaseCredential } from "./credentials"; export { createLakebasePool } from "./pool"; export { diff --git a/packages/lakebase/src/token-refresh.ts b/packages/lakebase/src/token-refresh.ts index d22bc7d4..1da75ba3 100644 --- a/packages/lakebase/src/token-refresh.ts +++ b/packages/lakebase/src/token-refresh.ts @@ -49,7 +49,7 @@ export function createTokenRefreshCallback( // Lazily initialize workspace client on first password fetch if (!workspaceClient) { try { - workspaceClient = await getWorkspaceClient(deps.userConfig); + workspaceClient = getWorkspaceClient(deps.userConfig); } catch (error) { deps.logger?.error("Failed to initialize workspace client: %O", error); throw error;