From d330655d70bb1e87afe48aced097a4958bca1476 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 7 Oct 2026 18:49:28 +0200 Subject: [PATCH 1/2] 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 c72820060..bf5de9260 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -606,9 +606,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 5fc3b680a0059dd999f757993c810996a49b70e1 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 7 Oct 2026 18:49:45 +0200 Subject: [PATCH 2/2] 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; }; }