From 01e05492c46c6de8f7e6cccada7b573933cf7967 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 10 Jul 2026 15:19:40 +0200 Subject: [PATCH] feat(appkit): introduce workspace-client wrapper over the Databricks SDK MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add packages/appkit/src/workspace-client/ — a facade over the legacy @databricks/sdk-experimental client that mirrors the modular Databricks SDK's multi-client shape. This prepares the library for an incremental, per-service migration to the modular SDK without migrating anything yet. - legacy.ts is the ONLY module importing @databricks/sdk-experimental; it constructs the client and re-exports the SDK symbols AppKit uses. - The WorkspaceClient facade exposes per-service accessors (files, warehouses, genie, jobs, statementExecution, servingEndpoints, currentUser, config, apiClient), each legacy-typed and delegating to the underlying client. Migrating a service later is a localized swap: one getter + its accessor type + one connector. - createWorkspaceClient() replaces every `new WorkspaceClient(...)` site (runtime + build-time type-generator); toLegacyWorkspaceClient() is the escape hatch for @databricks/lakebase. - createApp({ client }) and the public index now expose the wrapper type instead of the raw SDK client. - A Biome noRestrictedImports boundary rule forbids importing @databricks/sdk-experimental outside workspace-client/, keeping the seam enforceable. Tests mock the wrapper rather than the SDK. Time is sourced off the SDK namespace (SDK.Time ?? SDK.default.Time) because the SDK's CommonJS `Time` export is a getter that defeats Node's static ESM named-export detection — a direct `export { Time }` throws at link time. Signed-off-by: MarioCadenas --- biome.json | 30 ++++- .../api/appkit/Class.DatabricksAdapter.md | 5 +- docs/docs/api/appkit/Function.createApp.md | 4 +- .../appkit/Function.createWorkspaceClient.md | 25 ++++ .../api/appkit/Interface.WorkspaceClient.md | 118 ++++++++++++++++++ .../Interface.WorkspaceClientOptions.md | 47 +++++++ docs/docs/api/appkit/index.md | 3 + docs/docs/api/appkit/typedoc-sidebar.ts | 15 +++ packages/appkit/src/agents/databricks.ts | 14 +-- .../src/agents/tests/databricks.test.ts | 15 ++- packages/appkit/src/cache/index.ts | 6 +- .../src/cache/tests/cache-manager.test.ts | 12 +- .../appkit/src/connectors/files/client.ts | 2 +- .../src/connectors/files/tests/client.test.ts | 15 ++- .../appkit/src/connectors/genie/client.ts | 14 +-- .../src/connectors/genie/tests/client.test.ts | 2 +- packages/appkit/src/connectors/jobs/client.ts | 10 +- .../appkit/src/connectors/serving/client.ts | 12 +- .../connectors/serving/tests/client.test.ts | 2 +- .../src/connectors/sql-warehouse/client.ts | 10 +- .../src/connectors/sql-warehouse/defaults.ts | 2 +- .../sql-warehouse/warehouse-status-emitter.ts | 2 +- .../src/connectors/vector-search/client.ts | 2 +- packages/appkit/src/context/client-options.ts | 2 +- .../appkit/src/context/service-context.ts | 30 ++--- .../src/context/tests/service-context.test.ts | 15 ++- .../tests/user-agent-wire.integration.test.ts | 12 +- packages/appkit/src/core/appkit.ts | 2 +- .../core/tests/appkit-as-user-exports.test.ts | 13 +- packages/appkit/src/index.ts | 8 ++ .../appkit/src/internal-telemetry/reporter.ts | 2 +- .../internal-telemetry/tests/reporter.test.ts | 2 +- .../src/plugin/tests/asUser-proxy.test.ts | 15 ++- .../appkit/src/plugin/tests/plugin.test.ts | 11 +- .../appkit/src/plugins/analytics/analytics.ts | 2 +- .../appkit/src/plugins/analytics/query.ts | 2 +- packages/appkit/src/plugins/files/plugin.ts | 2 +- .../src/plugins/files/tests/delete.test.ts | 13 +- .../files/tests/download-endpoint.test.ts | 13 +- .../files/tests/error-handling.test.ts | 13 +- .../src/plugins/files/tests/mkdir.test.ts | 13 +- .../files/tests/path-validation.test.ts | 13 +- .../files/tests/plugin.integration.test.ts | 4 +- .../src/plugins/files/tests/plugin.test.ts | 13 +- .../plugins/files/tests/raw-endpoint.test.ts | 13 +- .../src/plugins/files/tests/shutdown.test.ts | 13 +- .../src/plugins/files/tests/upload.test.ts | 13 +- .../plugins/files/tests/volume-config.test.ts | 13 +- packages/appkit/src/plugins/files/types.ts | 2 +- packages/appkit/src/plugins/jobs/plugin.ts | 2 +- .../src/plugins/jobs/tests/plugin.test.ts | 15 ++- packages/appkit/src/plugins/jobs/types.ts | 2 +- .../appkit/src/plugins/lakebase/lakebase.ts | 17 ++- .../tests/lakebase-agent-tool.test.ts | 6 +- .../src/stream/arrow-stream-processor.ts | 2 +- .../tests/arrow-stream-processor.test.ts | 2 +- packages/appkit/src/type-generator/index.ts | 7 +- .../type-generator/mv-registry/describe.ts | 2 +- .../src/type-generator/query-registry.ts | 4 +- .../src/type-generator/serving/fetcher.ts | 2 +- .../src/type-generator/serving/generator.ts | 7 +- .../serving/tests/fetcher.test.ts | 2 +- .../serving/tests/generator.test.ts | 13 +- .../src/type-generator/statement-result.ts | 2 +- .../tests/generate-queries.test.ts | 17 ++- .../src/type-generator/tests/index.test.ts | 35 +++--- .../tests/statement-result.test.ts | 2 +- .../type-generator/tests/vite-plugin.test.ts | 15 ++- .../tests/warehouse-status.test.ts | 2 +- .../appkit/src/type-generator/vite-plugin.ts | 4 +- .../src/type-generator/warehouse-status.ts | 2 +- .../appkit/src/workspace-client/client.ts | 71 +++++++++++ .../appkit/src/workspace-client/errors.ts | 14 +++ .../appkit/src/workspace-client/factory.ts | 23 ++++ packages/appkit/src/workspace-client/index.ts | 29 +++++ .../appkit/src/workspace-client/legacy.ts | 88 +++++++++++++ packages/appkit/src/workspace-client/types.ts | 77 ++++++++++++ 77 files changed, 877 insertions(+), 208 deletions(-) create mode 100644 docs/docs/api/appkit/Function.createWorkspaceClient.md create mode 100644 docs/docs/api/appkit/Interface.WorkspaceClient.md create mode 100644 docs/docs/api/appkit/Interface.WorkspaceClientOptions.md create mode 100644 packages/appkit/src/workspace-client/client.ts create mode 100644 packages/appkit/src/workspace-client/errors.ts create mode 100644 packages/appkit/src/workspace-client/factory.ts create mode 100644 packages/appkit/src/workspace-client/index.ts create mode 100644 packages/appkit/src/workspace-client/legacy.ts create mode 100644 packages/appkit/src/workspace-client/types.ts diff --git a/biome.json b/biome.json index c9b8a4a1..1edd6278 100644 --- a/biome.json +++ b/biome.json @@ -41,6 +41,22 @@ }, "security": { "noDangerouslySetInnerHtml": "off" + }, + "style": { + "noRestrictedImports": { + "level": "error", + "options": { + "patterns": [ + { + "group": [ + "@databricks/sdk-experimental", + "@databricks/sdk-experimental/**" + ], + "message": "Import the Databricks SDK only through the wrapper in packages/appkit/src/workspace-client. Add a re-export there if you need a new symbol." + } + ] + } + } } } }, @@ -56,5 +72,17 @@ "organizeImports": "on" } } - } + }, + "overrides": [ + { + "includes": ["packages/appkit/src/workspace-client/**"], + "linter": { + "rules": { + "style": { + "noRestrictedImports": "off" + } + } + } + } + ] } diff --git a/docs/docs/api/appkit/Class.DatabricksAdapter.md b/docs/docs/api/appkit/Class.DatabricksAdapter.md index ba4a8a18..4833629a 100644 --- a/docs/docs/api/appkit/Class.DatabricksAdapter.md +++ b/docs/docs/api/appkit/Class.DatabricksAdapter.md @@ -12,12 +12,11 @@ fallback parsing for models that output tool calls as text. ## Examples ```ts -import { createApp, createAgent, agents } from "@databricks/appkit"; +import { createApp, createAgent, agents, createWorkspaceClient } from "@databricks/appkit"; import { DatabricksAdapter } from "@databricks/appkit/beta"; -import { WorkspaceClient } from "@databricks/sdk-experimental"; const adapter = DatabricksAdapter.fromServingEndpoint({ - workspaceClient: new WorkspaceClient({}), + workspaceClient: createWorkspaceClient(), endpointName: "my-endpoint", }); diff --git a/docs/docs/api/appkit/Function.createApp.md b/docs/docs/api/appkit/Function.createApp.md index 85cb584b..bc656537 100644 --- a/docs/docs/api/appkit/Function.createApp.md +++ b/docs/docs/api/appkit/Function.createApp.md @@ -31,9 +31,9 @@ with an `asUser(req)` method for user-scoped execution. | Parameter | Type | | ------ | ------ | -| `config` | \{ `cache?`: [`CacheConfig`](Interface.CacheConfig.md); `client?`: `WorkspaceClient`; `disableInternalTelemetry?`: `boolean`; `onPluginsReady?`: (`appkit`: `PluginMap`\<`T`\>) => `void` \| `Promise`\<`void`\>; `plugins?`: `T`; `telemetry?`: [`TelemetryConfig`](Interface.TelemetryConfig.md); \} | +| `config` | \{ `cache?`: [`CacheConfig`](Interface.CacheConfig.md); `client?`: [`WorkspaceClient`](Interface.WorkspaceClient.md); `disableInternalTelemetry?`: `boolean`; `onPluginsReady?`: (`appkit`: `PluginMap`\<`T`\>) => `void` \| `Promise`\<`void`\>; `plugins?`: `T`; `telemetry?`: [`TelemetryConfig`](Interface.TelemetryConfig.md); \} | | `config.cache?` | [`CacheConfig`](Interface.CacheConfig.md) | -| `config.client?` | `WorkspaceClient` | +| `config.client?` | [`WorkspaceClient`](Interface.WorkspaceClient.md) | | `config.disableInternalTelemetry?` | `boolean` | | `config.onPluginsReady?` | (`appkit`: `PluginMap`\<`T`\>) => `void` \| `Promise`\<`void`\> | | `config.plugins?` | `T` | diff --git a/docs/docs/api/appkit/Function.createWorkspaceClient.md b/docs/docs/api/appkit/Function.createWorkspaceClient.md new file mode 100644 index 00000000..8ad87f41 --- /dev/null +++ b/docs/docs/api/appkit/Function.createWorkspaceClient.md @@ -0,0 +1,25 @@ +# Function: createWorkspaceClient() + +```ts +function createWorkspaceClient(opts: WorkspaceClientOptions): WorkspaceClient; +``` + +Construct an AppKit workspace client. + +Auth resolution: + - If `opts.token` is set, uses PAT credentials. + - Otherwise walks the SDK default auth chain (env vars + ~/.databrickscfg). + +Host resolution: + - Explicit `opts.host` → use it. + - Otherwise resolved by the SDK from `DATABRICKS_HOST` / profile. + +## Parameters + +| Parameter | Type | +| ------ | ------ | +| `opts` | [`WorkspaceClientOptions`](Interface.WorkspaceClientOptions.md) | + +## Returns + +[`WorkspaceClient`](Interface.WorkspaceClient.md) diff --git a/docs/docs/api/appkit/Interface.WorkspaceClient.md b/docs/docs/api/appkit/Interface.WorkspaceClient.md new file mode 100644 index 00000000..bf508bbe --- /dev/null +++ b/docs/docs/api/appkit/Interface.WorkspaceClient.md @@ -0,0 +1,118 @@ +# Interface: WorkspaceClient + +AppKit's workspace client facade. Mirrors the multi-client shape of the +modular Databricks SDK: each service is its own accessor, so services can be +migrated one at a time behind this stable interface. + +Accessors are legacy-typed for now (delegated to the underlying legacy SDK +client); see the module docblock. + +## Properties + +### apiClient + +```ts +readonly apiClient: ApiClient; +``` + +Low-level HTTP transport (`apiClient.request(...)`). Used for endpoints +without a typed service method: SCIM header probe, warehouse listing, +serving SSE streaming, vector search, internal telemetry. + +*** + +### config + +```ts +readonly config: Config; +``` + +SDK `Config` — exposes `host` and `authenticate(headers)`. Used by the +files-upload path and agents auth-header stamping, which bypass the typed +services. + +*** + +### currentUser + +```ts +readonly currentUser: CurrentUserService; +``` + +Current user. + +*** + +### files + +```ts +readonly files: FilesService; +``` + +UC Volumes / Files API. + +*** + +### genie + +```ts +readonly genie: GenieService; +``` + +Genie / dashboards. + +*** + +### jobs + +```ts +readonly jobs: JobsService; +``` + +Jobs. + +*** + +### servingEndpoints + +```ts +readonly servingEndpoints: ServingEndpointsService; +``` + +Serving Endpoints. + +*** + +### statementExecution + +```ts +readonly statementExecution: StatementExecutionService; +``` + +Statement Execution. + +*** + +### warehouses + +```ts +readonly warehouses: WarehousesService; +``` + +SQL Warehouses. + +## Methods + +### toLegacyWorkspaceClient() + +```ts +toLegacyWorkspaceClient(): WorkspaceClient; +``` + +Returns the underlying legacy `@databricks/sdk-experimental` +`WorkspaceClient`, for handoff to code still typed against the old SDK +(`@databricks/lakebase`). Transitional. + +#### Returns + +`WorkspaceClient` diff --git a/docs/docs/api/appkit/Interface.WorkspaceClientOptions.md b/docs/docs/api/appkit/Interface.WorkspaceClientOptions.md new file mode 100644 index 00000000..29ef96aa --- /dev/null +++ b/docs/docs/api/appkit/Interface.WorkspaceClientOptions.md @@ -0,0 +1,47 @@ +# Interface: WorkspaceClientOptions + +Options used to construct the wrapper. Mirrors the subset of the old SDK's +`Config` + `ClientOptions` that AppKit relies on today; we deliberately do +NOT re-expose every old-SDK config knob. + +## Properties + +### authType? + +```ts +optional authType: "pat"; +``` + +Authentication strategy passed to the legacy client. + +*** + +### clientOptions? + +```ts +optional clientOptions: ClientOptions; +``` + +SDK client options (product / productVersion / userAgentExtra) used to +stamp the outbound User-Agent. Produced by `getClientOptions()`; omitted +by build-time callers that don't stamp a User-Agent. + +*** + +### host? + +```ts +optional host: string; +``` + +Databricks host, e.g. https://my-workspace.cloud.databricks.com. Defaults to DATABRICKS_HOST / profile resolution. + +*** + +### token? + +```ts +optional token: string; +``` + +Bearer token. When set, `authType` defaults to "pat". diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index 4d995051..969e3162 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -81,6 +81,8 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [ToolkitOptions](Interface.ToolkitOptions.md) | - | | [ToolProvider](Interface.ToolProvider.md) | - | | [ValidationResult](Interface.ValidationResult.md) | Result of validating all registered resources against the environment. | +| [WorkspaceClient](Interface.WorkspaceClient.md) | AppKit's workspace client facade. Mirrors the multi-client shape of the modular Databricks SDK: each service is its own accessor, so services can be migrated one at a time behind this stable interface. | +| [WorkspaceClientOptions](Interface.WorkspaceClientOptions.md) | Options used to construct the wrapper. Mirrors the subset of the old SDK's `Config` + `ClientOptions` that AppKit relies on today; we deliberately do NOT re-expose every old-SDK config knob. | ## Type Aliases @@ -128,6 +130,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [createApp](Function.createApp.md) | Bootstraps AppKit with the provided configuration. | | [createLakebasePool](Function.createLakebasePool.md) | Create a Lakebase pool with appkit's logger integration. Telemetry automatically uses appkit's OpenTelemetry configuration via global registry. | | [createLakebasePoolManager](Function.createLakebasePoolManager.md) | Create a pool manager that maintains per-key Lakebase connection pools. | +| [createWorkspaceClient](Function.createWorkspaceClient.md) | Construct an AppKit workspace client. | | [defineTool](Function.defineTool.md) | Defines a single tool entry for a plugin's internal registry. | | [executeFromRegistry](Function.executeFromRegistry.md) | Validates tool-call arguments against the entry's schema and invokes its handler. On validation failure, returns an LLM-friendly error string (matching the behavior of `tool()`) rather than throwing, so the model can self-correct on its next turn. | | [extractServingEndpoints](Function.extractServingEndpoints.md) | Extract serving endpoint config from a server file by AST-parsing it. Looks for `serving({ endpoints: { alias: { env: "..." }, ... } })` calls and extracts the endpoint alias names and their environment variable mappings. | diff --git a/docs/docs/api/appkit/typedoc-sidebar.ts b/docs/docs/api/appkit/typedoc-sidebar.ts index c74d6232..c488f71c 100644 --- a/docs/docs/api/appkit/typedoc-sidebar.ts +++ b/docs/docs/api/appkit/typedoc-sidebar.ts @@ -336,6 +336,16 @@ const typedocSidebar: SidebarsConfig = { type: "doc", id: "api/appkit/Interface.ValidationResult", label: "ValidationResult" + }, + { + type: "doc", + id: "api/appkit/Interface.WorkspaceClient", + label: "WorkspaceClient" + }, + { + type: "doc", + id: "api/appkit/Interface.WorkspaceClientOptions", + label: "WorkspaceClientOptions" } ] }, @@ -515,6 +525,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Function.createLakebasePoolManager", label: "createLakebasePoolManager" }, + { + type: "doc", + id: "api/appkit/Function.createWorkspaceClient", + label: "createWorkspaceClient" + }, { type: "doc", id: "api/appkit/Function.defineTool", diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 83f89a59..41b89104 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -7,6 +7,7 @@ import type { } from "shared"; import { stream as servingStream } from "../connectors/serving/client"; import { APPKIT_USER_AGENT, getClientOptions } from "../context/client-options"; +import { createWorkspaceClient } from "../workspace-client"; /** Default cap for a single incomplete SSE line tail (DoS guard). */ const DEFAULT_MAX_SSE_LINE_CHARS = 1024 * 1024; @@ -203,12 +204,11 @@ interface DeltaToolCall { * * @example Using the factory (recommended) * ```ts - * import { createApp, createAgent, agents } from "@databricks/appkit"; + * import { createApp, createAgent, agents, createWorkspaceClient } from "@databricks/appkit"; * import { DatabricksAdapter } from "@databricks/appkit/beta"; - * import { WorkspaceClient } from "@databricks/sdk-experimental"; * * const adapter = DatabricksAdapter.fromServingEndpoint({ - * workspaceClient: new WorkspaceClient({}), + * workspaceClient: createWorkspaceClient(), * endpointName: "my-endpoint", * }); * @@ -358,11 +358,9 @@ export class DatabricksAdapter implements AgentAdapter { let workspaceClient: WorkspaceClientLike | undefined = options?.workspaceClient; if (!workspaceClient) { - const sdk = await import("@databricks/sdk-experimental"); - workspaceClient = new sdk.WorkspaceClient( - {}, - getClientOptions(), - ) as unknown as WorkspaceClientLike; + workspaceClient = createWorkspaceClient({ + clientOptions: getClientOptions(), + }) as unknown as WorkspaceClientLike; } return DatabricksAdapter.fromServingEndpoint({ diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index 84f0c671..f680fd46 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -944,11 +944,16 @@ describe("DatabricksAdapter.fromModelServing", () => { test("reads endpoint from DATABRICKS_SERVING_ENDPOINT_NAME env var", async () => { process.env.DATABRICKS_SERVING_ENDPOINT_NAME = "my-model"; - vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn().mockImplementation(() => ({ - apiClient: { request: vi.fn() }, - })), - })); + vi.mock("../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: vi.fn().mockImplementation(() => ({ + apiClient: { request: vi.fn() }, + })), + }; + }); const adapter = await DatabricksAdapter.fromModelServing(); expect(adapter).toBeInstanceOf(DatabricksAdapter); diff --git a/packages/appkit/src/cache/index.ts b/packages/appkit/src/cache/index.ts index 641b59fa..e2b79d04 100644 --- a/packages/appkit/src/cache/index.ts +++ b/packages/appkit/src/cache/index.ts @@ -1,5 +1,4 @@ import { createHash } from "node:crypto"; -import { ApiError, WorkspaceClient } from "@databricks/sdk-experimental"; import type { CacheConfig, CacheEntry, CacheStorage } from "shared"; import { createLakebasePool } from "../connectors/lakebase"; import { getClientOptions } from "../context/client-options"; @@ -8,6 +7,7 @@ import { createLogger } from "../logging/logger"; import type { Counter, TelemetryProvider } from "../telemetry"; import { SpanStatusCode, TelemetryManager } from "../telemetry"; import { deepMerge } from "../utils"; +import { ApiError, createWorkspaceClient } from "../workspace-client"; import { cacheDefaults } from "./defaults"; import { InMemoryStorage, PersistentStorage } from "./storage"; @@ -171,7 +171,9 @@ export class CacheManager { // try to use lakebase storage try { - const workspaceClient = new WorkspaceClient({}, getClientOptions()); + const workspaceClient = createWorkspaceClient({ + clientOptions: getClientOptions(), + }).toLegacyWorkspaceClient(); const pool = createLakebasePool({ workspaceClient }); const persistentStorage = new PersistentStorage(config, pool); diff --git a/packages/appkit/src/cache/tests/cache-manager.test.ts b/packages/appkit/src/cache/tests/cache-manager.test.ts index 3f8fa5db..9f668988 100644 --- a/packages/appkit/src/cache/tests/cache-manager.test.ts +++ b/packages/appkit/src/cache/tests/cache-manager.test.ts @@ -48,13 +48,15 @@ vi.mock("../storage/persistent", () => ({ }), })); -// Mock WorkspaceClient (preserve Time, TimeUnits, and other SDK exports) -vi.mock("@databricks/sdk-experimental", async (importOriginal) => { +// Mock createWorkspaceClient (preserve Time, TimeUnits, and other wrapper exports) +vi.mock("../../workspace-client", async (importOriginal) => { const actual = - await importOriginal(); + await importOriginal(); return { ...actual, - WorkspaceClient: vi.fn().mockImplementation(() => ({})), + createWorkspaceClient: vi.fn().mockImplementation(() => ({ + toLegacyWorkspaceClient: () => ({}), + })), }; }); @@ -976,7 +978,7 @@ describe("CacheManager", () => { }); test("should re-throw ApiError without wrapping", async () => { - const { ApiError } = await import("@databricks/sdk-experimental"); + const { ApiError } = await import("../../workspace-client"); const cache = await CacheManager.getInstance({ storage: createMockStorage(), }); diff --git a/packages/appkit/src/connectors/files/client.ts b/packages/appkit/src/connectors/files/client.ts index ddb7ecde..1d160ee0 100644 --- a/packages/appkit/src/connectors/files/client.ts +++ b/packages/appkit/src/connectors/files/client.ts @@ -1,5 +1,4 @@ import { AsyncLocalStorage } from "node:async_hooks"; -import { ApiError, type WorkspaceClient } from "@databricks/sdk-experimental"; import type { TelemetryOptions } from "shared"; import { createLogger } from "../../logging/logger"; import type { @@ -17,6 +16,7 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; +import { ApiError, type WorkspaceClient } from "../../workspace-client"; import { contentTypeFromPath, FILES_MAX_READ_SIZE, diff --git a/packages/appkit/src/connectors/files/tests/client.test.ts b/packages/appkit/src/connectors/files/tests/client.test.ts index 56076470..4d2b9d07 100644 --- a/packages/appkit/src/connectors/files/tests/client.test.ts +++ b/packages/appkit/src/connectors/files/tests/client.test.ts @@ -1,6 +1,6 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import { createMockTelemetry } from "@tools/test-helpers"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import type { WorkspaceClient } from "../../../workspace-client"; import { FilesConnector } from "../client"; import { streamFromChunks, streamFromString } from "./utils"; @@ -50,10 +50,15 @@ const { mockFilesApi, mockConfig, mockClient, MockApiError } = vi.hoisted( }, ); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: () => mockClient, + ApiError: MockApiError, + }; +}); const mockTelemetry = createMockTelemetry(); diff --git a/packages/appkit/src/connectors/genie/client.ts b/packages/appkit/src/connectors/genie/client.ts index 5be91b6e..ec8a987d 100644 --- a/packages/appkit/src/connectors/genie/client.ts +++ b/packages/appkit/src/connectors/genie/client.ts @@ -1,8 +1,11 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; -import * as SDK from "@databricks/sdk-experimental"; -import type { GenieMessage } from "@databricks/sdk-experimental/dist/apis/dashboards"; -import type { Waiter } from "@databricks/sdk-experimental/dist/wait"; import { createLogger } from "../../logging"; +import { + type GenieMessage, + Time, + TimeUnits, + type Waiter, + type WorkspaceClient, +} from "../../workspace-client"; import { genieConnectorDefaults } from "./defaults"; import { pollWaiter } from "./poll-waiter"; import type { @@ -13,9 +16,6 @@ import type { GenieStreamEvent, } from "./types"; -const { TimeUnits } = SDK; -const Time = SDK.Time ?? (SDK as any).default.Time; - const logger = createLogger("connectors:genie"); const GenieErrors = { diff --git a/packages/appkit/src/connectors/genie/tests/client.test.ts b/packages/appkit/src/connectors/genie/tests/client.test.ts index 62fc3578..c2f0dbba 100644 --- a/packages/appkit/src/connectors/genie/tests/client.test.ts +++ b/packages/appkit/src/connectors/genie/tests/client.test.ts @@ -1,5 +1,5 @@ -import type { GenieMessage } from "@databricks/sdk-experimental/dist/apis/dashboards"; import { beforeEach, describe, expect, test, vi } from "vitest"; +import type { GenieMessage } from "../../../workspace-client"; import { GenieConnector } from "../client"; import type { GenieStreamEvent } from "../types"; diff --git a/packages/appkit/src/connectors/jobs/client.ts b/packages/appkit/src/connectors/jobs/client.ts index a2914535..a00474d2 100644 --- a/packages/appkit/src/connectors/jobs/client.ts +++ b/packages/appkit/src/connectors/jobs/client.ts @@ -1,8 +1,3 @@ -import { - Context, - type jobs, - type WorkspaceClient, -} from "@databricks/sdk-experimental"; import { AppKitError, ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import type { TelemetryProvider } from "../../telemetry"; @@ -14,6 +9,11 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; +import { + Context, + type jobs, + type WorkspaceClient, +} from "../../workspace-client"; import type { JobsConnectorConfig } from "./types"; const logger = createLogger("connectors:jobs"); diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index 83f065e6..f75a258f 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -1,10 +1,10 @@ -import type { - CancellationToken, - serving, - WorkspaceClient, -} from "@databricks/sdk-experimental"; -import { Context } from "@databricks/sdk-experimental"; import { createLogger } from "../../logging/logger"; +import { + type CancellationToken, + Context, + type serving, + type WorkspaceClient, +} from "../../workspace-client"; const logger = createLogger("connectors:serving"); diff --git a/packages/appkit/src/connectors/serving/tests/client.test.ts b/packages/appkit/src/connectors/serving/tests/client.test.ts index d243621e..34bfd743 100644 --- a/packages/appkit/src/connectors/serving/tests/client.test.ts +++ b/packages/appkit/src/connectors/serving/tests/client.test.ts @@ -1,5 +1,5 @@ -import { Context } from "@databricks/sdk-experimental"; import { afterEach, describe, expect, test, vi } from "vitest"; +import { Context } from "../../../workspace-client"; import { invoke, stream } from "../client"; function createMockClient(host = "https://test.databricks.com") { diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 2f6a9f73..a722d85f 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -1,8 +1,3 @@ -import { - Context, - type sql, - type WorkspaceClient, -} from "@databricks/sdk-experimental"; import type { TelemetryOptions } from "shared"; import { AppKitError, @@ -22,6 +17,11 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; +import { + Context, + type sql, + type WorkspaceClient, +} from "../../workspace-client"; import { executeStatementDefaults } from "./defaults"; import { WarehousePollBackoff } from "./warehouse-poll-backoff"; import { WarehouseStatusEmitter } from "./warehouse-status-emitter"; diff --git a/packages/appkit/src/connectors/sql-warehouse/defaults.ts b/packages/appkit/src/connectors/sql-warehouse/defaults.ts index 994f11da..b046a5c4 100644 --- a/packages/appkit/src/connectors/sql-warehouse/defaults.ts +++ b/packages/appkit/src/connectors/sql-warehouse/defaults.ts @@ -1,4 +1,4 @@ -import type { sql } from "@databricks/sdk-experimental"; +import type { sql } from "../../workspace-client"; interface ExecuteStatementDefaults { wait_timeout: string; diff --git a/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts b/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts index 5f310c80..aba8488c 100644 --- a/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts +++ b/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts @@ -1,5 +1,5 @@ -import type { sql } from "@databricks/sdk-experimental"; import type { Span } from "../../telemetry"; +import type { sql } from "../../workspace-client"; import type { WarehouseStatusUpdate } from "./client"; /** diff --git a/packages/appkit/src/connectors/vector-search/client.ts b/packages/appkit/src/connectors/vector-search/client.ts index a20a59ed..f424b17e 100644 --- a/packages/appkit/src/connectors/vector-search/client.ts +++ b/packages/appkit/src/connectors/vector-search/client.ts @@ -1,4 +1,3 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import { createLogger } from "../../logging/logger"; import type { TelemetryProvider } from "../../telemetry"; import { @@ -7,6 +6,7 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; +import type { WorkspaceClient } from "../../workspace-client"; import type { VectorSearchConnectorConfig, VsNextPageParams, diff --git a/packages/appkit/src/context/client-options.ts b/packages/appkit/src/context/client-options.ts index 673984bb..35482899 100644 --- a/packages/appkit/src/context/client-options.ts +++ b/packages/appkit/src/context/client-options.ts @@ -1,9 +1,9 @@ -import type { ClientOptions } from "@databricks/sdk-experimental"; import { coerce } from "semver"; import { name as productName, version as productVersion, } from "../../package.json"; +import type { ClientOptions } from "../workspace-client"; const normalizedVersion = (coerce(productVersion)?.version ?? productVersion) as ClientOptions["productVersion"]; diff --git a/packages/appkit/src/context/service-context.ts b/packages/appkit/src/context/service-context.ts index 22c92f22..b5840bbc 100644 --- a/packages/appkit/src/context/service-context.ts +++ b/packages/appkit/src/context/service-context.ts @@ -1,15 +1,16 @@ import { createHash } from "node:crypto"; -import { - type ClientOptions, - ConfigError, - type sql, - WorkspaceClient, -} from "@databricks/sdk-experimental"; import { AuthenticationError, ConfigurationError, InitializationError, } from "../errors"; +import { + type ClientOptions, + ConfigError, + createWorkspaceClient, + type sql, + type WorkspaceClient, +} from "../workspace-client"; import { getClientOptions } from "./client-options"; import type { UserContext } from "./user-context"; @@ -113,14 +114,12 @@ export class ServiceContext { // Create user client with the OAuth token from Databricks Apps // Note: We use authType: "pat" because the token is passed as a Bearer token // just like a PAT, even though it's technically an OAuth token - const userClient = new WorkspaceClient( - { - token, - host, - authType: "pat", - }, - getClientOptions(), - ); + const userClient = createWorkspaceClient({ + token, + host, + authType: "pat", + clientOptions: getClientOptions(), + }); const tokenFingerprint = createHash("sha256") .update(token) @@ -152,7 +151,8 @@ export class ServiceContext { client?: WorkspaceClient, ): Promise { try { - const wsClient = client ?? new WorkspaceClient({}, getClientOptions()); + const wsClient = + client ?? createWorkspaceClient({ clientOptions: getClientOptions() }); const [resolvedWorkspaceId, currentUser, resolvedWarehouseId] = await Promise.all([ diff --git a/packages/appkit/src/context/tests/service-context.test.ts b/packages/appkit/src/context/tests/service-context.test.ts index 7f4dcdf3..26a53b99 100644 --- a/packages/appkit/src/context/tests/service-context.test.ts +++ b/packages/appkit/src/context/tests/service-context.test.ts @@ -7,7 +7,7 @@ import { } from "../../errors"; import { ServiceContext } from "../service-context"; -// ── Mock @databricks/sdk-experimental ────────────────────────────── +// ── Mock the workspace-client wrapper ────────────────────────────── const { mockMe, mockApiRequest, MockWorkspaceClient, MockConfigError } = vi.hoisted(() => { @@ -30,10 +30,15 @@ const { mockMe, mockApiRequest, MockWorkspaceClient, MockConfigError } = return { mockMe, mockApiRequest, MockWorkspaceClient, MockConfigError }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: MockWorkspaceClient, - ConfigError: MockConfigError, -})); +vi.mock("../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (...args: unknown[]) => MockWorkspaceClient(...args), + ConfigError: MockConfigError, + }; +}); // ── Helpers ──────────────────────────────────────────────────────── diff --git a/packages/appkit/src/context/tests/user-agent-wire.integration.test.ts b/packages/appkit/src/context/tests/user-agent-wire.integration.test.ts index 0c531531..063842c4 100644 --- a/packages/appkit/src/context/tests/user-agent-wire.integration.test.ts +++ b/packages/appkit/src/context/tests/user-agent-wire.integration.test.ts @@ -1,5 +1,4 @@ import http, { type Server } from "node:http"; -import { WorkspaceClient } from "@databricks/sdk-experimental"; import { afterAll, beforeAll, @@ -9,6 +8,7 @@ import { test, } from "vitest"; import { FilesConnector } from "../../connectors/files/client"; +import { createWorkspaceClient } from "../../workspace-client"; import { getClientOptions } from "../client-options"; /** @@ -24,10 +24,12 @@ describe("User-Agent attribution (wire)", () => { let capturedUserAgents: string[]; const newClient = () => - new WorkspaceClient( - { host, token: "test-token", authType: "pat" }, - getClientOptions(), - ); + createWorkspaceClient({ + host, + token: "test-token", + authType: "pat", + clientOptions: getClientOptions(), + }); beforeAll(async () => { server = http.createServer((req, res) => { diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 2973ee26..cc574cb6 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -1,4 +1,3 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import type { BasePlugin, CacheConfig, @@ -20,6 +19,7 @@ import { isPlainObject } from "../plugin/plugin"; import { ResourceRegistry, ResourceType } from "../registry"; import type { TelemetryConfig } from "../telemetry"; import { TelemetryManager } from "../telemetry"; +import type { WorkspaceClient } from "../workspace-client"; import { isToolProvider, PluginContext } from "./plugin-context"; const logger = createLogger("appkit"); diff --git a/packages/appkit/src/core/tests/appkit-as-user-exports.test.ts b/packages/appkit/src/core/tests/appkit-as-user-exports.test.ts index e00a37e3..7cbadcd3 100644 --- a/packages/appkit/src/core/tests/appkit-as-user-exports.test.ts +++ b/packages/appkit/src/core/tests/appkit-as-user-exports.test.ts @@ -101,10 +101,15 @@ const { MockWorkspaceClient } = vi.hoisted(() => { return { MockWorkspaceClient }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: MockWorkspaceClient, - ConfigError: class extends Error {}, -})); +vi.mock("../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (...args: unknown[]) => MockWorkspaceClient(...args), + ConfigError: class extends Error {}, + }; +}); // ── Test plugin ───────────────────────────────────────────────────── diff --git a/packages/appkit/src/index.ts b/packages/appkit/src/index.ts index b2538073..b798256f 100644 --- a/packages/appkit/src/index.ts +++ b/packages/appkit/src/index.ts @@ -116,3 +116,11 @@ export { export { appKitServingTypesPlugin } from "./type-generator/serving/vite-plugin"; // Vite plugin and type generation export { appKitTypesPlugin } from "./type-generator/vite-plugin"; +// Workspace client wrapper (facade over the Databricks SDK). Consumers pass a +// `createWorkspaceClient()` instance to `createApp({ client })`. +export { + ApiError, + createWorkspaceClient, + type WorkspaceClient, + type WorkspaceClientOptions, +} from "./workspace-client"; diff --git a/packages/appkit/src/internal-telemetry/reporter.ts b/packages/appkit/src/internal-telemetry/reporter.ts index d6edbc5e..fefb788f 100644 --- a/packages/appkit/src/internal-telemetry/reporter.ts +++ b/packages/appkit/src/internal-telemetry/reporter.ts @@ -1,4 +1,4 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; +import type { WorkspaceClient } from "../workspace-client"; import { type AppkitLog, buildAppkitPayload, diff --git a/packages/appkit/src/internal-telemetry/tests/reporter.test.ts b/packages/appkit/src/internal-telemetry/tests/reporter.test.ts index e43cdb0f..01f2b36f 100644 --- a/packages/appkit/src/internal-telemetry/tests/reporter.test.ts +++ b/packages/appkit/src/internal-telemetry/tests/reporter.test.ts @@ -1,5 +1,5 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import { afterEach, describe, expect, test, vi } from "vitest"; +import type { WorkspaceClient } from "../../workspace-client"; import { TelemetryReporter } from "../reporter"; type RequestSpy = ReturnType; diff --git a/packages/appkit/src/plugin/tests/asUser-proxy.test.ts b/packages/appkit/src/plugin/tests/asUser-proxy.test.ts index d67c4748..50d76f98 100644 --- a/packages/appkit/src/plugin/tests/asUser-proxy.test.ts +++ b/packages/appkit/src/plugin/tests/asUser-proxy.test.ts @@ -41,11 +41,16 @@ import type { ITelemetry, TelemetryProvider } from "../../telemetry"; import { TelemetryManager } from "../../telemetry"; import { isDevOboFallback, Plugin } from "../plugin"; -vi.mock("@databricks/sdk-experimental", () => ({ - ApiError: class extends Error { - statusCode = 500; - }, -})); +vi.mock("../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + ApiError: class extends Error { + statusCode = 500; + }, + }; +}); vi.mock("../../app"); vi.mock("../../cache", () => ({ diff --git a/packages/appkit/src/plugin/tests/plugin.test.ts b/packages/appkit/src/plugin/tests/plugin.test.ts index 7675f57a..2362fb8c 100644 --- a/packages/appkit/src/plugin/tests/plugin.test.ts +++ b/packages/appkit/src/plugin/tests/plugin.test.ts @@ -49,9 +49,14 @@ const { MockApiError } = vi.hoisted(() => { return { MockApiError }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - ApiError: MockApiError, -})); +vi.mock("../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + ApiError: MockApiError, + }; +}); // Mock all dependencies vi.mock("../../app"); diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 282e3cbc..8712b7c7 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1,4 +1,3 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import type express from "express"; import type { AgentToolDefinition, @@ -26,6 +25,7 @@ import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; +import type { WorkspaceClient } from "../../workspace-client"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; diff --git a/packages/appkit/src/plugins/analytics/query.ts b/packages/appkit/src/plugins/analytics/query.ts index 1075bac1..b98d6199 100644 --- a/packages/appkit/src/plugins/analytics/query.ts +++ b/packages/appkit/src/plugins/analytics/query.ts @@ -1,8 +1,8 @@ import { createHash } from "node:crypto"; -import type { sql } from "@databricks/sdk-experimental"; import { isSQLTypeMarker, type SQLTypeMarker, sql as sqlHelpers } from "shared"; import { getWorkspaceId } from "../../context"; import { ValidationError } from "../../errors"; +import type { sql } from "../../workspace-client"; type SQLParameterValue = SQLTypeMarker | null | undefined; diff --git a/packages/appkit/src/plugins/files/plugin.ts b/packages/appkit/src/plugins/files/plugin.ts index 58c3894e..bec04f24 100644 --- a/packages/appkit/src/plugins/files/plugin.ts +++ b/packages/appkit/src/plugins/files/plugin.ts @@ -1,6 +1,5 @@ import { STATUS_CODES } from "node:http"; import { Readable } from "node:stream"; -import { ApiError } from "@databricks/sdk-experimental"; import type express from "express"; import type { AgentToolDefinition, @@ -38,6 +37,7 @@ import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest, ResourceRequirement } from "../../registry"; import { ResourceType } from "../../registry"; +import { ApiError } from "../../workspace-client"; import { FILES_DOWNLOAD_DEFAULTS, FILES_MAX_UPLOAD_SIZE, diff --git a/packages/appkit/src/plugins/files/tests/delete.test.ts b/packages/appkit/src/plugins/files/tests/delete.test.ts index 748152b6..54ba4f7b 100644 --- a/packages/appkit/src/plugins/files/tests/delete.test.ts +++ b/packages/appkit/src/plugins/files/tests/delete.test.ts @@ -46,10 +46,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/download-endpoint.test.ts b/packages/appkit/src/plugins/files/tests/download-endpoint.test.ts index 550c20be..58072387 100644 --- a/packages/appkit/src/plugins/files/tests/download-endpoint.test.ts +++ b/packages/appkit/src/plugins/files/tests/download-endpoint.test.ts @@ -47,10 +47,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/error-handling.test.ts b/packages/appkit/src/plugins/files/tests/error-handling.test.ts index a16e40d0..6426b3c7 100644 --- a/packages/appkit/src/plugins/files/tests/error-handling.test.ts +++ b/packages/appkit/src/plugins/files/tests/error-handling.test.ts @@ -46,10 +46,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/mkdir.test.ts b/packages/appkit/src/plugins/files/tests/mkdir.test.ts index 8f667f5c..3aaf2670 100644 --- a/packages/appkit/src/plugins/files/tests/mkdir.test.ts +++ b/packages/appkit/src/plugins/files/tests/mkdir.test.ts @@ -46,10 +46,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/path-validation.test.ts b/packages/appkit/src/plugins/files/tests/path-validation.test.ts index 8ef0cdb6..fdefa075 100644 --- a/packages/appkit/src/plugins/files/tests/path-validation.test.ts +++ b/packages/appkit/src/plugins/files/tests/path-validation.test.ts @@ -47,10 +47,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts b/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts index 4552d7fe..7832e482 100644 --- a/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts +++ b/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts @@ -48,9 +48,9 @@ const { mockFilesApi, mockSdkClient, MockApiError } = vi.hoisted(() => { return { mockFilesApi, mockSdkClient, MockApiError }; }); -vi.mock("@databricks/sdk-experimental", async (importOriginal) => { +vi.mock("../../../workspace-client", async (importOriginal) => { const actual = - await importOriginal(); + await importOriginal(); return { ...actual, ApiError: MockApiError, diff --git a/packages/appkit/src/plugins/files/tests/plugin.test.ts b/packages/appkit/src/plugins/files/tests/plugin.test.ts index a90c4f3e..4deb4261 100644 --- a/packages/appkit/src/plugins/files/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/files/tests/plugin.test.ts @@ -54,10 +54,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockFilesApi, mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/raw-endpoint.test.ts b/packages/appkit/src/plugins/files/tests/raw-endpoint.test.ts index 99641bce..8e5851e0 100644 --- a/packages/appkit/src/plugins/files/tests/raw-endpoint.test.ts +++ b/packages/appkit/src/plugins/files/tests/raw-endpoint.test.ts @@ -47,10 +47,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/shutdown.test.ts b/packages/appkit/src/plugins/files/tests/shutdown.test.ts index 99eaa1e8..3360686b 100644 --- a/packages/appkit/src/plugins/files/tests/shutdown.test.ts +++ b/packages/appkit/src/plugins/files/tests/shutdown.test.ts @@ -39,10 +39,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/upload.test.ts b/packages/appkit/src/plugins/files/tests/upload.test.ts index 725756a9..482298f9 100644 --- a/packages/appkit/src/plugins/files/tests/upload.test.ts +++ b/packages/appkit/src/plugins/files/tests/upload.test.ts @@ -47,10 +47,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/tests/volume-config.test.ts b/packages/appkit/src/plugins/files/tests/volume-config.test.ts index 1f676afc..ed13b7ad 100644 --- a/packages/appkit/src/plugins/files/tests/volume-config.test.ts +++ b/packages/appkit/src/plugins/files/tests/volume-config.test.ts @@ -39,10 +39,15 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { return { mockClient, MockApiError, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - ApiError: MockApiError, -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: (..._args: unknown[]) => mockClient, + ApiError: MockApiError, + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); diff --git a/packages/appkit/src/plugins/files/types.ts b/packages/appkit/src/plugins/files/types.ts index d912efff..b5075751 100644 --- a/packages/appkit/src/plugins/files/types.ts +++ b/packages/appkit/src/plugins/files/types.ts @@ -1,5 +1,5 @@ -import type { files } from "@databricks/sdk-experimental"; import type { BasePluginConfig, IAppRequest } from "shared"; +import type { files } from "../../workspace-client"; import type { FilePolicy } from "./policy"; /** diff --git a/packages/appkit/src/plugins/jobs/plugin.ts b/packages/appkit/src/plugins/jobs/plugin.ts index e9a53a8b..14ca754a 100644 --- a/packages/appkit/src/plugins/jobs/plugin.ts +++ b/packages/appkit/src/plugins/jobs/plugin.ts @@ -1,5 +1,4 @@ import { STATUS_CODES } from "node:http"; -import type { jobs as jobsTypes } from "@databricks/sdk-experimental"; import type express from "express"; import type { IAppRequest, @@ -16,6 +15,7 @@ import type { ExecutionResult } from "../../plugin"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest, ResourceRequirement } from "../../registry"; import { ResourceType } from "../../registry"; +import type { jobs as jobsTypes } from "../../workspace-client"; import { JOBS_READ_DEFAULTS, JOBS_STREAM_DEFAULTS, diff --git a/packages/appkit/src/plugins/jobs/tests/plugin.test.ts b/packages/appkit/src/plugins/jobs/tests/plugin.test.ts index 9fff3535..9c27960f 100644 --- a/packages/appkit/src/plugins/jobs/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/jobs/tests/plugin.test.ts @@ -45,10 +45,15 @@ const { mockClient, mockCacheInstance } = vi.hoisted(() => { return { mockJobsApi, mockClient, mockCacheInstance }; }); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => mockClient), - Context: vi.fn(), -})); +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: () => mockClient, + Context: vi.fn(), + }; +}); vi.mock("../../../context", async (importOriginal) => { const actual = await importOriginal(); @@ -685,7 +690,7 @@ describe("JobsPlugin", () => { test("connector's cancellation token reflects signal state live", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - const { Context } = await import("@databricks/sdk-experimental"); + const { Context } = await import("../../../workspace-client"); const mockContext = Context as unknown as ReturnType; mockContext.mockClear(); diff --git a/packages/appkit/src/plugins/jobs/types.ts b/packages/appkit/src/plugins/jobs/types.ts index ca68ba0b..d10d5684 100644 --- a/packages/appkit/src/plugins/jobs/types.ts +++ b/packages/appkit/src/plugins/jobs/types.ts @@ -1,7 +1,7 @@ -import type { jobs } from "@databricks/sdk-experimental"; import type { BasePluginConfig, IAppRequest } from "shared"; import type { z } from "zod"; import type { ExecutionResult } from "../../plugin"; +import type { jobs } from "../../workspace-client"; /** Supported task types for job parameter mapping. */ export type TaskType = diff --git a/packages/appkit/src/plugins/lakebase/lakebase.ts b/packages/appkit/src/plugins/lakebase/lakebase.ts index 9fac029a..69389f4f 100644 --- a/packages/appkit/src/plugins/lakebase/lakebase.ts +++ b/packages/appkit/src/plugins/lakebase/lakebase.ts @@ -1,4 +1,3 @@ -import { WorkspaceClient } from "@databricks/sdk-experimental"; import type { QueryResult, QueryResultRow } from "pg"; import type { AgentToolDefinition, ToolProvider } from "shared"; import { z } from "zod"; @@ -24,6 +23,7 @@ import { assertReadOnlySql } from "../../core/agent/tools/sql-policy"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; +import { createWorkspaceClient } from "../../workspace-client"; import manifest from "./manifest.json"; import type { ILakebaseConfig } from "./types"; @@ -84,7 +84,9 @@ export class LakebasePlugin extends Plugin implements ToolProvider { ...this.config.pool, workspaceClient: this.config.pool?.workspaceClient ?? - new WorkspaceClient({}, getClientOptions()), + createWorkspaceClient({ + clientOptions: getClientOptions(), + }).toLegacyWorkspaceClient(), }; const user = await getUsernameWithApiLookup(poolConfig); @@ -105,7 +107,10 @@ export class LakebasePlugin extends Plugin implements ToolProvider { const isNew = !oboManager.hasPool(userKey); const pool = oboManager.getPool( userKey, - { workspaceClient: ctx.client, user: userKey }, + { + workspaceClient: ctx.client.toLegacyWorkspaceClient(), + user: userKey, + }, ctx.tokenFingerprint, ); if (isNew) { @@ -285,7 +290,11 @@ export class LakebasePlugin extends Plugin implements ToolProvider { const ctx = getUserContext(); if (ctx) { const user = ctx.userEmail ?? ctx.userId; - return { ...this.config.pool, workspaceClient: ctx.client, user }; + return { + ...this.config.pool, + workspaceClient: ctx.client.toLegacyWorkspaceClient(), + user, + }; } return this.config.pool; } 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 5f96d6ef..28024e01 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 @@ -289,7 +289,7 @@ describe("LakebasePlugin — OBO via RoutingPool", () => { await plugin.setup(); const userCtx = { - client: {} as any, + client: { toLegacyWorkspaceClient: () => ({}) } as any, userId: "user-123", userEmail: "alice@example.com", workspaceId: Promise.resolve("ws-1"), @@ -319,7 +319,7 @@ describe("LakebasePlugin — OBO via RoutingPool", () => { await plugin.setup(); const userCtx = { - client: {} as any, + client: { toLegacyWorkspaceClient: () => ({}) } as any, userId: "user-123", userEmail: "alice@example.com", workspaceId: Promise.resolve("ws-1"), @@ -349,7 +349,7 @@ describe("LakebasePlugin — OBO via RoutingPool", () => { await plugin.setup(); const userCtx = { - client: {} as any, + client: { toLegacyWorkspaceClient: () => ({}) } as any, userId: "user-123", workspaceId: Promise.resolve("ws-1"), isUserContext: true as const, diff --git a/packages/appkit/src/stream/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index d3118e7e..96101d76 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -1,6 +1,6 @@ -import type { sql } from "@databricks/sdk-experimental"; import { ExecutionError, ValidationError } from "../errors"; import { createLogger } from "../logging/logger"; +import type { sql } from "../workspace-client"; const logger = createLogger("stream:arrow"); diff --git a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts index 8a95fbdf..bd91617d 100644 --- a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts +++ b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts @@ -1,5 +1,5 @@ -import type { sql } from "@databricks/sdk-experimental"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import type { sql } from "../../workspace-client"; import { ArrowStreamProcessor } from "../arrow-stream-processor"; type ResultSchema = sql.ResultManifest["schema"]; diff --git a/packages/appkit/src/type-generator/index.ts b/packages/appkit/src/type-generator/index.ts index 0a185ae8..2034c996 100644 --- a/packages/appkit/src/type-generator/index.ts +++ b/packages/appkit/src/type-generator/index.ts @@ -1,9 +1,12 @@ import fs from "node:fs/promises"; import path from "node:path"; -import { WorkspaceClient } from "@databricks/sdk-experimental"; import dotenv from "dotenv"; import pc from "picocolors"; import { createLogger } from "../logging/logger"; +import { + createWorkspaceClient, + type WorkspaceClient, +} from "../workspace-client"; import { isRevivableMetricCacheEntry, loadCache, @@ -484,7 +487,7 @@ export async function syncMetricViewsTypes(options: { let mvClient: WorkspaceClient | undefined; const getMvClient = (): WorkspaceClient => { - mvClient ??= new WorkspaceClient({}); + mvClient ??= createWorkspaceClient(); return mvClient; }; diff --git a/packages/appkit/src/type-generator/mv-registry/describe.ts b/packages/appkit/src/type-generator/mv-registry/describe.ts index db3231d6..9f479d61 100644 --- a/packages/appkit/src/type-generator/mv-registry/describe.ts +++ b/packages/appkit/src/type-generator/mv-registry/describe.ts @@ -1,4 +1,4 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; +import type { WorkspaceClient } from "../../workspace-client"; import { type DescribeFormatMemo, describeAdaptive } from "../statement-result"; import type { DatabricksStatementExecutionResponse } from "../types"; import { isValidFqn } from "./config"; diff --git a/packages/appkit/src/type-generator/query-registry.ts b/packages/appkit/src/type-generator/query-registry.ts index 1844720f..a482c291 100644 --- a/packages/appkit/src/type-generator/query-registry.ts +++ b/packages/appkit/src/type-generator/query-registry.ts @@ -1,8 +1,8 @@ import fs from "node:fs/promises"; import path from "node:path"; -import { WorkspaceClient } from "@databricks/sdk-experimental"; import pc from "picocolors"; import { createLogger } from "../logging/logger"; +import { createWorkspaceClient } from "../workspace-client"; import { CACHE_VERSION, hashSQL, loadCache, saveCache } from "./cache"; import { getErrorDiagnostic, isConnectivityError } from "./errors"; import { decidePreflight, type PreflightMode } from "./preflight"; @@ -512,7 +512,7 @@ export async function generateQueriesFromDescribe( const queryFiles = allFiles.filter((file) => file.endsWith(".sql")); logger.debug("Found %d SQL queries", queryFiles.length); - const client = new WorkspaceClient({}); + const client = createWorkspaceClient(); const spinner = new Spinner(); // Read all SQL files in parallel diff --git a/packages/appkit/src/type-generator/serving/fetcher.ts b/packages/appkit/src/type-generator/serving/fetcher.ts index c47775d7..d970c4ed 100644 --- a/packages/appkit/src/type-generator/serving/fetcher.ts +++ b/packages/appkit/src/type-generator/serving/fetcher.ts @@ -1,5 +1,5 @@ -import { ApiError, type WorkspaceClient } from "@databricks/sdk-experimental"; import { createLogger } from "../../logging/logger"; +import { ApiError, type WorkspaceClient } from "../../workspace-client"; const logger = createLogger("type-generator:serving:fetcher"); diff --git a/packages/appkit/src/type-generator/serving/generator.ts b/packages/appkit/src/type-generator/serving/generator.ts index b377c659..da23441c 100644 --- a/packages/appkit/src/type-generator/serving/generator.ts +++ b/packages/appkit/src/type-generator/serving/generator.ts @@ -1,9 +1,12 @@ import fs from "node:fs/promises"; import path from "node:path"; -import { WorkspaceClient } from "@databricks/sdk-experimental"; import pc from "picocolors"; import { createLogger } from "../../logging/logger"; import type { EndpointConfig } from "../../plugins/serving/types"; +import { + createWorkspaceClient, + type WorkspaceClient, +} from "../../workspace-client"; import { migrateProjectConfig, removeOldGeneratedTypes, @@ -81,7 +84,7 @@ export async function generateServingTypes( }> = []; for (const [alias, config] of Object.entries(endpoints)) { - client ??= new WorkspaceClient({}); + client ??= createWorkspaceClient(); const result = await processEndpoint(alias, config, client, cache); if (result.cacheUpdated) updated = true; registryEntries.push(result.entry); diff --git a/packages/appkit/src/type-generator/serving/tests/fetcher.test.ts b/packages/appkit/src/type-generator/serving/tests/fetcher.test.ts index cae3ec7b..2f977ab4 100644 --- a/packages/appkit/src/type-generator/serving/tests/fetcher.test.ts +++ b/packages/appkit/src/type-generator/serving/tests/fetcher.test.ts @@ -1,5 +1,5 @@ -import { ApiError } from "@databricks/sdk-experimental"; import { afterEach, describe, expect, test, vi } from "vitest"; +import { ApiError } from "../../../workspace-client"; import { fetchOpenApiSchema } from "../fetcher"; function makeValidSpec( diff --git a/packages/appkit/src/type-generator/serving/tests/generator.test.ts b/packages/appkit/src/type-generator/serving/tests/generator.test.ts index edd32590..05e22428 100644 --- a/packages/appkit/src/type-generator/serving/tests/generator.test.ts +++ b/packages/appkit/src/type-generator/serving/tests/generator.test.ts @@ -18,10 +18,15 @@ vi.mock("../fetcher", () => ({ fetchOpenApiSchema: (...args: any[]) => mockFetchOpenApiSchema(...args), })); -// Mock WorkspaceClient -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => ({ config: {} })), -})); +// Mock the wrapper's client factory +vi.mock("../../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: () => ({ config: {} }), + }; +}); const CHAT_OPENAPI_SPEC = { openapi: "3.1.0", diff --git a/packages/appkit/src/type-generator/statement-result.ts b/packages/appkit/src/type-generator/statement-result.ts index 0c0c5d20..7ae091aa 100644 --- a/packages/appkit/src/type-generator/statement-result.ts +++ b/packages/appkit/src/type-generator/statement-result.ts @@ -1,5 +1,5 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import { createLogger } from "../logging/logger"; +import type { WorkspaceClient } from "../workspace-client"; import { getErrorMessage } from "./errors"; import type { DatabricksStatementExecutionResponse } from "./types"; diff --git a/packages/appkit/src/type-generator/tests/generate-queries.test.ts b/packages/appkit/src/type-generator/tests/generate-queries.test.ts index 22a054f8..48ff4fb6 100644 --- a/packages/appkit/src/type-generator/tests/generate-queries.test.ts +++ b/packages/appkit/src/type-generator/tests/generate-queries.test.ts @@ -23,12 +23,17 @@ vi.mock("node:fs/promises", () => ({ }, })); -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => ({ - statementExecution: { executeStatement: mocks.executeStatement }, - warehouses: { get: mocks.getWarehouse, start: mocks.startWarehouse }, - })), -})); +vi.mock("../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: () => ({ + statementExecution: { executeStatement: mocks.executeStatement }, + warehouses: { get: mocks.getWarehouse, start: mocks.startWarehouse }, + }), + }; +}); vi.mock("../spinner", () => ({ Spinner: vi.fn(() => ({ diff --git a/packages/appkit/src/type-generator/tests/index.test.ts b/packages/appkit/src/type-generator/tests/index.test.ts index c469bb11..adea0048 100644 --- a/packages/appkit/src/type-generator/tests/index.test.ts +++ b/packages/appkit/src/type-generator/tests/index.test.ts @@ -82,18 +82,23 @@ vi.mock("../warehouse-status", async (importOriginal) => { }; }); -// index.ts lazily constructs at most ONE `new WorkspaceClient({})` per pass -// for the whole metric path (status probe + blocking preflight + default -// DESCRIBE fetcher share it). Stub the SDK so no real credentials are needed; -// `executeStatement` doubles as the "was any metric DESCRIBE actually -// issued?" spy and the constructor mock doubles as the client-count spy. -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: vi.fn(() => ({ - statementExecution: { executeStatement: mocks.executeStatement }, - })), -})); +// index.ts lazily constructs at most ONE client per pass for the whole metric +// path (status probe + blocking preflight + default DESCRIBE fetcher share it). +// Stub the wrapper so no real credentials are needed; `executeStatement` +// doubles as the "was any metric DESCRIBE actually issued?" spy and the +// createWorkspaceClient mock doubles as the client-count spy. +vi.mock("../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: vi.fn(() => ({ + statementExecution: { executeStatement: mocks.executeStatement }, + })), + }; +}); -const { WorkspaceClient } = await import("@databricks/sdk-experimental"); +const { createWorkspaceClient } = await import("../../workspace-client"); const { generateFromEntryPoint, TypegenFatalError, TypegenSyntaxError } = await import("../index"); // The "../cache" mock spreads the actual module, so this is the real hashSQL — @@ -919,7 +924,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { expect(mocks.getWarehouseState).not.toHaveBeenCalled(); expect(mocks.startWarehouse).not.toHaveBeenCalled(); expect(mocks.waitUntilRunning).not.toHaveBeenCalled(); - expect(vi.mocked(WorkspaceClient)).not.toHaveBeenCalled(); + expect(vi.mocked(createWorkspaceClient)).not.toHaveBeenCalled(); expect(fs.readFileSync(metricFile, "utf-8")).toContain( '"total_revenue": number', ); @@ -944,7 +949,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { // The pass both probed AND described — one shared client total. expect(mocks.getWarehouseState).toHaveBeenCalledTimes(1); expect(mocks.executeStatement).toHaveBeenCalledTimes(1); - expect(vi.mocked(WorkspaceClient)).toHaveBeenCalledTimes(1); + expect(vi.mocked(createWorkspaceClient)).toHaveBeenCalledTimes(1); }); test("empty metricViews map: no probe, no preflight, no client — empty artifacts still ship", async () => { @@ -964,7 +969,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { expect(mocks.getWarehouseState).not.toHaveBeenCalled(); expect(mocks.executeStatement).not.toHaveBeenCalled(); - expect(vi.mocked(WorkspaceClient)).not.toHaveBeenCalled(); + expect(vi.mocked(createWorkspaceClient)).not.toHaveBeenCalled(); expect(fs.existsSync(metricFile)).toBe(true); // An empty metricViews map emits an empty MetricRegistry augmentation. expect(fs.readFileSync(metricFile, "utf-8")).toContain( @@ -1147,7 +1152,7 @@ describe("generateFromEntryPoint — metric cache section", () => { // All keys were hits, so the gate never even probed the warehouse ... expect(mocks.getWarehouseState).not.toHaveBeenCalled(); // ... and the whole pass constructed zero SDK clients. - expect(vi.mocked(WorkspaceClient)).not.toHaveBeenCalled(); + expect(vi.mocked(createWorkspaceClient)).not.toHaveBeenCalled(); expect(fs.readFileSync(metricFile, "utf-8")).toBe(firstDeclarations); }); diff --git a/packages/appkit/src/type-generator/tests/statement-result.test.ts b/packages/appkit/src/type-generator/tests/statement-result.test.ts index 9712d944..1a19ae5e 100644 --- a/packages/appkit/src/type-generator/tests/statement-result.test.ts +++ b/packages/appkit/src/type-generator/tests/statement-result.test.ts @@ -1,7 +1,7 @@ import fs from "node:fs"; import path from "node:path"; -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import { describe, expect, test } from "vitest"; +import type { WorkspaceClient } from "../../workspace-client"; import { type DescribeFormatMemo, describeAdaptive, diff --git a/packages/appkit/src/type-generator/tests/vite-plugin.test.ts b/packages/appkit/src/type-generator/tests/vite-plugin.test.ts index 75522556..372750a3 100644 --- a/packages/appkit/src/type-generator/tests/vite-plugin.test.ts +++ b/packages/appkit/src/type-generator/tests/vite-plugin.test.ts @@ -30,11 +30,16 @@ vi.mock("../warehouse-status", () => ({ waitUntilRunning: mocks.waitUntilRunning, })); -// armWarehouseWatch constructs `new WorkspaceClient({})`. Stub the SDK so that -// constructor is inert in unit tests. -vi.mock("@databricks/sdk-experimental", () => ({ - WorkspaceClient: class {}, -})); +// armWarehouseWatch constructs a client via createWorkspaceClient(). Stub the +// wrapper so that factory is inert in unit tests. +vi.mock("../../workspace-client", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createWorkspaceClient: () => ({}), + }; +}); const { appKitTypesPlugin } = await import("../vite-plugin"); // Real constant values: the "../index" mock spreads the actual module, so these diff --git a/packages/appkit/src/type-generator/tests/warehouse-status.test.ts b/packages/appkit/src/type-generator/tests/warehouse-status.test.ts index 700df259..ab01395d 100644 --- a/packages/appkit/src/type-generator/tests/warehouse-status.test.ts +++ b/packages/appkit/src/type-generator/tests/warehouse-status.test.ts @@ -1,5 +1,5 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import type { WorkspaceClient } from "../../workspace-client"; import { getWarehouseState, type WarehouseState, diff --git a/packages/appkit/src/type-generator/vite-plugin.ts b/packages/appkit/src/type-generator/vite-plugin.ts index b9d89401..5a937b3e 100644 --- a/packages/appkit/src/type-generator/vite-plugin.ts +++ b/packages/appkit/src/type-generator/vite-plugin.ts @@ -1,8 +1,8 @@ import { existsSync } from "node:fs"; import path from "node:path"; -import { WorkspaceClient } from "@databricks/sdk-experimental"; import type { Plugin } from "vite"; import { createLogger } from "../logging/logger"; +import { createWorkspaceClient } from "../workspace-client"; import { ANALYTICS_TYPES_FILE, generateFromEntryPoint, @@ -223,7 +223,7 @@ export function appKitTypesPlugin(options?: AppKitTypesPluginOptions): Plugin { void (async () => { try { - const client = new WorkspaceClient({}); + const client = createWorkspaceClient(); const state = await getWarehouseState(client, warehouseId); // A deleted/deleting warehouse can't be started and blocking typegen diff --git a/packages/appkit/src/type-generator/warehouse-status.ts b/packages/appkit/src/type-generator/warehouse-status.ts index aa50089e..27a0afeb 100644 --- a/packages/appkit/src/type-generator/warehouse-status.ts +++ b/packages/appkit/src/type-generator/warehouse-status.ts @@ -1,4 +1,4 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; +import type { WorkspaceClient } from "../workspace-client"; /** * Lifecycle states a SQL warehouse can report. Mirrors the Databricks SDK diff --git a/packages/appkit/src/workspace-client/client.ts b/packages/appkit/src/workspace-client/client.ts new file mode 100644 index 00000000..18ef76a1 --- /dev/null +++ b/packages/appkit/src/workspace-client/client.ts @@ -0,0 +1,71 @@ +/** + * `AppKitWorkspaceClient` — the facade implementation. Construct via + * `createWorkspaceClient(...)`; this class is internal. + * + * Every service accessor delegates to a single lazily-constructed legacy SDK + * client. This is the seam: to migrate a service to the modular SDK, replace + * its getter here with a modular client instance (and update its connector + + * the accessor type in `types.ts`). No other AppKit module touches the SDK. + */ +import { + buildLegacyWorkspaceClient, + type LegacyWorkspaceClient, + type WorkspaceClientOptions, +} from "./legacy"; +import type { WorkspaceClient } from "./types"; + +export class AppKitWorkspaceClient implements WorkspaceClient { + readonly #opts: WorkspaceClientOptions; + #legacy?: LegacyWorkspaceClient; + + constructor(opts: WorkspaceClientOptions) { + this.#opts = opts; + } + + get files() { + return this.#getLegacy().files; + } + + get warehouses() { + return this.#getLegacy().warehouses; + } + + get genie() { + return this.#getLegacy().genie; + } + + get jobs() { + return this.#getLegacy().jobs; + } + + get statementExecution() { + return this.#getLegacy().statementExecution; + } + + get servingEndpoints() { + return this.#getLegacy().servingEndpoints; + } + + get currentUser() { + return this.#getLegacy().currentUser; + } + + get config() { + return this.#getLegacy().config; + } + + get apiClient() { + return this.#getLegacy().apiClient; + } + + toLegacyWorkspaceClient(): LegacyWorkspaceClient { + return this.#getLegacy(); + } + + #getLegacy(): LegacyWorkspaceClient { + if (!this.#legacy) { + this.#legacy = buildLegacyWorkspaceClient(this.#opts); + } + return this.#legacy; + } +} diff --git a/packages/appkit/src/workspace-client/errors.ts b/packages/appkit/src/workspace-client/errors.ts new file mode 100644 index 00000000..01427e19 --- /dev/null +++ b/packages/appkit/src/workspace-client/errors.ts @@ -0,0 +1,14 @@ +/** + * API error type surfaced by the AppKit workspace client. + * + * Re-exports the legacy `@databricks/sdk-experimental` `ApiError` so AppKit + * catch sites (`instanceof ApiError`, `.statusCode`) and the one construction + * site in `connectors/files/client.ts` keep working unchanged while importing + * from the wrapper instead of the SDK. + * + * When the first modular service lands (throwing `@databricks/sdk-core`'s + * distinct `ApiError`), add an `isApiError(e)` / `getApiErrorStatusCode(e)` + * predicate here to unify the two shapes — `instanceof` can't, since they're + * classes from different packages. + */ +export { ApiError } from "@databricks/sdk-experimental"; diff --git a/packages/appkit/src/workspace-client/factory.ts b/packages/appkit/src/workspace-client/factory.ts new file mode 100644 index 00000000..e1c38fb3 --- /dev/null +++ b/packages/appkit/src/workspace-client/factory.ts @@ -0,0 +1,23 @@ +/** + * Public factory for constructing a wrapper instance. + */ +import { AppKitWorkspaceClient } from "./client"; +import type { WorkspaceClientOptions } from "./legacy"; +import type { WorkspaceClient } from "./types"; + +/** + * Construct an AppKit workspace client. + * + * Auth resolution: + * - If `opts.token` is set, uses PAT credentials. + * - Otherwise walks the SDK default auth chain (env vars + ~/.databrickscfg). + * + * Host resolution: + * - Explicit `opts.host` → use it. + * - Otherwise resolved by the SDK from `DATABRICKS_HOST` / profile. + */ +export function createWorkspaceClient( + opts: WorkspaceClientOptions = {}, +): WorkspaceClient { + return new AppKitWorkspaceClient(opts); +} diff --git a/packages/appkit/src/workspace-client/index.ts b/packages/appkit/src/workspace-client/index.ts new file mode 100644 index 00000000..152e8080 --- /dev/null +++ b/packages/appkit/src/workspace-client/index.ts @@ -0,0 +1,29 @@ +/** + * AppKit workspace-client wrapper — the single entry point every AppKit module + * uses to reach a Databricks SDK client. Isolates `@databricks/sdk-experimental` + * to `./legacy.ts` so services can migrate to the modular SDK incrementally + * behind a stable facade. + */ +export { ApiError } from "./errors"; +export { createWorkspaceClient } from "./factory"; +export type { + CancellationToken, + ClientOptions, + GenieMessage, + Waiter, + WorkspaceClientOptions, +} from "./legacy"; +// SDK value + type re-exports so AppKit modules import them from the wrapper. +export { + ConfigError, + Context, + Time, + TimeUnits, +} from "./legacy"; +export type { + files, + jobs, + serving, + sql, + WorkspaceClient, +} from "./types"; diff --git a/packages/appkit/src/workspace-client/legacy.ts b/packages/appkit/src/workspace-client/legacy.ts new file mode 100644 index 00000000..15dffa73 --- /dev/null +++ b/packages/appkit/src/workspace-client/legacy.ts @@ -0,0 +1,88 @@ +/** + * The single module in AppKit allowed to import `@databricks/sdk-experimental` + * directly (enforced by the Biome `noRestrictedImports` boundary rule). Every + * other AppKit module reaches the SDK through the wrapper's re-exports and the + * {@link WorkspaceClient} facade. + * + * Isolating the SDK here is what makes the incremental migration to the modular + * Databricks SDK a localized change: to migrate a service, swap its getter in + * `client.ts` from the legacy delegate to the modular client — nothing else in + * the codebase imports the SDK, so the blast radius is one connector + its test. + */ + +import type { + ClientOptions, + WorkspaceClient as SdkWorkspaceClient, +} from "@databricks/sdk-experimental"; +import * as SDK from "@databricks/sdk-experimental"; + +const { WorkspaceClient: SdkWorkspaceClientCtor } = SDK; + +/** The concrete legacy SDK client type. */ +export type LegacyWorkspaceClient = SdkWorkspaceClient; + +/** + * Options used to construct the wrapper. Mirrors the subset of the old SDK's + * `Config` + `ClientOptions` that AppKit relies on today; we deliberately do + * NOT re-expose every old-SDK config knob. + */ +export interface WorkspaceClientOptions { + /** Databricks host, e.g. https://my-workspace.cloud.databricks.com. Defaults to DATABRICKS_HOST / profile resolution. */ + host?: string; + /** Bearer token. When set, `authType` defaults to "pat". */ + token?: string; + /** Authentication strategy passed to the legacy client. */ + authType?: "pat"; + /** + * SDK client options (product / productVersion / userAgentExtra) used to + * stamp the outbound User-Agent. Produced by `getClientOptions()`; omitted + * by build-time callers that don't stamp a User-Agent. + */ + clientOptions?: ClientOptions; +} + +/** + * Construct a legacy `WorkspaceClient` from wrapper options. + * + * Centralised so the wrapper facade, the `.toLegacyWorkspaceClient()` escape + * hatch, and any per-request OBO client all build it the same way. + */ +export function buildLegacyWorkspaceClient( + opts: WorkspaceClientOptions, +): LegacyWorkspaceClient { + const cfg = opts.token + ? { host: opts.host, token: opts.token, authType: opts.authType ?? "pat" } + : opts.host + ? { host: opts.host } + : {}; + return new SdkWorkspaceClientCtor(cfg, opts.clientOptions); +} + +// ── SDK type re-exports ────────────────────────────────────────────── +export type { + CancellationToken, + ClientOptions, +} from "@databricks/sdk-experimental"; +// ── SDK value re-exports ───────────────────────────────────────────── +// +// AppKit modules import these from the wrapper instead of the SDK so the +// boundary rule holds. `Context` bridges AbortSignal → CancellationToken +// (serving/jobs/sql-warehouse); `Time`/`TimeUnits` drive genie polling; +// `ConfigError` is matched in service-context's auth-failure handling. +// +// These are sourced off the namespace import rather than `export { ... } from` +// because the SDK is CommonJS: its `Time` export is emitted as a getter that +// calls `__importDefault(...)`, which defeats Node's static named-export +// detection — a direct `export { Time }` throws "does not provide an export +// named 'Time'" at ESM link time. `Time` is only reachable via the module +// object, so we fall back to `SDK.default.Time` (matching the original genie +// connector's `SDK.Time ?? SDK.default.Time` guard). +export const { ConfigError, Context, TimeUnits } = SDK; +export const Time = + SDK.Time ?? (SDK as unknown as { default: typeof SDK }).default.Time; + +// Deep-import types used by the genie connector's waiter idiom. Not exposed +// on the SDK's top-level index, so re-exported here to keep the genie +// connector off a direct `@databricks/sdk-experimental/dist/**` import. +export type { GenieMessage } from "@databricks/sdk-experimental/dist/apis/dashboards"; +export type { Waiter } from "@databricks/sdk-experimental/dist/wait"; diff --git a/packages/appkit/src/workspace-client/types.ts b/packages/appkit/src/workspace-client/types.ts new file mode 100644 index 00000000..6d4db6ae --- /dev/null +++ b/packages/appkit/src/workspace-client/types.ts @@ -0,0 +1,77 @@ +/** + * The wrapper facade type and the SDK type-namespace re-exports used across + * AppKit. + * + * Every service accessor is currently typed against the legacy SDK and + * delegated through `toLegacyWorkspaceClient()` — this PR introduces the seam + * only, it does not migrate any service to the modular SDK. Migrating a service + * later means changing that accessor's type here + its getter in `client.ts` + * and updating the one connector that consumes it. + * + * Type-namespace re-exports (`files`, `jobs`, `serving`, `sql`) point at + * `@databricks/sdk-experimental` so existing AppKit call-site shapes stay + * stable. They move to the modular `@databricks/sdk-` model exports + * as each service migrates. + */ +import type { LegacyWorkspaceClient } from "./legacy"; + +// SDK type namespaces, re-exported so AppKit modules import them from the +// wrapper rather than the SDK directly. +export type { + files, + jobs, + serving, + sql, +} from "@databricks/sdk-experimental"; + +/** + * AppKit's workspace client facade. Mirrors the multi-client shape of the + * modular Databricks SDK: each service is its own accessor, so services can be + * migrated one at a time behind this stable interface. + * + * Accessors are legacy-typed for now (delegated to the underlying legacy SDK + * client); see the module docblock. + */ +export interface WorkspaceClient { + /** UC Volumes / Files API. */ + readonly files: LegacyWorkspaceClient["files"]; + + /** SQL Warehouses. */ + readonly warehouses: LegacyWorkspaceClient["warehouses"]; + + /** Genie / dashboards. */ + readonly genie: LegacyWorkspaceClient["genie"]; + + /** Jobs. */ + readonly jobs: LegacyWorkspaceClient["jobs"]; + + /** Statement Execution. */ + readonly statementExecution: LegacyWorkspaceClient["statementExecution"]; + + /** Serving Endpoints. */ + readonly servingEndpoints: LegacyWorkspaceClient["servingEndpoints"]; + + /** Current user. */ + readonly currentUser: LegacyWorkspaceClient["currentUser"]; + + /** + * SDK `Config` — exposes `host` and `authenticate(headers)`. Used by the + * files-upload path and agents auth-header stamping, which bypass the typed + * services. + */ + readonly config: LegacyWorkspaceClient["config"]; + + /** + * Low-level HTTP transport (`apiClient.request(...)`). Used for endpoints + * without a typed service method: SCIM header probe, warehouse listing, + * serving SSE streaming, vector search, internal telemetry. + */ + readonly apiClient: LegacyWorkspaceClient["apiClient"]; + + /** + * Returns the underlying legacy `@databricks/sdk-experimental` + * `WorkspaceClient`, for handoff to code still typed against the old SDK + * (`@databricks/lakebase`). Transitional. + */ + toLegacyWorkspaceClient(): LegacyWorkspaceClient; +}