Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -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/`
Expand Down
31 changes: 31 additions & 0 deletions docs/docs/api/appkit/Function.getUsernameWithApiLookup.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
# Function: getUsernameWithApiLookup()

```ts
function getUsernameWithApiLookup(config?: Partial<LakebasePoolConfig>): Promise<string | undefined>;
```

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`\>
4 changes: 2 additions & 2 deletions docs/docs/api/appkit/Function.getWorkspaceClient.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# Function: getWorkspaceClient()

```ts
function getWorkspaceClient(config: Partial<LakebasePoolConfig>): Promise<WorkspaceClient>;
function getWorkspaceClient(config: Partial<LakebasePoolConfig>): WorkspaceClient;
```

Get workspace client from config or SDK default auth chain
Expand All @@ -14,4 +14,4 @@ Get workspace client from config or SDK default auth chain

## Returns

`Promise`\<`WorkspaceClient`\>
`WorkspaceClient`
1 change: 1 addition & 0 deletions docs/docs/api/appkit/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 |
5 changes: 5 additions & 0 deletions docs/docs/api/appkit/typedoc-sidebar.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
1 change: 1 addition & 0 deletions packages/appkit/src/connectors/lakebase/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ export {
generateDatabaseCredential,
getLakebaseOrmConfig,
getLakebasePgConfig,
getUsernameWithApiLookup,
getWorkspaceClient,
type LakebasePoolConfig,
type Logger,
Expand Down
1 change: 1 addition & 0 deletions packages/appkit/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ export {
generateDatabaseCredential,
getLakebaseOrmConfig,
getLakebasePgConfig,
getUsernameWithApiLookup,
getWorkspaceClient,
RequestedClaimsPermissionSet,
} from "./connectors/lakebase";
Expand Down
39 changes: 31 additions & 8 deletions packages/lakebase/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand Down Expand Up @@ -85,14 +79,43 @@ 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 |
| ------------------------- | ---------------------------------- | --------------------------------------- | ----------------------- |
| `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` |
Expand Down
2 changes: 1 addition & 1 deletion packages/lakebase/src/__tests__/pool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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", () => {
Expand Down
41 changes: 38 additions & 3 deletions packages/lakebase/src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<LakebasePoolConfig>,
): Promise<WorkspaceClient> {
): WorkspaceClient {
// Priority 1: Explicit workspaceClient in config
if (config.workspaceClient) {
return config.workspaceClient;
Expand Down Expand Up @@ -140,6 +140,41 @@ export function getUsernameSync(config: Partial<LakebasePoolConfig>): 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<LakebasePoolConfig>,
): Promise<string | undefined> {
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;
}
}
2 changes: 1 addition & 1 deletion packages/lakebase/src/index.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
export { getWorkspaceClient } from "./config";
export { getUsernameWithApiLookup, getWorkspaceClient } from "./config";
export { generateDatabaseCredential } from "./credentials";
export { createLakebasePool } from "./pool";
export {
Expand Down
2 changes: 1 addition & 1 deletion packages/lakebase/src/token-refresh.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down