From 49bea9d51cdbd9565e07a37948115f4594405e7b Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 7 Oct 2026 18:49:28 +0200 Subject: [PATCH 1/5] feat: drop @databricks/sdk-experimental from @databricks/lakebase Lakebase now takes a structural workspace client: either the legacy WorkspaceClient shape (currentUser.me + apiClient.request) or anything with an authenticated request() (the AppKit modular client). Credential generation and the SCIM Me lookup go through request() when present. The default client (no workspaceClient passed) uses the modular sdk-core profile resolution and sdk-auth default credentials, loaded via dynamic import so the CJS build keeps working with the ESM-only SDK. Token refresh timing and caching are unchanged. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/lakebase/package.json | 3 +- .../src/__tests__/credentials.test.ts | 140 ++++++++++-------- packages/lakebase/src/config.ts | 88 ++++++++++- packages/lakebase/src/credentials.ts | 19 ++- packages/lakebase/src/index.ts | 4 + packages/lakebase/src/token-refresh.ts | 12 +- packages/lakebase/src/types.ts | 43 +++++- pnpm-lock.yaml | 9 +- 8 files changed, 237 insertions(+), 81 deletions(-) diff --git a/packages/lakebase/package.json b/packages/lakebase/package.json index fdf8f471d..3d8c95914 100644 --- a/packages/lakebase/package.json +++ b/packages/lakebase/package.json @@ -56,7 +56,8 @@ "release:sbom": "pnpm exec cdxgen -t js --no-recurse --required-only -o tmp/sbom.cdx.json ." }, "dependencies": { - "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-auth": "0.51.0", + "@databricks/sdk-core": "0.51.0", "pg": "8.18.0", "@opentelemetry/api": "1.9.0" }, diff --git a/packages/lakebase/src/__tests__/credentials.test.ts b/packages/lakebase/src/__tests__/credentials.test.ts index 3ca9e7930..a723f48f1 100644 --- a/packages/lakebase/src/__tests__/credentials.test.ts +++ b/packages/lakebase/src/__tests__/credentials.test.ts @@ -1,47 +1,24 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; -import { ApiClient, Config } from "@databricks/sdk-experimental"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { getUsernameWithApiLookup } from "../config"; import { generateDatabaseCredential } from "../credentials"; import { type DatabaseCredential, + type LegacyWorkspaceClientLike, RequestedClaimsPermissionSet, } from "../types"; -// Mock the @databricks/sdk-experimental module -vi.mock("@databricks/sdk-experimental", () => { - const mockRequest = vi.fn(); - - return { - Config: vi.fn(), - ApiClient: vi.fn().mockImplementation(() => ({ - request: mockRequest, - })), - }; -}); - describe("Lakebase Authentication", () => { - let mockWorkspaceClient: WorkspaceClient; - let mockApiClient: ApiClient; + let mockWorkspaceClient: LegacyWorkspaceClientLike; + let mockApiClient: LegacyWorkspaceClientLike["apiClient"]; beforeEach(() => { vi.clearAllMocks(); - - // Get the mocked ApiClient constructor - const ApiClientConstructor = ApiClient as unknown as ReturnType< - typeof vi.fn - >; - mockApiClient = new ApiClientConstructor( - new Config({ host: "https://test.databricks.com" }), - ); - - // Setup mock workspace client with apiClient + mockApiClient = { request: vi.fn() }; mockWorkspaceClient = { - config: { - host: "https://test.databricks.com", - }, + currentUser: { me: vi.fn() }, apiClient: mockApiClient, - } as WorkspaceClient; + }; }); describe("generateDatabaseCredential", () => { @@ -145,44 +122,87 @@ describe("Lakebase Authentication", () => { }), ).rejects.toThrow("API request failed"); }); + }); - it("should use correct workspace host for API calls", async () => { - const customHost = "https://custom-workspace.databricks.com"; + describe("request-capable (AppKit modular) client", () => { + const credential: DatabaseCredential = { + token: "modular-token", + expire_time: "2026-02-06T18:00:00Z", + }; + + it("posts the snake_case request body through request()", async () => { + const request = vi.fn(async () => Response.json(credential)); + const claims = [ + { + permission_set: RequestedClaimsPermissionSet.READ_ONLY, + resources: [{ table_name: "catalog.schema.users" }], + }, + ]; - // Create a new mock API client for the custom workspace - const ApiClientConstructor = ApiClient as unknown as ReturnType< - typeof vi.fn - >; - const customApiClient = new ApiClientConstructor( - new Config({ host: customHost }), + const result = await generateDatabaseCredential( + { request }, + { endpoint: "projects/p/branches/main/endpoints/primary", claims }, ); - const customWorkspaceClient = { - config: { host: customHost }, - apiClient: customApiClient, - } as WorkspaceClient; + expect(result).toEqual(credential); + expect(request).toHaveBeenCalledWith({ + method: "POST", + path: "/api/2.0/postgres/credentials", + headers: { + Accept: "application/json", + "Content-Type": "application/json", + }, + body: JSON.stringify({ + endpoint: "projects/p/branches/main/endpoints/primary", + claims, + }), + }); + }); - const mockCredential: DatabaseCredential = { - token: "mock-token", - expire_time: "2026-02-06T18:00:00Z", - }; + it("prefers request() over a legacy apiClient on the same object", async () => { + // The AppKit facade exposes both; the modular path must win. + const request = vi.fn(async () => Response.json(credential)); + const client = { ...mockWorkspaceClient, request }; - vi.mocked(customApiClient.request).mockResolvedValue(mockCredential); + await generateDatabaseCredential(client, { endpoint: "e" }); - await generateDatabaseCredential(customWorkspaceClient, { - endpoint: "projects/test/branches/main/endpoints/primary", - }); + expect(request).toHaveBeenCalledOnce(); + expect(mockApiClient.request).not.toHaveBeenCalled(); + }); - // Verify the request was made with the correct workspace client - expect(customApiClient.request).toHaveBeenCalledWith({ - path: "/api/2.0/postgres/credentials", - method: "POST", - headers: expect.any(Headers), - raw: false, - payload: { - endpoint: "projects/test/branches/main/endpoints/primary", - }, - }); + it("rejects a malformed credential response", async () => { + const request = vi.fn(async () => Response.json({ token: "t" })); + await expect( + generateDatabaseCredential({ request }, { endpoint: "e" }), + ).rejects.toThrow(); + }); + + it("resolves the username via the SCIM Me endpoint", async () => { + const prev = { + PGUSER: process.env.PGUSER, + DATABRICKS_CLIENT_ID: process.env.DATABRICKS_CLIENT_ID, + }; + delete process.env.PGUSER; + delete process.env.DATABRICKS_CLIENT_ID; + try { + const request = vi.fn(async () => + Response.json({ userName: "someone@example.com" }), + ); + await expect( + getUsernameWithApiLookup({ workspaceClient: { request } }), + ).resolves.toBe("someone@example.com"); + expect(request).toHaveBeenCalledWith( + expect.objectContaining({ + method: "GET", + path: "/api/2.0/preview/scim/v2/Me", + }), + ); + } finally { + Object.assign(process.env, prev); + for (const [k, v] of Object.entries(prev)) { + if (v === undefined) delete process.env[k]; + } + } }); }); }); diff --git a/packages/lakebase/src/config.ts b/packages/lakebase/src/config.ts index bca7c9fa9..8a1de2c88 100644 --- a/packages/lakebase/src/config.ts +++ b/packages/lakebase/src/config.ts @@ -1,8 +1,11 @@ -import { WorkspaceClient } from "@databricks/sdk-experimental"; import type pg from "pg"; import { ConfigurationError, ValidationError } from "./errors"; -import type { LakebasePoolConfig } from "./types"; +import type { + LakebasePoolConfig, + LakebaseWorkspaceClient, + RequestCapableWorkspaceClient, +} from "./types"; /** Default configuration values for the Lakebase connector */ const defaults = { @@ -109,16 +112,78 @@ function validateSslMode(value: string | undefined): SslMode | undefined { /** Get workspace client from config or SDK default auth chain */ export function getWorkspaceClient( config: Partial, -): WorkspaceClient { +): LakebaseWorkspaceClient { // Priority 1: Explicit workspaceClient in config if (config.workspaceClient) { return config.workspaceClient; } - // Priority 2: Create with SDK default auth chain - // Use empty config to let SDK use .databrickscfg, DATABRICKS_HOST, DATABRICKS_TOKEN, etc. + // Priority 2: SDK default auth chain (.databrickscfg, DATABRICKS_HOST, DATABRICKS_TOKEN, etc.) // NOTE: config.host is the PostgreSQL host (PGHOST), not the Databricks workspace host - return new WorkspaceClient({}); + return createDefaultWorkspaceClient(); +} + +/** True when the client exposes the modular `request` capability (AppKit client). */ +export function isRequestCapable( + client: LakebaseWorkspaceClient, +): client is RequestCapableWorkspaceClient { + return ( + typeof (client as RequestCapableWorkspaceClient).request === "function" + ); +} + +/** + * Default-auth client built on the modular SDK. The SDK is ESM-only, so it is + * loaded with a dynamic import (keeps the CJS build working) and only when no + * client was passed in. Auth is resolved once, lazily, and retried on failure. + */ +function createDefaultWorkspaceClient(): RequestCapableWorkspaceClient { + let resolved: + | Promise<{ + host: string; + authHeaders: () => Promise<{ key: string; value: string }[]>; + }> + | undefined; + const resolveOnce = () => { + resolved ??= (async () => { + const [{ resolve }, { defaultCredentials }] = await Promise.all([ + import("@databricks/sdk-core/profiles"), + import("@databricks/sdk-auth/credentials"), + ]); + const profile = await resolve(); + const rawHost = profile.host?.trim(); + if (!rawHost) throw ConfigurationError.missingEnvVar("DATABRICKS_HOST"); + // The legacy SDK prepended https:// to a scheme-less host; the modular SDK does not. + const host = ( + /^https?:\/\//i.test(rawHost) ? rawHost : `https://${rawHost}` + ).replace(/\/+$/, ""); + const credentials = defaultCredentials({ profile: { ...profile, host } }); + return { host, authHeaders: () => credentials.authHeaders() }; + })().catch((error) => { + resolved = undefined; + throw error; + }); + return resolved; + }; + + return { + async request(req) { + const { host, authHeaders } = await resolveOnce(); + const headers = new Headers(req.headers); + for (const h of await authHeaders()) headers.set(h.key, h.value); + const response = await fetch(new URL(req.path, host), { + method: req.method, + headers, + body: req.body, + }); + if (!response.ok) { + throw new Error( + `Databricks API ${req.method} ${req.path} failed with ${response.status}: ${await response.text()}`, + ); + } + return response; + }, + }; } /** Get username synchronously from config or environment */ @@ -173,7 +238,16 @@ export async function getUsernameWithApiLookup( try { const workspaceClient = getWorkspaceClient(config ?? {}); - const me = await workspaceClient.currentUser.me(); + if (!isRequestCapable(workspaceClient)) { + const me = await workspaceClient.currentUser.me(); + return me.userName ?? undefined; + } + const response = await workspaceClient.request({ + method: "GET", + path: "/api/2.0/preview/scim/v2/Me", + headers: { Accept: "application/json" }, + }); + const me = (await response.json()) as { userName?: string }; return me.userName ?? undefined; } catch { return undefined; diff --git a/packages/lakebase/src/credentials.ts b/packages/lakebase/src/credentials.ts index ae1a8c9b6..eca506444 100644 --- a/packages/lakebase/src/credentials.ts +++ b/packages/lakebase/src/credentials.ts @@ -1,9 +1,9 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; - +import { isRequestCapable } from "./config"; import { ValidationError } from "./errors"; import type { DatabaseCredential, GenerateDatabaseCredentialRequest, + LakebaseWorkspaceClient, } from "./types"; /** @@ -48,11 +48,24 @@ import type { * ``` */ export async function generateDatabaseCredential( - workspaceClient: WorkspaceClient, + workspaceClient: LakebaseWorkspaceClient, request: GenerateDatabaseCredentialRequest, ): Promise { const apiPath = "/api/2.0/postgres/credentials"; + if (isRequestCapable(workspaceClient)) { + const response = await workspaceClient.request({ + method: "POST", + path: apiPath, + headers: { + Accept: "application/json", + "Content-Type": "application/json", + }, + body: JSON.stringify(request), + }); + return validateCredentialResponse(await response.json()); + } + const response = await workspaceClient.apiClient.request({ path: apiPath, method: "POST", diff --git a/packages/lakebase/src/index.ts b/packages/lakebase/src/index.ts index b2b9f5e9c..8e00ad4c1 100644 --- a/packages/lakebase/src/index.ts +++ b/packages/lakebase/src/index.ts @@ -8,10 +8,14 @@ export { createTokenRefreshCallback } from "./token-refresh"; export type { DatabaseCredential, GenerateDatabaseCredentialRequest, + LakebaseWorkspaceClient, + LakebaseWorkspaceRequest, + LegacyWorkspaceClientLike, LakebasePoolConfig, Logger, LoggerConfig, RequestedClaims, + RequestCapableWorkspaceClient, RequestedResource, } from "./types"; export { RequestedClaimsPermissionSet } from "./types"; diff --git a/packages/lakebase/src/token-refresh.ts b/packages/lakebase/src/token-refresh.ts index 793e25611..7fc2963e3 100644 --- a/packages/lakebase/src/token-refresh.ts +++ b/packages/lakebase/src/token-refresh.ts @@ -1,9 +1,11 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; - import { getWorkspaceClient } from "./config"; import { generateDatabaseCredential } from "./credentials"; import { type DriverTelemetry, SpanStatusCode } from "./telemetry"; -import type { LakebasePoolConfig, Logger } from "./types"; +import type { + LakebasePoolConfig, + LakebaseWorkspaceClient, + 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 @@ -18,7 +20,7 @@ export interface TokenRefreshDeps { /** Fetch a fresh OAuth token from Databricks */ async function refreshToken( - workspaceClient: WorkspaceClient, + workspaceClient: LakebaseWorkspaceClient, endpoint: string, ): Promise<{ token: string; expiresAt: number }> { const credential = await generateDatabaseCredential(workspaceClient, { @@ -43,7 +45,7 @@ export function createTokenRefreshCallback( ): () => Promise { let cachedToken: string | undefined; let tokenExpiresAt = 0; - let workspaceClient: WorkspaceClient | null = null; + let workspaceClient: LakebaseWorkspaceClient | null = null; let refreshPromise: Promise | null = null; return async (): Promise => { diff --git a/packages/lakebase/src/types.ts b/packages/lakebase/src/types.ts index 18e51e15e..ed9105252 100644 --- a/packages/lakebase/src/types.ts +++ b/packages/lakebase/src/types.ts @@ -1,4 +1,3 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import type { PoolConfig } from "pg"; /** @@ -55,7 +54,7 @@ export interface LakebasePoolConfig extends PoolConfig { * * Note: If password is provided, OAuth auth is not used */ - workspaceClient?: WorkspaceClient; + workspaceClient?: LakebaseWorkspaceClient; /** * Endpoint resource path for OAuth token generation. @@ -209,3 +208,43 @@ export interface GenerateDatabaseCredentialRequest { */ claims?: RequestedClaims[]; } + +/** A raw REST call against the workspace host. */ +export interface LakebaseWorkspaceRequest { + method: string; + /** Path on the workspace host, e.g. `/api/2.0/postgres/credentials`. */ + path: string; + headers?: Record; + body?: string; +} + +/** + * Client exposing an authenticated raw `request` (the AppKit workspace client). + * Must throw on a non-2xx response. + */ +export interface RequestCapableWorkspaceClient { + request(req: LakebaseWorkspaceRequest): Promise; +} + +/** The subset of the legacy `@databricks/sdk-experimental` `WorkspaceClient` lakebase uses. */ +export interface LegacyWorkspaceClientLike { + currentUser: { me(): Promise<{ userName?: string }> }; + apiClient: { + request(options: { + path: string; + method: string; + headers: Headers; + raw: boolean; + payload?: unknown; + }): Promise; + }; +} + +/** + * Workspace client accepted by lakebase: a legacy `@databricks/sdk-experimental` + * `WorkspaceClient`, or anything with an authenticated `request` (the AppKit + * modular workspace client). Structural, so lakebase depends on neither SDK's client. + */ +export type LakebaseWorkspaceClient = + | RequestCapableWorkspaceClient + | LegacyWorkspaceClientLike; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index d09f608b1..2f3964679 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -610,9 +610,12 @@ importers: packages/lakebase: dependencies: - '@databricks/sdk-experimental': - specifier: 0.17.0 - version: 0.17.0 + '@databricks/sdk-auth': + specifier: 0.51.0 + version: 0.51.0 + '@databricks/sdk-core': + specifier: 0.51.0 + version: 0.51.0 '@opentelemetry/api': specifier: 1.9.0 version: 1.9.0 From e01cf0f8635cd0210214796ba88b4dd3454a9dd5 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 7 Oct 2026 18:49:45 +0200 Subject: [PATCH 2/5] feat(appkit): pass the modular workspace client to lakebase The lakebase plugin and initializeLakebasePool hand the AppKit client to @databricks/lakebase directly instead of toLegacyWorkspaceClient(). Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../src/connectors/lakebase/endpoint-host.ts | 32 ++++++--- .../appkit/src/connectors/lakebase/index.ts | 2 +- .../lakebase/tests/endpoint-host.test.ts | 66 +++++++++++++++++++ .../lakebase/tests/initialize-pool.test.ts | 8 +-- .../appkit/src/plugins/lakebase/lakebase.ts | 6 +- .../tests/lakebase-agent-tool.test.ts | 6 +- packages/lakebase/src/types.ts | 18 +++-- 7 files changed, 110 insertions(+), 28 deletions(-) diff --git a/packages/appkit/src/connectors/lakebase/endpoint-host.ts b/packages/appkit/src/connectors/lakebase/endpoint-host.ts index a3020901d..f2516d661 100644 --- a/packages/appkit/src/connectors/lakebase/endpoint-host.ts +++ b/packages/appkit/src/connectors/lakebase/endpoint-host.ts @@ -74,16 +74,30 @@ async function compareHosts( reason, ); try { + const path = `/api/2.0/postgres/${endpoint}`; + const lookup = + "request" in client + ? // AppKit's modular client: throws ApiError (with statusCode) on non-2xx. + // `signal` isn't in lakebase's public request type but AppKit honors it. + client + .request({ + method: "GET", + path, + headers: { Accept: "application/json" }, + signal: controller.signal, + } as Parameters[0]) + .then((res) => res.json()) + : client.apiClient.request( + { + path, + method: "GET", + headers: new Headers({ Accept: "application/json" }), + raw: false, + }, + contextFromAbortSignal(controller.signal), + ); const response = await Promise.race([ - client.apiClient.request( - { - path: `/api/2.0/postgres/${endpoint}`, - method: "GET", - headers: new Headers({ Accept: "application/json" }), - raw: false, - }, - contextFromAbortSignal(controller.signal), - ), + lookup, new Promise((resolve) => { timer = setTimeout(() => { controller.abort(); diff --git a/packages/appkit/src/connectors/lakebase/index.ts b/packages/appkit/src/connectors/lakebase/index.ts index 08e169d14..67d0cca50 100644 --- a/packages/appkit/src/connectors/lakebase/index.ts +++ b/packages/appkit/src/connectors/lakebase/index.ts @@ -46,7 +46,7 @@ export async function initializeLakebasePool( const client = ServiceContext.isInitialized() ? ServiceContext.get().client : createWorkspaceClient({ clientOptions: getClientOptions() }); - resolved.workspaceClient = client.toLegacyWorkspaceClient(); + resolved.workspaceClient = client; } const [user] = await Promise.all([ getUsernameWithApiLookup(resolved), diff --git a/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts b/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts index b66c0567d..8afc790ef 100644 --- a/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts +++ b/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts @@ -316,3 +316,69 @@ describe("assertEndpointHostMatches", () => { ]); }); }); + +// AppKit passes its modular workspace client (with `request()`) to lakebase, +// so in apps the lookup goes through `request`, not the legacy `apiClient`. +describe("assertEndpointHostMatches with the modular request() client", () => { + function modularClient( + impl: (req: { signal?: AbortSignal }) => Promise, + ) { + const request = vi.fn(impl); + return { request, client: { request } as unknown as Client }; + } + + test("rejects a mismatch read through request()", async () => { + const { endpoint, host } = names(); + const { request, client } = modularClient(async () => + Response.json(endpointWith({ host: "ep-expected.database.example.com" })), + ); + vi.spyOn(console, "error").mockImplementation(() => undefined); + + await expect( + assertEndpointHostMatches({ endpoint, host, workspaceClient: client }), + ).rejects.toThrow(host); + expect(request).toHaveBeenCalledWith( + expect.objectContaining({ + method: "GET", + path: `/api/2.0/postgres/${endpoint}`, + signal: expect.any(AbortSignal), + }), + ); + }); + + test("warns when request() reports the endpoint no longer exists (404)", async () => { + const { endpoint, host } = names(); + const { client } = modularClient(async () => { + throw Object.assign(new Error("not found"), { statusCode: 404 }); + }); + + await assertEndpointHostMatches({ + endpoint, + host, + workspaceClient: client, + }); + expect(warnings()).toContain("was not found"); + }); + + test("aborts the request() lookup when the deadline expires", async () => { + vi.useFakeTimers(); + const { endpoint, host } = names(); + let seen: AbortSignal | undefined; + const { client } = modularClient( + (req) => + new Promise(() => { + seen = req.signal; + }), + ); + + const check = assertEndpointHostMatches({ + endpoint, + host, + workspaceClient: client, + }); + await vi.advanceTimersByTimeAsync(3_000); + await check; + expect(seen?.aborted).toBe(true); + expect(warnings()).toContain("timed out"); + }); +}); diff --git a/packages/appkit/src/connectors/lakebase/tests/initialize-pool.test.ts b/packages/appkit/src/connectors/lakebase/tests/initialize-pool.test.ts index 2c4f5989b..29715fce4 100644 --- a/packages/appkit/src/connectors/lakebase/tests/initialize-pool.test.ts +++ b/packages/appkit/src/connectors/lakebase/tests/initialize-pool.test.ts @@ -15,9 +15,7 @@ const mocks = vi.hoisted(() => { request, client, createPool: vi.fn(), - createWorkspaceClient: vi.fn(() => ({ - toLegacyWorkspaceClient: () => client, - })), + createWorkspaceClient: vi.fn(() => client), }; }); @@ -90,7 +88,7 @@ describe("AppKit Lakebase connector initialization", () => { }; vi.mocked(ServiceContext.isInitialized).mockReturnValue(true); vi.spyOn(ServiceContext, "get").mockReturnValue({ - client: { toLegacyWorkspaceClient: () => legacy }, + client: legacy, } as unknown as ReturnType); await initializeLakebasePool(); expect(poolConfig()).toMatchObject({ @@ -125,7 +123,7 @@ describe("AppKit Lakebase connector initialization", () => { const requestClient = { currentUser: { me: requestLookup } }; await runInUserContext( { - client: { toLegacyWorkspaceClient: () => requestClient }, + client: requestClient, userId: "request-user-id", userEmail: "request-user@example.test", workspaceId: Promise.resolve("workspace"), diff --git a/packages/appkit/src/plugins/lakebase/lakebase.ts b/packages/appkit/src/plugins/lakebase/lakebase.ts index dd4fe5a92..bb0d66a55 100644 --- a/packages/appkit/src/plugins/lakebase/lakebase.ts +++ b/packages/appkit/src/plugins/lakebase/lakebase.ts @@ -84,7 +84,7 @@ export class LakebasePlugin extends Plugin implements ToolProvider { this.config.pool?.workspaceClient ?? createWorkspaceClient({ clientOptions: getClientOptions(), - }).toLegacyWorkspaceClient(), + }), }; const [user] = await Promise.all([ getUsernameWithApiLookup(poolConfig), @@ -109,7 +109,7 @@ export class LakebasePlugin extends Plugin implements ToolProvider { const pool = oboManager.getPool( userKey, { - workspaceClient: ctx.client.toLegacyWorkspaceClient(), + workspaceClient: ctx.client, user: userKey, }, ctx.tokenFingerprint, @@ -305,7 +305,7 @@ export class LakebasePlugin extends Plugin implements ToolProvider { const user = ctx.principal.userEmail ?? ctx.principal.userId; return { ...this.config.pool, - workspaceClient: ctx.client.toLegacyWorkspaceClient(), + workspaceClient: ctx.client, user, }; } diff --git a/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts b/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts index 214d236ee..9b9bbf616 100644 --- a/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts +++ b/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts @@ -373,7 +373,7 @@ describe("LakebasePlugin - OBO via RoutingPool", () => { await plugin.setup(); const userCtx = { - client: { toLegacyWorkspaceClient: () => ({}) } as any, + client: {} as any, userId: "user-123", userEmail: "alice@example.com", workspaceId: Promise.resolve("ws-1"), @@ -402,7 +402,7 @@ describe("LakebasePlugin - OBO via RoutingPool", () => { await plugin.setup(); const userCtx = { - client: { toLegacyWorkspaceClient: () => ({}) } as any, + client: {} as any, userId: "user-123", userEmail: "alice@example.com", workspaceId: Promise.resolve("ws-1"), @@ -431,7 +431,7 @@ describe("LakebasePlugin - OBO via RoutingPool", () => { await plugin.setup(); const userCtx = { - client: { toLegacyWorkspaceClient: () => ({}) } as any, + client: {} as any, userId: "user-123", workspaceId: Promise.resolve("ws-1"), isUserContext: true as const, diff --git a/packages/lakebase/src/types.ts b/packages/lakebase/src/types.ts index ed9105252..041b02dfb 100644 --- a/packages/lakebase/src/types.ts +++ b/packages/lakebase/src/types.ts @@ -230,13 +230,17 @@ export interface RequestCapableWorkspaceClient { export interface LegacyWorkspaceClientLike { currentUser: { me(): Promise<{ userName?: string }> }; apiClient: { - request(options: { - path: string; - method: string; - headers: Headers; - raw: boolean; - payload?: unknown; - }): Promise; + request( + options: { + path: string; + method: string; + headers: Headers; + raw: boolean; + payload?: unknown; + }, + /** The legacy SDK's cancellation `Context` (optional). */ + context?: unknown, + ): Promise; }; } From 3ecf45f39f916687fa904b163647aef319d24e95 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 13 Oct 2026 17:35:26 +0200 Subject: [PATCH 3/5] feat(appkit): look up lakebase endpoint hosts through the typed postgres client The PGHOST check now calls postgres.getEndpoint on the AppKit workspace client instead of a raw GET, keeping the raw request and legacy paths for other clients. A 404 from the typed client's ApiError (httpStatusCode) is still treated as a missing endpoint. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../src/connectors/lakebase/endpoint-host.ts | 28 +++++- .../lakebase/tests/endpoint-host.test.ts | 91 +++++++++++++++++++ 2 files changed, 114 insertions(+), 5 deletions(-) diff --git a/packages/appkit/src/connectors/lakebase/endpoint-host.ts b/packages/appkit/src/connectors/lakebase/endpoint-host.ts index f2516d661..bbb6fd56f 100644 --- a/packages/appkit/src/connectors/lakebase/endpoint-host.ts +++ b/packages/appkit/src/connectors/lakebase/endpoint-host.ts @@ -58,6 +58,18 @@ export function assertEndpointHostMatches( return check; } +/** AppKit's workspace client exposes the typed Lakebase Postgres client. */ +function hasPostgresClient(client: Client): client is Client & { + postgres: { + getEndpoint( + req: { name: string }, + options: { signal: AbortSignal }, + ): Promise; + }; +} { + return "postgres" in client; +} + async function compareHosts( client: NonNullable, endpoint: string, @@ -75,9 +87,14 @@ async function compareHosts( ); try { const path = `/api/2.0/postgres/${endpoint}`; - const lookup = - "request" in client - ? // AppKit's modular client: throws ApiError (with statusCode) on non-2xx. + const lookup = hasPostgresClient(client) + ? // AppKit's workspace client: the typed Lakebase Postgres client. + client.postgres.getEndpoint( + { name: endpoint }, + { signal: controller.signal }, + ) + : "request" in client + ? // A raw request-capable client: throws ApiError (with statusCode) on non-2xx. // `signal` isn't in lakebase's public request type but AppKit honors it. client .request({ @@ -111,8 +128,9 @@ async function compareHosts( if ( typeof error === "object" && error !== null && - "statusCode" in error && - error.statusCode === 404 + // The typed client's sdk-core ApiError reports `httpStatusCode`. + (("statusCode" in error && error.statusCode === 404) || + ("httpStatusCode" in error && error.httpStatusCode === 404)) ) { logger.warn( "LAKEBASE_ENDPOINT %s was not found. Check the project, branch, and endpoint names; `databricks postgres list-endpoints projects/{project}/branches/{branch}` lists them with their hosts.", diff --git a/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts b/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts index 8afc790ef..d78cbb020 100644 --- a/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts +++ b/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts @@ -382,3 +382,94 @@ describe("assertEndpointHostMatches with the modular request() client", () => { expect(warnings()).toContain("timed out"); }); }); + +// AppKit's workspace client carries the typed Lakebase Postgres client, which +// wins over its raw `request()` seam. +describe("assertEndpointHostMatches with the typed postgres client", () => { + function postgresClient( + impl: ( + req: { name: string }, + options: { signal: AbortSignal }, + ) => Promise, + ) { + const getEndpoint = vi.fn(impl); + const request = vi.fn(); + return { + getEndpoint, + request, + client: { postgres: { getEndpoint }, request } as unknown as Client, + }; + } + + test("rejects a mismatch read through postgres.getEndpoint", async () => { + const { endpoint, host } = names(); + const { getEndpoint, request, client } = postgresClient(async () => ({ + status: { + hosts: { + host: "ep-expected.database.example.com", + readOnlyHost: "ep-ro.database.example.com", + }, + }, + })); + vi.spyOn(console, "error").mockImplementation(() => undefined); + + await expect( + assertEndpointHostMatches({ endpoint, host, workspaceClient: client }), + ).rejects.toThrow("ep-expected.database.example.com"); + expect(getEndpoint).toHaveBeenCalledWith( + { name: endpoint }, + { signal: expect.any(AbortSignal) }, + ); + expect(request).not.toHaveBeenCalled(); + }); + + test("accepts the camelCase read-only host", async () => { + const { endpoint, host } = names(); + const { client } = postgresClient(async () => ({ + status: { hosts: { host: "ep-rw.example.com", readOnlyHost: host } }, + })); + + await assertEndpointHostMatches({ + endpoint, + host, + workspaceClient: client, + }); + expect(warn).not.toHaveBeenCalled(); + }); + + test("warns when the typed client reports a 404 via httpStatusCode", async () => { + const { endpoint, host } = names(); + const { client } = postgresClient(async () => { + throw Object.assign(new Error("not found"), { httpStatusCode: 404 }); + }); + + await assertEndpointHostMatches({ + endpoint, + host, + workspaceClient: client, + }); + expect(warnings()).toContain("was not found"); + }); + + test("aborts the typed lookup when the deadline expires", async () => { + vi.useFakeTimers(); + const { endpoint, host } = names(); + let seen: AbortSignal | undefined; + const { client } = postgresClient( + (_req, options) => + new Promise(() => { + seen = options.signal; + }), + ); + + const check = assertEndpointHostMatches({ + endpoint, + host, + workspaceClient: client, + }); + await vi.advanceTimersByTimeAsync(3_000); + await check; + expect(seen?.aborted).toBe(true); + expect(warnings()).toContain("timed out"); + }); +}); From 4d8148e6575ffe066a42b8dada41df99838a35bb Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 13 Oct 2026 18:02:01 +0200 Subject: [PATCH 4/5] feat(appkit): use the typed postgres client for lakebase credentials Lakebase's request-capable client type gains an optional `postgres` (the subset of sdk-postgres PostgresClient it uses). When present, credentials go through postgres.generateDatabaseCredential, mapping the snake_case public request to the camelCase model and expireTime back to an ISO string. Claims using unspecified_resource_name fall back to the raw request: the typed model only has tableName and would drop them. The endpoint host check drops its raw request() branch and the cast; a request-only client now skips verification with a warning. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../src/connectors/lakebase/endpoint-host.ts | 57 ++++-------- .../lakebase/tests/endpoint-host.test.ts | 74 +++------------- .../src/__tests__/credentials.test.ts | 88 +++++++++++++++++++ packages/lakebase/src/credentials.ts | 44 ++++++++++ packages/lakebase/src/index.ts | 1 + packages/lakebase/src/types.ts | 25 ++++++ 6 files changed, 187 insertions(+), 102 deletions(-) diff --git a/packages/appkit/src/connectors/lakebase/endpoint-host.ts b/packages/appkit/src/connectors/lakebase/endpoint-host.ts index bbb6fd56f..dfa94c9e1 100644 --- a/packages/appkit/src/connectors/lakebase/endpoint-host.ts +++ b/packages/appkit/src/connectors/lakebase/endpoint-host.ts @@ -58,18 +58,6 @@ export function assertEndpointHostMatches( return check; } -/** AppKit's workspace client exposes the typed Lakebase Postgres client. */ -function hasPostgresClient(client: Client): client is Client & { - postgres: { - getEndpoint( - req: { name: string }, - options: { signal: AbortSignal }, - ): Promise; - }; -} { - return "postgres" in client; -} - async function compareHosts( client: NonNullable, endpoint: string, @@ -86,33 +74,24 @@ async function compareHosts( reason, ); try { - const path = `/api/2.0/postgres/${endpoint}`; - const lookup = hasPostgresClient(client) - ? // AppKit's workspace client: the typed Lakebase Postgres client. - client.postgres.getEndpoint( - { name: endpoint }, - { signal: controller.signal }, - ) - : "request" in client - ? // A raw request-capable client: throws ApiError (with statusCode) on non-2xx. - // `signal` isn't in lakebase's public request type but AppKit honors it. - client - .request({ - method: "GET", - path, - headers: { Accept: "application/json" }, - signal: controller.signal, - } as Parameters[0]) - .then((res) => res.json()) - : client.apiClient.request( - { - path, - method: "GET", - headers: new Headers({ Accept: "application/json" }), - raw: false, - }, - contextFromAbortSignal(controller.signal), - ); + const lookup = + "postgres" in client && client.postgres + ? // AppKit's workspace client: the typed Lakebase Postgres client. + client.postgres.getEndpoint( + { name: endpoint }, + { signal: controller.signal }, + ) + : "apiClient" in client + ? client.apiClient.request( + { + path: `/api/2.0/postgres/${endpoint}`, + method: "GET", + headers: new Headers({ Accept: "application/json" }), + raw: false, + }, + contextFromAbortSignal(controller.signal), + ) + : Promise.resolve(undefined); const response = await Promise.race([ lookup, new Promise((resolve) => { diff --git a/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts b/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts index d78cbb020..5fa2846b6 100644 --- a/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts +++ b/packages/appkit/src/connectors/lakebase/tests/endpoint-host.test.ts @@ -317,70 +317,18 @@ describe("assertEndpointHostMatches", () => { }); }); -// AppKit passes its modular workspace client (with `request()`) to lakebase, -// so in apps the lookup goes through `request`, not the legacy `apiClient`. -describe("assertEndpointHostMatches with the modular request() client", () => { - function modularClient( - impl: (req: { signal?: AbortSignal }) => Promise, - ) { - const request = vi.fn(impl); - return { request, client: { request } as unknown as Client }; - } - - test("rejects a mismatch read through request()", async () => { - const { endpoint, host } = names(); - const { request, client } = modularClient(async () => - Response.json(endpointWith({ host: "ep-expected.database.example.com" })), - ); - vi.spyOn(console, "error").mockImplementation(() => undefined); - - await expect( - assertEndpointHostMatches({ endpoint, host, workspaceClient: client }), - ).rejects.toThrow(host); - expect(request).toHaveBeenCalledWith( - expect.objectContaining({ - method: "GET", - path: `/api/2.0/postgres/${endpoint}`, - signal: expect.any(AbortSignal), - }), - ); - }); - - test("warns when request() reports the endpoint no longer exists (404)", async () => { - const { endpoint, host } = names(); - const { client } = modularClient(async () => { - throw Object.assign(new Error("not found"), { statusCode: 404 }); - }); - - await assertEndpointHostMatches({ - endpoint, - host, - workspaceClient: client, - }); - expect(warnings()).toContain("was not found"); - }); - - test("aborts the request() lookup when the deadline expires", async () => { - vi.useFakeTimers(); - const { endpoint, host } = names(); - let seen: AbortSignal | undefined; - const { client } = modularClient( - (req) => - new Promise(() => { - seen = req.signal; - }), - ); - - const check = assertEndpointHostMatches({ - endpoint, - host, - workspaceClient: client, - }); - await vi.advanceTimersByTimeAsync(3_000); - await check; - expect(seen?.aborted).toBe(true); - expect(warnings()).toContain("timed out"); +// A request-only client (no typed postgres, no legacy apiClient) can't read the +// endpoint, so the check is skipped with a warning instead of blocking startup. +test("skips verification for a request-only client", async () => { + const { endpoint, host } = names(); + const request = vi.fn(); + await assertEndpointHostMatches({ + endpoint, + host, + workspaceClient: { request } as unknown as Client, }); + expect(request).not.toHaveBeenCalled(); + expect(warnings()).toContain("Could not verify"); }); // AppKit's workspace client carries the typed Lakebase Postgres client, which diff --git a/packages/lakebase/src/__tests__/credentials.test.ts b/packages/lakebase/src/__tests__/credentials.test.ts index a723f48f1..3a9d61d2e 100644 --- a/packages/lakebase/src/__tests__/credentials.test.ts +++ b/packages/lakebase/src/__tests__/credentials.test.ts @@ -170,6 +170,94 @@ describe("Lakebase Authentication", () => { expect(mockApiClient.request).not.toHaveBeenCalled(); }); + it("uses the typed postgres client when present, mapping to camelCase", async () => { + const request = vi.fn(); + const generate = vi.fn(async () => ({ + token: "typed-token", + // Stand-in for Temporal.Instant: only toString() is relied on. + expireTime: { toString: () => "2026-02-06T18:00:00Z" }, + })); + const client = { + request, + postgres: { + generateDatabaseCredential: generate, + getEndpoint: vi.fn(), + }, + }; + + const result = await generateDatabaseCredential(client, { + endpoint: "projects/p/branches/main/endpoints/primary", + claims: [ + { + permission_set: RequestedClaimsPermissionSet.READ_ONLY, + resources: [{ table_name: "catalog.schema.users" }], + }, + ], + }); + + expect(result).toEqual({ + token: "typed-token", + expire_time: "2026-02-06T18:00:00Z", + }); + expect(generate).toHaveBeenCalledWith({ + endpoint: "projects/p/branches/main/endpoints/primary", + claims: [ + { + permissionSet: "READ_ONLY", + resources: [ + { + resourceName: { + $case: "tableName", + tableName: "catalog.schema.users", + }, + }, + ], + }, + ], + }); + expect(request).not.toHaveBeenCalled(); + }); + + it("rejects a typed response missing expireTime", async () => { + const client = { + request: vi.fn(), + postgres: { + generateDatabaseCredential: vi.fn(async () => ({ token: "t" })), + getEndpoint: vi.fn(), + }, + }; + await expect( + generateDatabaseCredential(client, { endpoint: "e" }), + ).rejects.toThrow(); + }); + + it("falls back to raw request() for unspecified_resource_name claims", async () => { + // The typed model only has tableName, so these would be silently dropped. + const request = vi.fn(async () => Response.json(credential)); + const generate = vi.fn(); + const claims = [ + { resources: [{ unspecified_resource_name: "some-resource" }] }, + ]; + + await generateDatabaseCredential( + { + request, + postgres: { + generateDatabaseCredential: generate, + getEndpoint: vi.fn(), + }, + }, + { endpoint: "e", claims }, + ); + + expect(generate).not.toHaveBeenCalled(); + expect(request).toHaveBeenCalledWith( + expect.objectContaining({ + body: JSON.stringify({ endpoint: "e", claims }), + }), + ); + }); + it("rejects a malformed credential response", async () => { const request = vi.fn(async () => Response.json({ token: "t" })); await expect( diff --git a/packages/lakebase/src/credentials.ts b/packages/lakebase/src/credentials.ts index eca506444..940cbc24f 100644 --- a/packages/lakebase/src/credentials.ts +++ b/packages/lakebase/src/credentials.ts @@ -53,6 +53,34 @@ export async function generateDatabaseCredential( ): Promise { const apiPath = "/api/2.0/postgres/credentials"; + if ( + isRequestCapable(workspaceClient) && + workspaceClient.postgres && + !hasUnspecifiedResource(request) + ) { + const credential = + await workspaceClient.postgres.generateDatabaseCredential({ + endpoint: request.endpoint, + claims: request.claims?.map((claim) => ({ + permissionSet: claim.permission_set, + resources: claim.resources?.map((resource) => ({ + resourceName: + resource.table_name === undefined + ? undefined + : { + $case: "tableName" as const, + tableName: resource.table_name, + }, + })), + })), + }); + return validateCredentialResponse({ + token: credential.token, + // Temporal.Instant -> ISO-8601, the shape the raw API returns. + expire_time: credential.expireTime?.toString(), + }); + } + if (isRequestCapable(workspaceClient)) { const response = await workspaceClient.request({ method: "POST", @@ -80,6 +108,22 @@ export async function generateDatabaseCredential( return validateCredentialResponse(response); } +/** + * The typed model has no `unspecified_resource_name` (only `tableName`), so those + * claims go through the raw request to avoid silently dropping them. + */ +function hasUnspecifiedResource( + request: GenerateDatabaseCredentialRequest, +): boolean { + return ( + request.claims?.some((claim) => + claim.resources?.some( + (resource) => resource.unspecified_resource_name !== undefined, + ), + ) ?? false + ); +} + /** Validate the API response has the expected shape */ function validateCredentialResponse(response: unknown): DatabaseCredential { if ( diff --git a/packages/lakebase/src/index.ts b/packages/lakebase/src/index.ts index 8e00ad4c1..1fd760bf6 100644 --- a/packages/lakebase/src/index.ts +++ b/packages/lakebase/src/index.ts @@ -12,6 +12,7 @@ export type { LakebaseWorkspaceRequest, LegacyWorkspaceClientLike, LakebasePoolConfig, + LakebasePostgresClient, Logger, LoggerConfig, RequestedClaims, diff --git a/packages/lakebase/src/types.ts b/packages/lakebase/src/types.ts index 041b02dfb..9a62cf8dd 100644 --- a/packages/lakebase/src/types.ts +++ b/packages/lakebase/src/types.ts @@ -224,6 +224,31 @@ export interface LakebaseWorkspaceRequest { */ export interface RequestCapableWorkspaceClient { request(req: LakebaseWorkspaceRequest): Promise; + /** Typed Lakebase Postgres client; used instead of raw requests when present. */ + postgres?: LakebasePostgresClient; +} + +/** + * The subset of the modular `@databricks/sdk-postgres` `PostgresClient` lakebase + * uses (the AppKit workspace client's `postgres`). Its model is camelCase. + */ +export interface LakebasePostgresClient { + generateDatabaseCredential( + req: { + endpoint?: string; + claims?: { + permissionSet?: string; + resources?: { + resourceName?: { $case: "tableName"; tableName: string }; + }[]; + }[]; + }, + options?: { signal?: AbortSignal }, + ): Promise<{ token?: string; expireTime?: { toString(): string } }>; + getEndpoint( + req: { name: string }, + options?: { signal?: AbortSignal }, + ): Promise; } /** The subset of the legacy `@databricks/sdk-experimental` `WorkspaceClient` lakebase uses. */ From 162c59ff482abb4b644f886ba018fbb8e3d80aaa Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 14 Oct 2026 16:03:31 +0200 Subject: [PATCH 5/5] fix(lakebase): send Headers through request-capable workspace clients AppKit's workspace client types its raw request headers as Headers, so the request type and its callers now use Headers too. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/lakebase/src/__tests__/credentials.test.ts | 12 ++++++++---- packages/lakebase/src/config.ts | 2 +- packages/lakebase/src/credentials.ts | 4 ++-- packages/lakebase/src/types.ts | 2 +- 4 files changed, 12 insertions(+), 8 deletions(-) diff --git a/packages/lakebase/src/__tests__/credentials.test.ts b/packages/lakebase/src/__tests__/credentials.test.ts index 3a9d61d2e..0d03bdbcc 100644 --- a/packages/lakebase/src/__tests__/credentials.test.ts +++ b/packages/lakebase/src/__tests__/credentials.test.ts @@ -148,15 +148,19 @@ describe("Lakebase Authentication", () => { expect(request).toHaveBeenCalledWith({ method: "POST", path: "/api/2.0/postgres/credentials", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - }, + headers: expect.any(Headers), body: JSON.stringify({ endpoint: "projects/p/branches/main/endpoints/primary", claims, }), }); + const [[sent]] = request.mock.calls as unknown as [ + [{ headers: Headers }], + ]; + expect(Object.fromEntries(sent.headers)).toEqual({ + accept: "application/json", + "content-type": "application/json", + }); }); it("prefers request() over a legacy apiClient on the same object", async () => { diff --git a/packages/lakebase/src/config.ts b/packages/lakebase/src/config.ts index 8a1de2c88..a8e779c3b 100644 --- a/packages/lakebase/src/config.ts +++ b/packages/lakebase/src/config.ts @@ -245,7 +245,7 @@ export async function getUsernameWithApiLookup( const response = await workspaceClient.request({ method: "GET", path: "/api/2.0/preview/scim/v2/Me", - headers: { Accept: "application/json" }, + headers: new Headers({ Accept: "application/json" }), }); const me = (await response.json()) as { userName?: string }; return me.userName ?? undefined; diff --git a/packages/lakebase/src/credentials.ts b/packages/lakebase/src/credentials.ts index 940cbc24f..b958eb89f 100644 --- a/packages/lakebase/src/credentials.ts +++ b/packages/lakebase/src/credentials.ts @@ -85,10 +85,10 @@ export async function generateDatabaseCredential( const response = await workspaceClient.request({ method: "POST", path: apiPath, - headers: { + headers: new Headers({ Accept: "application/json", "Content-Type": "application/json", - }, + }), body: JSON.stringify(request), }); return validateCredentialResponse(await response.json()); diff --git a/packages/lakebase/src/types.ts b/packages/lakebase/src/types.ts index 9a62cf8dd..d123e588c 100644 --- a/packages/lakebase/src/types.ts +++ b/packages/lakebase/src/types.ts @@ -214,7 +214,7 @@ export interface LakebaseWorkspaceRequest { method: string; /** Path on the workspace host, e.g. `/api/2.0/postgres/credentials`. */ path: string; - headers?: Record; + headers?: Headers; body?: string; }