diff --git a/.gitignore b/.gitignore index 3b6cc969..fdbf6744 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,7 @@ coverage *.tsbuildinfo .turbo + +# AI generated files +.serena +docs/plans/ \ No newline at end of file diff --git a/packages/appkit/src/context/execution-context.ts b/packages/appkit/src/context/execution-context.ts index e201fec6..663511bb 100644 --- a/packages/appkit/src/context/execution-context.ts +++ b/packages/appkit/src/context/execution-context.ts @@ -1,4 +1,5 @@ import { AsyncLocalStorage } from "node:async_hooks"; +import { ConfigurationError } from "../errors"; import { ServiceContext } from "./service-context"; import { type ExecutionContext, @@ -64,7 +65,14 @@ export function getWorkspaceClient() { * Get the warehouse ID promise. */ export function getWarehouseId(): Promise { - return getExecutionContext().warehouseId; + const ctx = getExecutionContext(); + if (!ctx.warehouseId) { + throw ConfigurationError.resourceNotFound( + "Warehouse ID", + 'Ensure a plugin declares static requiredResources = ["warehouseId"] or set DATABRICKS_WAREHOUSE_ID', + ); + } + return ctx.warehouseId; } /** diff --git a/packages/appkit/src/context/service-context.ts b/packages/appkit/src/context/service-context.ts index ee4b2453..fd606bd5 100644 --- a/packages/appkit/src/context/service-context.ts +++ b/packages/appkit/src/context/service-context.ts @@ -4,6 +4,7 @@ import { WorkspaceClient, } from "@databricks/sdk-experimental"; import { coerce } from "semver"; +import type { ServiceContextResource } from "shared"; import { name as productName, version as productVersion, @@ -24,8 +25,8 @@ export interface ServiceContextState { client: WorkspaceClient; /** The service principal's user ID */ serviceUserId: string; - /** Promise that resolves to the warehouse ID */ - warehouseId: Promise; + /** Promise that resolves to the warehouse ID (only present when a plugin requires it) */ + warehouseId?: Promise; /** Promise that resolves to the workspace ID */ workspaceId: Promise; } @@ -62,6 +63,7 @@ export class ServiceContext { * of creating one from environment credentials. */ static async initialize( + requiredResources: ServiceContextResource[] = [], client?: WorkspaceClient, ): Promise { if (ServiceContext.instance) { @@ -72,7 +74,10 @@ export class ServiceContext { return ServiceContext.initPromise; } - ServiceContext.initPromise = ServiceContext.createContext(client); + ServiceContext.initPromise = ServiceContext.createContext( + requiredResources, + client, + ); ServiceContext.instance = await ServiceContext.initPromise; return ServiceContext.instance; } @@ -153,11 +158,15 @@ export class ServiceContext { } private static async createContext( + requiredResources: ServiceContextResource[] = [], client?: WorkspaceClient, ): Promise { const wsClient = client ?? new WorkspaceClient({}, getClientOptions()); - const warehouseId = ServiceContext.getWarehouseId(wsClient); + const warehouseId = requiredResources.includes("warehouseId") + ? ServiceContext.getWarehouseId(wsClient) + : undefined; + const workspaceId = ServiceContext.getWorkspaceId(wsClient); const currentUser = await wsClient.currentUser.me(); diff --git a/packages/appkit/src/context/user-context.ts b/packages/appkit/src/context/user-context.ts index e106510e..7b409f9b 100644 --- a/packages/appkit/src/context/user-context.ts +++ b/packages/appkit/src/context/user-context.ts @@ -11,8 +11,8 @@ export interface UserContext { userId: string; /** The user's name (from request headers) */ userName?: string; - /** Promise that resolves to the warehouse ID (inherited from service context) */ - warehouseId: Promise; + /** Promise that resolves to the warehouse ID (inherited from service context, only present when a plugin requires it) */ + warehouseId?: Promise; /** Promise that resolves to the workspace ID (inherited from service context) */ workspaceId: Promise; /** Flag indicating this is a user context */ diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index ed226b36..017966a6 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -7,6 +7,7 @@ import type { PluginConstructor, PluginData, PluginMap, + ServiceContextResource, } from "shared"; import { CacheManager } from "../cache"; import { ServiceContext } from "../context"; @@ -149,11 +150,13 @@ export class AppKit { TelemetryManager.initialize(config?.telemetry); await CacheManager.getInstance(config?.cache); - // Initialize ServiceContext for Databricks client management - // This provides the service principal client and shared resources - await ServiceContext.initialize(config?.client); - + // Collect required resources from all plugins before initializing context const rawPlugins = config.plugins as T; + const requiredResources = AppKit.collectRequiredResources(rawPlugins); + + // Initialize ServiceContext for Databricks client management + // Only resolves resources that plugins actually need + await ServiceContext.initialize(requiredResources, config?.client); const preparedPlugins = AppKit.preparePlugins(rawPlugins); const mergedConfig = { plugins: preparedPlugins, @@ -178,6 +181,18 @@ export class AppKit { } return result; } + + private static collectRequiredResources( + plugins: PluginData[], + ): ServiceContextResource[] { + const resources = new Set(); + for (const { plugin } of plugins) { + for (const req of plugin.requiredResources ?? []) { + resources.add(req); + } + } + return [...resources]; + } } /** diff --git a/packages/appkit/src/plugin/plugin.ts b/packages/appkit/src/plugin/plugin.ts index c5050ca9..3a1e8161 100644 --- a/packages/appkit/src/plugin/plugin.ts +++ b/packages/appkit/src/plugin/plugin.ts @@ -8,6 +8,7 @@ import type { PluginExecutionSettings, PluginPhase, RouteConfig, + ServiceContextResource, StreamExecuteHandler, StreamExecutionSettings, } from "shared"; @@ -76,6 +77,7 @@ export abstract class Plugin< private registeredEndpoints: PluginEndpointMap = {}; static phase: PluginPhase = "normal"; + static requiredResources: ServiceContextResource[] = []; name: string; constructor(protected config: TConfig) { diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index a631a776..43b269ca 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -3,6 +3,7 @@ import type express from "express"; import type { IAppRouter, PluginExecuteConfig, + ServiceContextResource, SQLTypeMarker, StreamExecutionSettings, } from "shared"; @@ -25,6 +26,7 @@ import type { const logger = createLogger("analytics"); export class AnalyticsPlugin extends Plugin { + static requiredResources: ServiceContextResource[] = ["warehouseId"]; name = "analytics"; protected envVars: string[] = []; diff --git a/packages/shared/src/plugin.ts b/packages/shared/src/plugin.ts index a30260aa..a596f1db 100644 --- a/packages/shared/src/plugin.ts +++ b/packages/shared/src/plugin.ts @@ -46,6 +46,8 @@ export interface PluginConfig { export type PluginPhase = "core" | "normal" | "deferred"; +export type ServiceContextResource = "warehouseId"; + export type PluginConstructor< C = BasePluginConfig, I extends BasePlugin = BasePlugin, @@ -54,6 +56,7 @@ export type PluginConstructor< ) => I) & { DEFAULT_CONFIG?: Record; phase?: PluginPhase; + requiredResources?: ServiceContextResource[]; }; export type ConfigFor = T extends { DEFAULT_CONFIG: infer D }