@databricks/appkit 0.34.1 → 0.35.1

package/dist/plugins/lakebase/lakebase.js

@@ -1,4 +1,7 @@
 import { createLogger } from "../../logging/logger.js";
+import { createLakebasePoolManager } from "../../connectors/lakebase/pool-manager.js";
+import { getUserContext, init_execution_context } from "../../context/execution-context.js";
+import { RoutingPool } from "../../connectors/lakebase/routing-pool.js";
 import { createLakebasePool, getLakebaseOrmConfig, getLakebasePgConfig, getUsernameWithApiLookup } from "../../connectors/lakebase/index.js";
 import { Plugin } from "../../plugin/plugin.js";
 import { toPlugin } from "../../plugin/to-plugin.js";
@@ -10,13 +13,26 @@ import manifest_default from "./manifest.js";
 import { z } from "zod";
 
 //#region src/plugins/lakebase/lakebase.ts
+init_execution_context();
 const logger = createLogger("lakebase");
+/** Default pool settings for per-user OBO pools. */
+const OBO_POOL_DEFAULTS = {
+	max: 3,
+	allowExitOnIdle: true,
+	idleTimeoutMillis: 3e4
+};
 /**
 * AppKit plugin for Databricks Lakebase Autoscaling.
 *
 * Wraps `@databricks/lakebase` to provide a standard `pg.Pool` with automatic
 * OAuth token refresh, integrated with AppKit's logger and OpenTelemetry setup.
 *
+* Supports On-Behalf-Of (OBO) via `asUser(req)` — each user gets a separate
+* `pg.Pool` authenticated with their Databricks identity, enabling features
+* like Row-Level Security (RLS). Routing is handled transparently by
+* {@link RoutingPool}, which reads the execution context set by the base
+* class `asUser()`.
+*
 * @example
 * ```ts
 * import { createApp, lakebase, server } from "@databricks/appkit";
@@ -25,32 +41,58 @@ const logger = createLogger("lakebase");
 *   plugins: [server(), lakebase()],
 * });
 *
+* // Service principal query
 * const result = await AppKit.lakebase.query("SELECT * FROM users WHERE id = $1", [userId]);
+*
+* // User-scoped query (per-user pool, RLS enforced)
+* const mine = await AppKit.lakebase.asUser(req).query("SELECT * FROM my_data");
 * ```
 */
 var LakebasePlugin = class extends Plugin {
 	/** Plugin manifest declaring metadata and resource requirements */
 	static manifest = manifest_default;
 	pool = null;
+	oboPoolManager = null;
 	/**
-	* Initializes the Lakebase connection pool.
+	* Initializes the Lakebase connection pool and OBO pool manager.
 	* Called automatically by AppKit during the plugin setup phase.
 	*
-	* Resolves the PostgreSQL username via {@link getUsernameWithApiLookup},
-	* which tries config, env vars, and finally the Databricks workspace API.
+	* Creates a {@link RoutingPool} that automatically routes queries to either
+	* the service-principal pool or a per-user pool based on the execution
+	* context (set by `Plugin.asUser(req)` via AsyncLocalStorage).
 	*/
 	async setup() {
 		const poolConfig = this.config.pool;
 		const user = await getUsernameWithApiLookup(poolConfig);
-		this.pool = createLakebasePool({
+		const spPool = createLakebasePool({
 			...poolConfig,
 			user
 		});
-		logger.info("Lakebase pool initialized");
+		logger.info("Lakebase SP pool initialized");
+		this.oboPoolManager = createLakebasePoolManager({
+			...poolConfig,
+			...OBO_POOL_DEFAULTS
+		});
+		logger.info("Lakebase OBO pool manager initialized");
+		const oboManager = this.oboPoolManager;
+		this.pool = new RoutingPool(spPool, (ctx) => {
+			if (!oboManager) throw new Error("OBO pool manager not initialized");
+			const userKey = ctx.userEmail ?? ctx.userId;
+			const isNew = !oboManager.hasPool(userKey);
+			const pool = oboManager.getPool(userKey, {
+				workspaceClient: ctx.client,
+				user: userKey
+			}, ctx.tokenFingerprint);
+			if (isNew) logger.debug("Created OBO pool for user (total: %d)", oboManager.size);
+			return pool;
+		});
 	}
 	/**
 	* Executes a parameterized SQL query against the Lakebase pool.
 	*
+	* When called inside `asUser(req)`, the query automatically routes to
+	* the per-user pool via {@link RoutingPool}.
+	*
 	* @param text - SQL query string, using `$1`, `$2`, ... placeholders
 	* @param values - Parameter values corresponding to placeholders
 	* @returns Query result with typed rows
@@ -94,28 +136,27 @@ var LakebasePlugin = class extends Plugin {
 		}
 	}
 	/**
-	* Gracefully drains and closes the connection pool.
+	* Gracefully drains and closes all connection pools (SP + OBO).
 	* Called automatically by AppKit during shutdown.
 	*/
 	abortActiveOperations() {
 		super.abortActiveOperations();
 		if (this.pool) {
-			logger.info("Closing Lakebase pool");
+			logger.info("Closing Lakebase SP pool");
 			this.pool.end().catch((err) => {
-				logger.error("Error closing Lakebase pool: %O", err);
+				logger.error("Error closing Lakebase SP pool: %O", err);
 			});
 			this.pool = null;
 		}
+		if (this.oboPoolManager) {
+			logger.info("Closing all Lakebase OBO pools (%d)", this.oboPoolManager.size);
+			this.oboPoolManager.closeAll().catch((err) => {
+				logger.error("Error closing Lakebase OBO pools: %O", err);
+			});
+			this.oboPoolManager = null;
+		}
 	}
 	/**
-	* Returns the plugin's public API, accessible via `AppKit.lakebase`.
-	*
-	* - `pool` — The raw `pg.Pool` instance, for use with ORMs or advanced scenarios
-	* - `query` — Convenience method for executing parameterized SQL queries
-	* - `getOrmConfig()` — Returns a config object compatible with Drizzle, TypeORM, Sequelize, etc.
-	* - `getPgConfig()` — Returns a `pg.PoolConfig` object for manual pool construction
-	*/
-	/**
 	* Agent tool registry. Empty by default — the Lakebase plugin does NOT
 	* expose its SQL connection to LLM agents unless the developer explicitly
 	* opts in via `config.exposeAsAgentTool`. See {@link buildQueryTool}.
@@ -126,20 +167,21 @@ var LakebasePlugin = class extends Plugin {
 		this.config = config;
 		if (config.exposeAsAgentTool) {
 			this.tools = { query: this.buildQueryTool(config.exposeAsAgentTool) };
-			logger.warn("Lakebase agent tool is enabled (readOnly=%s). Every agent with access to this plugin can execute SQL against the Lakebase database as the service principal.", config.exposeAsAgentTool.readOnly !== false);
+			logger.warn("Lakebase agent tool is enabled (readOnly=%s). Every agent with access to this plugin can execute SQL against the Lakebase database as the requesting user's identity.", config.exposeAsAgentTool.readOnly !== false);
 		}
 	}
 	buildQueryTool(opt) {
 		const readOnly = opt.readOnly !== false;
 		return defineTool({
-			description: readOnly ? "Execute a read-only SQL query against the Lakebase PostgreSQL database. Only SELECT, WITH, SHOW, EXPLAIN, and DESCRIBE statements are accepted. Use $1, $2, etc. as placeholders and pass values separately. Runs as the application's service principal." : "Execute a parameterized SQL statement against the Lakebase PostgreSQL database. Use $1, $2, etc. as placeholders and pass values separately. Runs as the application's service principal. This tool can modify data; every invocation requires explicit human approval.",
+			description: readOnly ? "Execute a read-only SQL query against the Lakebase PostgreSQL database. Only SELECT, WITH, SHOW, EXPLAIN, and DESCRIBE statements are accepted. Use $1, $2, etc. as placeholders and pass values separately." : "Execute a parameterized SQL statement against the Lakebase PostgreSQL database. Use $1, $2, etc. as placeholders and pass values separately. This tool can modify data; every invocation requires explicit human approval.",
 			schema: z.object({
 				text: z.string().describe("SQL statement with $1, $2, ... placeholders for parameters"),
 				values: z.array(z.unknown()).optional().describe("Parameter values corresponding to placeholders")
 			}),
 			annotations: {
 				effect: readOnly ? "read" : "destructive",
-				idempotent: false
+				idempotent: false,
+				requiresUserContext: true
 			},
 			execute: async (args, signal) => {
 				signal?.throwIfAborted();
@@ -160,12 +202,40 @@ var LakebasePlugin = class extends Plugin {
 	toolkit(opts) {
 		return buildToolkitEntries(this.name, this.tools, opts);
 	}
+	/**
+	* Returns the pool config for the current execution context.
+	* Inside `asUser(req)`, returns user-scoped config; otherwise SP config.
+	*/
+	activePoolConfig() {
+		const ctx = getUserContext();
+		if (ctx) {
+			const user = ctx.userEmail ?? ctx.userId;
+			return {
+				...this.config.pool,
+				workspaceClient: ctx.client,
+				user
+			};
+		}
+		return this.config.pool;
+	}
+	/**
+	* Returns the plugin's public API, accessible via `AppKit.lakebase`.
+	*
+	* - `pool` — The connection pool (routes to per-user pool when inside `asUser(req)`)
+	* - `query` — Convenience method for executing parameterized SQL queries
+	* - `getOrmConfig()` — Returns a config object compatible with Drizzle, TypeORM, Sequelize, etc.
+	*   Inside `asUser(req)`, returns user-scoped config.
+	* - `getPgConfig()` — Returns a `pg.PoolConfig` object for manual pool construction.
+	*   Inside `asUser(req)`, returns user-scoped config.
+	*
+	* Use `AppKit.lakebase.asUser(req)` to get the same API backed by a per-user pool.
+	*/
 	exports() {
 		return {
 			pool: this.pool,
 			query: this.query.bind(this),
-			getOrmConfig: () => getLakebaseOrmConfig(this.config.pool),
-			getPgConfig: () => getLakebasePgConfig(this.config.pool)
+			getOrmConfig: () => getLakebaseOrmConfig(this.activePoolConfig()),
+			getPgConfig: () => getLakebasePgConfig(this.activePoolConfig())
 		};
 	}
 };

package/dist/plugins/lakebase/lakebase.js.map

@@ -1 +1 @@
-{"version":3,"file":"lakebase.js","names":["manifest"],"sources":["../../../src/plugins/lakebase/lakebase.ts"],"sourcesContent":["import type { Pool, QueryResult, QueryResultRow } from \"pg\";\nimport type { AgentToolDefinition, ToolProvider } from \"shared\";\nimport { z } from \"zod\";\nimport {\n  createLakebasePool,\n  getLakebaseOrmConfig,\n  getLakebasePgConfig,\n  getUsernameWithApiLookup,\n} from \"../../connectors/lakebase\";\nimport { buildToolkitEntries } from \"../../core/agent/build-toolkit\";\nimport {\n  defineTool,\n  executeFromRegistry,\n  toolsFromRegistry,\n} from \"../../core/agent/tools/define-tool\";\nimport { assertReadOnlySql } from \"../../core/agent/tools/sql-policy\";\nimport { createLogger } from \"../../logging/logger\";\nimport { Plugin, toPlugin } from \"../../plugin\";\nimport type { PluginManifest } from \"../../registry\";\nimport manifest from \"./manifest.json\";\nimport type { ILakebaseConfig } from \"./types\";\n\nconst logger = createLogger(\"lakebase\");\n\n/**\n * AppKit plugin for Databricks Lakebase Autoscaling.\n *\n * Wraps `@databricks/lakebase` to provide a standard `pg.Pool` with automatic\n * OAuth token refresh, integrated with AppKit's logger and OpenTelemetry setup.\n *\n * @example\n * ```ts\n * import { createApp, lakebase, server } from \"@databricks/appkit\";\n *\n * const AppKit = await createApp({\n *   plugins: [server(), lakebase()],\n * });\n *\n * const result = await AppKit.lakebase.query(\"SELECT * FROM users WHERE id = $1\", [userId]);\n * ```\n */\nexport class LakebasePlugin extends Plugin implements ToolProvider {\n  /** Plugin manifest declaring metadata and resource requirements */\n  static manifest = manifest as PluginManifest<\"lakebase\">;\n\n  protected declare config: ILakebaseConfig;\n  private pool: Pool | null = null;\n\n  /**\n   * Initializes the Lakebase connection pool.\n   * Called automatically by AppKit during the plugin setup phase.\n   *\n   * Resolves the PostgreSQL username via {@link getUsernameWithApiLookup},\n   * which tries config, env vars, and finally the Databricks workspace API.\n   */\n  async setup() {\n    const poolConfig = this.config.pool;\n    const user = await getUsernameWithApiLookup(poolConfig);\n    this.pool = createLakebasePool({ ...poolConfig, user });\n    logger.info(\"Lakebase pool initialized\");\n  }\n\n  /**\n   * Executes a parameterized SQL query against the Lakebase pool.\n   *\n   * @param text - SQL query string, using `$1`, `$2`, ... placeholders\n   * @param values - Parameter values corresponding to placeholders\n   * @returns Query result with typed rows\n   *\n   * @example\n   * ```ts\n   * const result = await AppKit.lakebase.query<{ id: number; name: string }>(\n   *   \"SELECT id, name FROM users WHERE active = $1\",\n   *   [true],\n   * );\n   * ```\n   */\n  async query<T extends QueryResultRow = any>(\n    text: string,\n    values?: unknown[],\n  ): Promise<QueryResult<T>> {\n    // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup(), which AppKit always awaits before exposing the plugin API\n    return this.pool!.query<T>(text, values);\n  }\n\n  /**\n   * Execute a single statement inside a `BEGIN READ ONLY … ROLLBACK`\n   * transaction on a dedicated client.\n   *\n   * The three commands MUST share a connection — a naive\n   * `pool.query(\"BEGIN READ ONLY; <stmt>; ROLLBACK\")` batch cannot accept\n   * parameter values (PostgreSQL's Extended Query protocol rejects multi-\n   * statement prepared queries), which would silently break every\n   * parameterized query the agent tool issues.\n   *\n   * Returns the raw `rows` array for the user's statement. Side effects the\n   * statement may attempt (writes, writable-function side effects) are\n   * rejected by PostgreSQL under the read-only transaction posture.\n   */\n  private async runReadOnlyStatement(\n    text: string,\n    values?: unknown[],\n  ): Promise<unknown[]> {\n    // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup()\n    const client = await this.pool!.connect();\n    try {\n      await client.query(\"BEGIN READ ONLY\");\n      const result = await client.query(text, values);\n      return result.rows;\n    } finally {\n      try {\n        await client.query(\"ROLLBACK\");\n      } finally {\n        client.release();\n      }\n    }\n  }\n\n  /**\n   * Gracefully drains and closes the connection pool.\n   * Called automatically by AppKit during shutdown.\n   */\n  abortActiveOperations(): void {\n    super.abortActiveOperations();\n    if (this.pool) {\n      logger.info(\"Closing Lakebase pool\");\n      this.pool.end().catch((err) => {\n        logger.error(\"Error closing Lakebase pool: %O\", err);\n      });\n      this.pool = null;\n    }\n  }\n\n  /**\n   * Returns the plugin's public API, accessible via `AppKit.lakebase`.\n   *\n   * - `pool` — The raw `pg.Pool` instance, for use with ORMs or advanced scenarios\n   * - `query` — Convenience method for executing parameterized SQL queries\n   * - `getOrmConfig()` — Returns a config object compatible with Drizzle, TypeORM, Sequelize, etc.\n   * - `getPgConfig()` — Returns a `pg.PoolConfig` object for manual pool construction\n   */\n\n  /**\n   * Agent tool registry. Empty by default — the Lakebase plugin does NOT\n   * expose its SQL connection to LLM agents unless the developer explicitly\n   * opts in via `config.exposeAsAgentTool`. See {@link buildQueryTool}.\n   */\n  private tools: Record<string, ReturnType<typeof this.buildQueryTool>> = {};\n\n  constructor(config: ILakebaseConfig) {\n    super(config);\n    this.config = config;\n    if (config.exposeAsAgentTool) {\n      this.tools = { query: this.buildQueryTool(config.exposeAsAgentTool) };\n      logger.warn(\n        \"Lakebase agent tool is enabled (readOnly=%s). Every agent with access to this plugin can execute SQL against the Lakebase database as the service principal.\",\n        config.exposeAsAgentTool.readOnly !== false,\n      );\n    }\n  }\n\n  private buildQueryTool(\n    opt: NonNullable<ILakebaseConfig[\"exposeAsAgentTool\"]>,\n  ) {\n    const readOnly = opt.readOnly !== false;\n    return defineTool({\n      description: readOnly\n        ? \"Execute a read-only SQL query against the Lakebase PostgreSQL database. Only SELECT, WITH, SHOW, EXPLAIN, and DESCRIBE statements are accepted. Use $1, $2, etc. as placeholders and pass values separately. Runs as the application's service principal.\"\n        : \"Execute a parameterized SQL statement against the Lakebase PostgreSQL database. Use $1, $2, etc. as placeholders and pass values separately. Runs as the application's service principal. This tool can modify data; every invocation requires explicit human approval.\",\n      schema: z.object({\n        text: z\n          .string()\n          .describe(\n            \"SQL statement with $1, $2, ... placeholders for parameters\",\n          ),\n        values: z\n          .array(z.unknown())\n          .optional()\n          .describe(\"Parameter values corresponding to placeholders\"),\n      }),\n      annotations: {\n        effect: readOnly ? \"read\" : \"destructive\",\n        idempotent: false,\n      },\n      execute: async (args, signal) => {\n        // Matches the files plugin pattern: the pg connection API\n        // doesn't accept AbortSignal in its current shape, so deeper\n        // mid-call cancellation needs a separate plumbing pass on the\n        // connector. This entry check still catches the common case —\n        // a tool dispatched after the user already cancelled the\n        // stream — and unwinds cleanly instead of running to\n        // completion against the SQL warehouse.\n        signal?.throwIfAborted();\n        if (readOnly) {\n          assertReadOnlySql(args.text);\n          return this.runReadOnlyStatement(args.text, args.values);\n        }\n        const result = await this.query(args.text, args.values);\n        return result.rows;\n      },\n    });\n  }\n\n  getAgentTools(): AgentToolDefinition[] {\n    return toolsFromRegistry(this.tools);\n  }\n\n  async executeAgentTool(\n    name: string,\n    args: unknown,\n    signal?: AbortSignal,\n  ): Promise<unknown> {\n    return executeFromRegistry(this.tools, name, args, signal);\n  }\n\n  toolkit(opts?: import(\"../../core/agent/types\").ToolkitOptions) {\n    return buildToolkitEntries(this.name, this.tools, opts);\n  }\n\n  exports() {\n    return {\n      // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup(), which AppKit always awaits before exposing the plugin API\n      pool: this.pool!,\n      query: this.query.bind(this),\n      getOrmConfig: () => getLakebaseOrmConfig(this.config.pool),\n      getPgConfig: () => getLakebasePgConfig(this.config.pool),\n    };\n  }\n}\n\n/**\n * @internal\n */\nexport const lakebase = toPlugin(LakebasePlugin);\n"],"mappings":";;;;;;;;;;;;AAsBA,MAAM,SAAS,aAAa,WAAW;;;;;;;;;;;;;;;;;;AAmBvC,IAAa,iBAAb,cAAoC,OAA+B;;CAEjE,OAAO,WAAWA;CAGlB,AAAQ,OAAoB;;;;;;;;CAS5B,MAAM,QAAQ;EACZ,MAAM,aAAa,KAAK,OAAO;EAC/B,MAAM,OAAO,MAAM,yBAAyB,WAAW;AACvD,OAAK,OAAO,mBAAmB;GAAE,GAAG;GAAY;GAAM,CAAC;AACvD,SAAO,KAAK,4BAA4B;;;;;;;;;;;;;;;;;CAkB1C,MAAM,MACJ,MACA,QACyB;AAEzB,SAAO,KAAK,KAAM,MAAS,MAAM,OAAO;;;;;;;;;;;;;;;;CAiB1C,MAAc,qBACZ,MACA,QACoB;EAEpB,MAAM,SAAS,MAAM,KAAK,KAAM,SAAS;AACzC,MAAI;AACF,SAAM,OAAO,MAAM,kBAAkB;AAErC,WADe,MAAM,OAAO,MAAM,MAAM,OAAO,EACjC;YACN;AACR,OAAI;AACF,UAAM,OAAO,MAAM,WAAW;aACtB;AACR,WAAO,SAAS;;;;;;;;CAStB,wBAA8B;AAC5B,QAAM,uBAAuB;AAC7B,MAAI,KAAK,MAAM;AACb,UAAO,KAAK,wBAAwB;AACpC,QAAK,KAAK,KAAK,CAAC,OAAO,QAAQ;AAC7B,WAAO,MAAM,mCAAmC,IAAI;KACpD;AACF,QAAK,OAAO;;;;;;;;;;;;;;;;CAkBhB,AAAQ,QAAgE,EAAE;CAE1E,YAAY,QAAyB;AACnC,QAAM,OAAO;AACb,OAAK,SAAS;AACd,MAAI,OAAO,mBAAmB;AAC5B,QAAK,QAAQ,EAAE,OAAO,KAAK,eAAe,OAAO,kBAAkB,EAAE;AACrE,UAAO,KACL,gKACA,OAAO,kBAAkB,aAAa,MACvC;;;CAIL,AAAQ,eACN,KACA;EACA,MAAM,WAAW,IAAI,aAAa;AAClC,SAAO,WAAW;GAChB,aAAa,WACT,8PACA;GACJ,QAAQ,EAAE,OAAO;IACf,MAAM,EACH,QAAQ,CACR,SACC,6DACD;IACH,QAAQ,EACL,MAAM,EAAE,SAAS,CAAC,CAClB,UAAU,CACV,SAAS,iDAAiD;IAC9D,CAAC;GACF,aAAa;IACX,QAAQ,WAAW,SAAS;IAC5B,YAAY;IACb;GACD,SAAS,OAAO,MAAM,WAAW;AAQ/B,YAAQ,gBAAgB;AACxB,QAAI,UAAU;AACZ,uBAAkB,KAAK,KAAK;AAC5B,YAAO,KAAK,qBAAqB,KAAK,MAAM,KAAK,OAAO;;AAG1D,YADe,MAAM,KAAK,MAAM,KAAK,MAAM,KAAK,OAAO,EACzC;;GAEjB,CAAC;;CAGJ,gBAAuC;AACrC,SAAO,kBAAkB,KAAK,MAAM;;CAGtC,MAAM,iBACJ,MACA,MACA,QACkB;AAClB,SAAO,oBAAoB,KAAK,OAAO,MAAM,MAAM,OAAO;;CAG5D,QAAQ,MAAwD;AAC9D,SAAO,oBAAoB,KAAK,MAAM,KAAK,OAAO,KAAK;;CAGzD,UAAU;AACR,SAAO;GAEL,MAAM,KAAK;GACX,OAAO,KAAK,MAAM,KAAK,KAAK;GAC5B,oBAAoB,qBAAqB,KAAK,OAAO,KAAK;GAC1D,mBAAmB,oBAAoB,KAAK,OAAO,KAAK;GACzD;;;;;;AAOL,MAAa,WAAW,SAAS,eAAe"}
+{"version":3,"file":"lakebase.js","names":["manifest"],"sources":["../../../src/plugins/lakebase/lakebase.ts"],"sourcesContent":["import type { QueryResult, QueryResultRow } from \"pg\";\nimport type { AgentToolDefinition, ToolProvider } from \"shared\";\nimport { z } from \"zod\";\nimport {\n  createLakebasePool,\n  createLakebasePoolManager,\n  getLakebaseOrmConfig,\n  getLakebasePgConfig,\n  getUsernameWithApiLookup,\n  type LakebasePool,\n  type LakebasePoolManager,\n  RoutingPool,\n} from \"../../connectors/lakebase\";\nimport { getUserContext } from \"../../context/execution-context\";\nimport { buildToolkitEntries } from \"../../core/agent/build-toolkit\";\nimport {\n  defineTool,\n  executeFromRegistry,\n  toolsFromRegistry,\n} from \"../../core/agent/tools/define-tool\";\nimport { assertReadOnlySql } from \"../../core/agent/tools/sql-policy\";\nimport { createLogger } from \"../../logging/logger\";\nimport { Plugin, toPlugin } from \"../../plugin\";\nimport type { PluginManifest } from \"../../registry\";\nimport manifest from \"./manifest.json\";\nimport type { ILakebaseConfig } from \"./types\";\n\nconst logger = createLogger(\"lakebase\");\n\n/** Default pool settings for per-user OBO pools. */\nconst OBO_POOL_DEFAULTS = {\n  max: 3,\n  allowExitOnIdle: true,\n  idleTimeoutMillis: 30_000,\n};\n\n/**\n * AppKit plugin for Databricks Lakebase Autoscaling.\n *\n * Wraps `@databricks/lakebase` to provide a standard `pg.Pool` with automatic\n * OAuth token refresh, integrated with AppKit's logger and OpenTelemetry setup.\n *\n * Supports On-Behalf-Of (OBO) via `asUser(req)` — each user gets a separate\n * `pg.Pool` authenticated with their Databricks identity, enabling features\n * like Row-Level Security (RLS). Routing is handled transparently by\n * {@link RoutingPool}, which reads the execution context set by the base\n * class `asUser()`.\n *\n * @example\n * ```ts\n * import { createApp, lakebase, server } from \"@databricks/appkit\";\n *\n * const AppKit = await createApp({\n *   plugins: [server(), lakebase()],\n * });\n *\n * // Service principal query\n * const result = await AppKit.lakebase.query(\"SELECT * FROM users WHERE id = $1\", [userId]);\n *\n * // User-scoped query (per-user pool, RLS enforced)\n * const mine = await AppKit.lakebase.asUser(req).query(\"SELECT * FROM my_data\");\n * ```\n */\nexport class LakebasePlugin extends Plugin implements ToolProvider {\n  /** Plugin manifest declaring metadata and resource requirements */\n  static manifest = manifest as PluginManifest<\"lakebase\">;\n\n  protected declare config: ILakebaseConfig;\n  private pool: RoutingPool | null = null;\n  private oboPoolManager: LakebasePoolManager | null = null;\n\n  /**\n   * Initializes the Lakebase connection pool and OBO pool manager.\n   * Called automatically by AppKit during the plugin setup phase.\n   *\n   * Creates a {@link RoutingPool} that automatically routes queries to either\n   * the service-principal pool or a per-user pool based on the execution\n   * context (set by `Plugin.asUser(req)` via AsyncLocalStorage).\n   */\n  async setup() {\n    const poolConfig = this.config.pool;\n    const user = await getUsernameWithApiLookup(poolConfig);\n\n    const spPool = createLakebasePool({ ...poolConfig, user });\n    logger.info(\"Lakebase SP pool initialized\");\n\n    this.oboPoolManager = createLakebasePoolManager({\n      ...poolConfig,\n      ...OBO_POOL_DEFAULTS,\n    });\n    logger.info(\"Lakebase OBO pool manager initialized\");\n\n    const oboManager = this.oboPoolManager;\n    this.pool = new RoutingPool(spPool, (ctx) => {\n      if (!oboManager) throw new Error(\"OBO pool manager not initialized\");\n      // Lakebase OAuth roles use email as the postgres role when available\n      const userKey = ctx.userEmail ?? ctx.userId;\n      const isNew = !oboManager.hasPool(userKey);\n      const pool = oboManager.getPool(\n        userKey,\n        { workspaceClient: ctx.client, user: userKey },\n        ctx.tokenFingerprint,\n      );\n      if (isNew) {\n        logger.debug(\"Created OBO pool for user (total: %d)\", oboManager.size);\n      }\n      return pool;\n    });\n  }\n\n  /**\n   * Executes a parameterized SQL query against the Lakebase pool.\n   *\n   * When called inside `asUser(req)`, the query automatically routes to\n   * the per-user pool via {@link RoutingPool}.\n   *\n   * @param text - SQL query string, using `$1`, `$2`, ... placeholders\n   * @param values - Parameter values corresponding to placeholders\n   * @returns Query result with typed rows\n   *\n   * @example\n   * ```ts\n   * const result = await AppKit.lakebase.query<{ id: number; name: string }>(\n   *   \"SELECT id, name FROM users WHERE active = $1\",\n   *   [true],\n   * );\n   * ```\n   */\n  async query<T extends QueryResultRow = any>(\n    text: string,\n    values?: unknown[],\n  ): Promise<QueryResult<T>> {\n    // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup(), which AppKit always awaits before exposing the plugin API\n    return this.pool!.query<T>(text, values);\n  }\n\n  /**\n   * Execute a single statement inside a `BEGIN READ ONLY … ROLLBACK`\n   * transaction on a dedicated client.\n   *\n   * The three commands MUST share a connection — a naive\n   * `pool.query(\"BEGIN READ ONLY; <stmt>; ROLLBACK\")` batch cannot accept\n   * parameter values (PostgreSQL's Extended Query protocol rejects multi-\n   * statement prepared queries), which would silently break every\n   * parameterized query the agent tool issues.\n   *\n   * Returns the raw `rows` array for the user's statement. Side effects the\n   * statement may attempt (writes, writable-function side effects) are\n   * rejected by PostgreSQL under the read-only transaction posture.\n   */\n  private async runReadOnlyStatement(\n    text: string,\n    values?: unknown[],\n  ): Promise<unknown[]> {\n    // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup()\n    const client = await this.pool!.connect();\n    try {\n      await client.query(\"BEGIN READ ONLY\");\n      const result = await client.query(text, values);\n      return result.rows;\n    } finally {\n      try {\n        await client.query(\"ROLLBACK\");\n      } finally {\n        client.release();\n      }\n    }\n  }\n\n  /**\n   * Gracefully drains and closes all connection pools (SP + OBO).\n   * Called automatically by AppKit during shutdown.\n   */\n  abortActiveOperations(): void {\n    super.abortActiveOperations();\n    if (this.pool) {\n      logger.info(\"Closing Lakebase SP pool\");\n      this.pool.end().catch((err) => {\n        logger.error(\"Error closing Lakebase SP pool: %O\", err);\n      });\n      this.pool = null;\n    }\n    if (this.oboPoolManager) {\n      logger.info(\n        \"Closing all Lakebase OBO pools (%d)\",\n        this.oboPoolManager.size,\n      );\n      this.oboPoolManager.closeAll().catch((err) => {\n        logger.error(\"Error closing Lakebase OBO pools: %O\", err);\n      });\n      this.oboPoolManager = null;\n    }\n  }\n\n  /**\n   * Agent tool registry. Empty by default — the Lakebase plugin does NOT\n   * expose its SQL connection to LLM agents unless the developer explicitly\n   * opts in via `config.exposeAsAgentTool`. See {@link buildQueryTool}.\n   */\n  private tools: Record<string, ReturnType<typeof this.buildQueryTool>> = {};\n\n  constructor(config: ILakebaseConfig) {\n    super(config);\n    this.config = config;\n    if (config.exposeAsAgentTool) {\n      this.tools = { query: this.buildQueryTool(config.exposeAsAgentTool) };\n      logger.warn(\n        \"Lakebase agent tool is enabled (readOnly=%s). Every agent with access to this plugin can execute SQL against the Lakebase database as the requesting user's identity.\",\n        config.exposeAsAgentTool.readOnly !== false,\n      );\n    }\n  }\n\n  private buildQueryTool(\n    opt: NonNullable<ILakebaseConfig[\"exposeAsAgentTool\"]>,\n  ) {\n    const readOnly = opt.readOnly !== false;\n    return defineTool({\n      description: readOnly\n        ? \"Execute a read-only SQL query against the Lakebase PostgreSQL database. Only SELECT, WITH, SHOW, EXPLAIN, and DESCRIBE statements are accepted. Use $1, $2, etc. as placeholders and pass values separately.\"\n        : \"Execute a parameterized SQL statement against the Lakebase PostgreSQL database. Use $1, $2, etc. as placeholders and pass values separately. This tool can modify data; every invocation requires explicit human approval.\",\n      schema: z.object({\n        text: z\n          .string()\n          .describe(\n            \"SQL statement with $1, $2, ... placeholders for parameters\",\n          ),\n        values: z\n          .array(z.unknown())\n          .optional()\n          .describe(\"Parameter values corresponding to placeholders\"),\n      }),\n      annotations: {\n        effect: readOnly ? \"read\" : \"destructive\",\n        idempotent: false,\n        requiresUserContext: true,\n      },\n      execute: async (args, signal) => {\n        // Matches the files plugin pattern: the pg connection API\n        // doesn't accept AbortSignal in its current shape, so deeper\n        // mid-call cancellation needs a separate plumbing pass on the\n        // connector. This entry check still catches the common case —\n        // a tool dispatched after the user already cancelled the\n        // stream — and unwinds cleanly instead of running to\n        // completion against the SQL warehouse.\n        signal?.throwIfAborted();\n        if (readOnly) {\n          assertReadOnlySql(args.text);\n          return this.runReadOnlyStatement(args.text, args.values);\n        }\n        const result = await this.query(args.text, args.values);\n        return result.rows;\n      },\n    });\n  }\n\n  getAgentTools(): AgentToolDefinition[] {\n    return toolsFromRegistry(this.tools);\n  }\n\n  async executeAgentTool(\n    name: string,\n    args: unknown,\n    signal?: AbortSignal,\n  ): Promise<unknown> {\n    return executeFromRegistry(this.tools, name, args, signal);\n  }\n\n  toolkit(opts?: import(\"../../core/agent/types\").ToolkitOptions) {\n    return buildToolkitEntries(this.name, this.tools, opts);\n  }\n\n  /**\n   * Returns the pool config for the current execution context.\n   * Inside `asUser(req)`, returns user-scoped config; otherwise SP config.\n   */\n  private activePoolConfig() {\n    const ctx = getUserContext();\n    if (ctx) {\n      const user = ctx.userEmail ?? ctx.userId;\n      return { ...this.config.pool, workspaceClient: ctx.client, user };\n    }\n    return this.config.pool;\n  }\n\n  /**\n   * Returns the plugin's public API, accessible via `AppKit.lakebase`.\n   *\n   * - `pool` — The connection pool (routes to per-user pool when inside `asUser(req)`)\n   * - `query` — Convenience method for executing parameterized SQL queries\n   * - `getOrmConfig()` — Returns a config object compatible with Drizzle, TypeORM, Sequelize, etc.\n   *   Inside `asUser(req)`, returns user-scoped config.\n   * - `getPgConfig()` — Returns a `pg.PoolConfig` object for manual pool construction.\n   *   Inside `asUser(req)`, returns user-scoped config.\n   *\n   * Use `AppKit.lakebase.asUser(req)` to get the same API backed by a per-user pool.\n   */\n  exports() {\n    return {\n      // biome-ignore lint/style/noNonNullAssertion: pool is guaranteed non-null after setup(), which AppKit always awaits before exposing the plugin API\n      pool: this.pool! as LakebasePool,\n      query: this.query.bind(this),\n      getOrmConfig: () => getLakebaseOrmConfig(this.activePoolConfig()),\n      getPgConfig: () => getLakebasePgConfig(this.activePoolConfig()),\n    };\n  }\n}\n\n/**\n * @internal\n */\nexport const lakebase = toPlugin(LakebasePlugin);\n"],"mappings":";;;;;;;;;;;;;;;wBAaiE;AAcjE,MAAM,SAAS,aAAa,WAAW;;AAGvC,MAAM,oBAAoB;CACxB,KAAK;CACL,iBAAiB;CACjB,mBAAmB;CACpB;;;;;;;;;;;;;;;;;;;;;;;;;;;;AA6BD,IAAa,iBAAb,cAAoC,OAA+B;;CAEjE,OAAO,WAAWA;CAGlB,AAAQ,OAA2B;CACnC,AAAQ,iBAA6C;;;;;;;;;CAUrD,MAAM,QAAQ;EACZ,MAAM,aAAa,KAAK,OAAO;EAC/B,MAAM,OAAO,MAAM,yBAAyB,WAAW;EAEvD,MAAM,SAAS,mBAAmB;GAAE,GAAG;GAAY;GAAM,CAAC;AAC1D,SAAO,KAAK,+BAA+B;AAE3C,OAAK,iBAAiB,0BAA0B;GAC9C,GAAG;GACH,GAAG;GACJ,CAAC;AACF,SAAO,KAAK,wCAAwC;EAEpD,MAAM,aAAa,KAAK;AACxB,OAAK,OAAO,IAAI,YAAY,SAAS,QAAQ;AAC3C,OAAI,CAAC,WAAY,OAAM,IAAI,MAAM,mCAAmC;GAEpE,MAAM,UAAU,IAAI,aAAa,IAAI;GACrC,MAAM,QAAQ,CAAC,WAAW,QAAQ,QAAQ;GAC1C,MAAM,OAAO,WAAW,QACtB,SACA;IAAE,iBAAiB,IAAI;IAAQ,MAAM;IAAS,EAC9C,IAAI,iBACL;AACD,OAAI,MACF,QAAO,MAAM,yCAAyC,WAAW,KAAK;AAExE,UAAO;IACP;;;;;;;;;;;;;;;;;;;;CAqBJ,MAAM,MACJ,MACA,QACyB;AAEzB,SAAO,KAAK,KAAM,MAAS,MAAM,OAAO;;;;;;;;;;;;;;;;CAiB1C,MAAc,qBACZ,MACA,QACoB;EAEpB,MAAM,SAAS,MAAM,KAAK,KAAM,SAAS;AACzC,MAAI;AACF,SAAM,OAAO,MAAM,kBAAkB;AAErC,WADe,MAAM,OAAO,MAAM,MAAM,OAAO,EACjC;YACN;AACR,OAAI;AACF,UAAM,OAAO,MAAM,WAAW;aACtB;AACR,WAAO,SAAS;;;;;;;;CAStB,wBAA8B;AAC5B,QAAM,uBAAuB;AAC7B,MAAI,KAAK,MAAM;AACb,UAAO,KAAK,2BAA2B;AACvC,QAAK,KAAK,KAAK,CAAC,OAAO,QAAQ;AAC7B,WAAO,MAAM,sCAAsC,IAAI;KACvD;AACF,QAAK,OAAO;;AAEd,MAAI,KAAK,gBAAgB;AACvB,UAAO,KACL,uCACA,KAAK,eAAe,KACrB;AACD,QAAK,eAAe,UAAU,CAAC,OAAO,QAAQ;AAC5C,WAAO,MAAM,wCAAwC,IAAI;KACzD;AACF,QAAK,iBAAiB;;;;;;;;CAS1B,AAAQ,QAAgE,EAAE;CAE1E,YAAY,QAAyB;AACnC,QAAM,OAAO;AACb,OAAK,SAAS;AACd,MAAI,OAAO,mBAAmB;AAC5B,QAAK,QAAQ,EAAE,OAAO,KAAK,eAAe,OAAO,kBAAkB,EAAE;AACrE,UAAO,KACL,yKACA,OAAO,kBAAkB,aAAa,MACvC;;;CAIL,AAAQ,eACN,KACA;EACA,MAAM,WAAW,IAAI,aAAa;AAClC,SAAO,WAAW;GAChB,aAAa,WACT,iNACA;GACJ,QAAQ,EAAE,OAAO;IACf,MAAM,EACH,QAAQ,CACR,SACC,6DACD;IACH,QAAQ,EACL,MAAM,EAAE,SAAS,CAAC,CAClB,UAAU,CACV,SAAS,iDAAiD;IAC9D,CAAC;GACF,aAAa;IACX,QAAQ,WAAW,SAAS;IAC5B,YAAY;IACZ,qBAAqB;IACtB;GACD,SAAS,OAAO,MAAM,WAAW;AAQ/B,YAAQ,gBAAgB;AACxB,QAAI,UAAU;AACZ,uBAAkB,KAAK,KAAK;AAC5B,YAAO,KAAK,qBAAqB,KAAK,MAAM,KAAK,OAAO;;AAG1D,YADe,MAAM,KAAK,MAAM,KAAK,MAAM,KAAK,OAAO,EACzC;;GAEjB,CAAC;;CAGJ,gBAAuC;AACrC,SAAO,kBAAkB,KAAK,MAAM;;CAGtC,MAAM,iBACJ,MACA,MACA,QACkB;AAClB,SAAO,oBAAoB,KAAK,OAAO,MAAM,MAAM,OAAO;;CAG5D,QAAQ,MAAwD;AAC9D,SAAO,oBAAoB,KAAK,MAAM,KAAK,OAAO,KAAK;;;;;;CAOzD,AAAQ,mBAAmB;EACzB,MAAM,MAAM,gBAAgB;AAC5B,MAAI,KAAK;GACP,MAAM,OAAO,IAAI,aAAa,IAAI;AAClC,UAAO;IAAE,GAAG,KAAK,OAAO;IAAM,iBAAiB,IAAI;IAAQ;IAAM;;AAEnE,SAAO,KAAK,OAAO;;;;;;;;;;;;;;CAerB,UAAU;AACR,SAAO;GAEL,MAAM,KAAK;GACX,OAAO,KAAK,MAAM,KAAK,KAAK;GAC5B,oBAAoB,qBAAqB,KAAK,kBAAkB,CAAC;GACjE,mBAAmB,oBAAoB,KAAK,kBAAkB,CAAC;GAChE;;;;;;AAOL,MAAa,WAAW,SAAS,eAAe"}

package/dist/plugins/serving/serving.js

@@ -1,9 +1,9 @@
 import { createLogger } from "../../logging/logger.js";
 import { getWorkspaceClient } from "../../context/execution-context.js";
 import { init_context } from "../../context/index.js";
+import { Plugin } from "../../plugin/plugin.js";
 import { ResourceType } from "../../registry/types.generated.js";
 import "../../registry/index.js";
-import { Plugin } from "../../plugin/plugin.js";
 import { toPlugin } from "../../plugin/to-plugin.js";
 import "../../plugin/index.js";
 import "../../logging/index.js";

package/docs/api/appkit/Function.createLakebasePoolManager.md

@@ -0,0 +1,36 @@
+# Function: createLakebasePoolManager()
+
+```ts
+function createLakebasePoolManager(baseConfig?: Partial<LakebasePoolConfig>): LakebasePoolManager;
+
+```
+
+Create a pool manager that maintains per-key Lakebase connection pools.
+
+Each pool is created via `createLakebasePool` with the base config merged with per-pool overrides (e.g. a user's `workspaceClient` and `user`).
+
+A periodic cleanup removes empty Pool objects (where all connections have been closed by pg's built-in `idleTimeoutMillis`) from the internal Map.
+
+## Parameters[​](#parameters "Direct link to Parameters")
+
+| Parameter     | Type                                                                                       |
+| ------------- | ------------------------------------------------------------------------------------------ |
+| `baseConfig?` | `Partial`<[`LakebasePoolConfig`](./docs/api/appkit/Interface.LakebasePoolConfig.md)> |
+
+## Returns[​](#returns "Direct link to Returns")
+
+[`LakebasePoolManager`](./docs/api/appkit/Interface.LakebasePoolManager.md)
+
+## Example[​](#example "Direct link to Example")
+
+```typescript
+const poolManager = createLakebasePoolManager();
+
+// In a route handler:
+const userPool = poolManager.getPool(userName, {
+  workspaceClient: new WorkspaceClient({ token: userToken, host, authType: "pat" }),
+  user: userName,
+});
+const result = await userPool.query("SELECT * FROM products");
+
+```

package/docs/api/appkit/Interface.LakebasePool.md

@@ -0,0 +1,84 @@
+# Interface: LakebasePool
+
+Subset of `pg.Pool` exposed by the Lakebase plugin.
+
+RoutingPool does not extend EventEmitter — event listener methods like `on('error', ...)` are not available. Use `query()`, `connect()`, and `end()` for all pool operations.
+
+## Properties[​](#properties "Direct link to Properties")
+
+### idleCount[​](#idlecount "Direct link to idleCount")
+
+```ts
+readonly idleCount: number;
+
+```
+
+***
+
+### totalCount[​](#totalcount "Direct link to totalCount")
+
+```ts
+readonly totalCount: number;
+
+```
+
+***
+
+### waitingCount[​](#waitingcount "Direct link to waitingCount")
+
+```ts
+readonly waitingCount: number;
+
+```
+
+## Methods[​](#methods "Direct link to Methods")
+
+### connect()[​](#connect "Direct link to connect()")
+
+```ts
+connect(): Promise<PoolClient>;
+
+```
+
+#### Returns[​](#returns "Direct link to Returns")
+
+`Promise`<`PoolClient`>
+
+***
+
+### end()[​](#end "Direct link to end()")
+
+```ts
+end(): Promise<void>;
+
+```
+
+#### Returns[​](#returns-1 "Direct link to Returns")
+
+`Promise`<`void`>
+
+***
+
+### query()[​](#query "Direct link to query()")
+
+```ts
+query<T>(text: string, values?: unknown[]): Promise<QueryResult<T>>;
+
+```
+
+#### Type Parameters[​](#type-parameters "Direct link to Type Parameters")
+
+| Type Parameter                 | Default type |
+| ------------------------------ | ------------ |
+| `T` *extends* `QueryResultRow` | `any`        |
+
+#### Parameters[​](#parameters "Direct link to Parameters")
+
+| Parameter | Type         |
+| --------- | ------------ |
+| `text`    | `string`     |
+| `values?` | `unknown`\[] |
+
+#### Returns[​](#returns-2 "Direct link to Returns")
+
+`Promise`<`QueryResult`<`T`>>

package/docs/api/appkit/Interface.LakebasePoolManager.md

@@ -0,0 +1,101 @@
+# Interface: LakebasePoolManager
+
+Manages multiple Lakebase connection pools keyed by an identifier (e.g. userId).
+
+Used for On-Behalf-Of (OBO) scenarios where each user needs their own pool with their own OAuth token refresh, enabling features like Row-Level Security.
+
+## Properties[​](#properties "Direct link to Properties")
+
+### size[​](#size "Direct link to size")
+
+```ts
+readonly size: number;
+
+```
+
+Number of active pools.
+
+## Methods[​](#methods "Direct link to Methods")
+
+### closeAll()[​](#closeall "Direct link to closeAll()")
+
+```ts
+closeAll(): Promise<void>;
+
+```
+
+Close all managed pools and stop cleanup (for graceful shutdown).
+
+#### Returns[​](#returns "Direct link to Returns")
+
+`Promise`<`void`>
+
+***
+
+### closePool()[​](#closepool "Direct link to closePool()")
+
+```ts
+closePool(key: string): Promise<void>;
+
+```
+
+Close and remove a specific pool.
+
+#### Parameters[​](#parameters "Direct link to Parameters")
+
+| Parameter | Type     |
+| --------- | -------- |
+| `key`     | `string` |
+
+#### Returns[​](#returns-1 "Direct link to Returns")
+
+`Promise`<`void`>
+
+***
+
+### getPool()[​](#getpool "Direct link to getPool()")
+
+```ts
+getPool(
+   key: string, 
+   perPoolConfig: Partial<LakebasePoolConfig>, 
+   tokenFingerprint?: string): Pool;
+
+```
+
+Get an existing pool or create a new one for the given key. When creating, merges `perPoolConfig` with the base config passed to the factory.
+
+If `tokenFingerprint` is provided and differs from the cached pool's fingerprint, the stale pool is closed and a fresh one is created with the new config (including the updated `workspaceClient`).
+
+#### Parameters[​](#parameters-1 "Direct link to Parameters")
+
+| Parameter           | Type                                                                                       |
+| ------------------- | ------------------------------------------------------------------------------------------ |
+| `key`               | `string`                                                                                   |
+| `perPoolConfig`     | `Partial`<[`LakebasePoolConfig`](./docs/api/appkit/Interface.LakebasePoolConfig.md)> |
+| `tokenFingerprint?` | `string`                                                                                   |
+
+#### Returns[​](#returns-2 "Direct link to Returns")
+
+`Pool`
+
+***
+
+### hasPool()[​](#haspool "Direct link to hasPool()")
+
+```ts
+hasPool(key: string): boolean;
+
+```
+
+Check whether a pool exists for the given key.
+
+#### Parameters[​](#parameters-2 "Direct link to Parameters")
+
+| Parameter | Type     |
+| --------- | -------- |
+| `key`     | `string` |
+
+#### Returns[​](#returns-3 "Direct link to Returns")
+
+`boolean`

package/docs/api/appkit.md

@@ -52,7 +52,9 @@ Documentation merge entry for Typedoc — combines the stable `@databricks/appki
 | [JobAPI](./docs/api/appkit/Interface.JobAPI.md)                                                       | User-facing API for a single configured job.                                                                                                                                                                                                                                                                                                                                                                                               |
 | [JobConfig](./docs/api/appkit/Interface.JobConfig.md)                                                 | Per-job configuration options.                                                                                                                                                                                                                                                                                                                                                                                                             |
 | [JobsConnectorConfig](./docs/api/appkit/Interface.JobsConnectorConfig.md)                             | -                                                                                                                                                                                                                                                                                                                                                                                                                                          |
+| [LakebasePool](./docs/api/appkit/Interface.LakebasePool.md)                                           | Subset of `pg.Pool` exposed by the Lakebase plugin.                                                                                                                                                                                                                                                                                                                                                                                        |
 | [LakebasePoolConfig](./docs/api/appkit/Interface.LakebasePoolConfig.md)                               | Configuration for creating a Lakebase connection pool                                                                                                                                                                                                                                                                                                                                                                                      |
+| [LakebasePoolManager](./docs/api/appkit/Interface.LakebasePoolManager.md)                             | Manages multiple Lakebase connection pools keyed by an identifier (e.g. userId).                                                                                                                                                                                                                                                                                                                                                           |
 | [McpConnectAllResult](./docs/api/appkit/Interface.McpConnectAllResult.md)                             | Per-endpoint outcome of [AppKitMcpClient.connectAll](./docs/api/appkit/Class.AppKitMcpClient.md#connectall). Callers (the agents plugin in particular) use the split to warn at startup when some MCP servers are unreachable without aborting boot for the rest.                                                                                                                                                                    |
 | [Message](./docs/api/appkit/Interface.Message.md)                                                     | -                                                                                                                                                                                                                                                                                                                                                                                                                                          |
 | [PluginManifest](./docs/api/appkit/Interface.PluginManifest.md)                                       | Plugin manifest that declares metadata and resource requirements. Attached to plugin classes as a static property. Extends the shared PluginManifest with strict resource types.                                                                                                                                                                                                                                                           |
@@ -124,6 +126,7 @@ Documentation merge entry for Typedoc — combines the stable `@databricks/appki
 | [createAgent](./docs/api/appkit/Function.createAgent.md)                               | Pure factory for agent definitions. Returns the passed-in definition after cycle-detecting the sub-agent graph. Accepts the full `AgentDefinition` shape and is safe to call at module top-level.                                                                                     |
 | [createApp](./docs/api/appkit/Function.createApp.md)                                   | Bootstraps AppKit with the provided configuration.                                                                                                                                                                                                                                    |
 | [createLakebasePool](./docs/api/appkit/Function.createLakebasePool.md)                 | Create a Lakebase pool with appkit's logger integration. Telemetry automatically uses appkit's OpenTelemetry configuration via global registry.                                                                                                                                       |
+| [createLakebasePoolManager](./docs/api/appkit/Function.createLakebasePoolManager.md)   | Create a pool manager that maintains per-key Lakebase connection pools.                                                                                                                                                                                                               |
 | [defineTool](./docs/api/appkit/Function.defineTool.md)                                 | Defines a single tool entry for a plugin's internal registry.                                                                                                                                                                                                                         |
 | [executeFromRegistry](./docs/api/appkit/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](./docs/api/appkit/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.                                                         |

package/docs/development/llm-guide.md

@@ -24,7 +24,6 @@ This guide is designed to work even when you *do not* have access to the AppKit
 
 * **Do not invent APIs**. If unsure, stick to the patterns shown in the documentation and only use documented exports from `@databricks/appkit` and `@databricks/appkit-ui`.
 * **`createApp()` is async**. Prefer **top-level `await createApp(...)`**. If you can't, use `void createApp(...)` and do not ignore promise rejection.
-* **Always memoize query parameters** passed to `useAnalyticsQuery` / charts to avoid refetch loops.
 * **Always handle loading/error/empty states** in UI (use `Skeleton`, error text, empty state).
 * **Always use `sql.*` helpers** for query parameters (do not pass raw strings/numbers unless the query expects none).
 * **Never construct SQL strings dynamically**. Use parameterized queries with `:paramName`.

package/docs/plugins/execution-context.md

@@ -51,6 +51,12 @@ The `plugin.execute` span created by the execution interceptor chain includes th
 
 These attributes are automatically added when your plugin uses `execute()` or `executeStream()`. All built-in plugins use these methods for their OBO operations. Custom plugins should do the same to get automatic telemetry instrumentation.
 
+## Lakebase per-user connections[​](#lakebase-per-user-connections "Direct link to Lakebase per-user connections")
+
+The Lakebase plugin uses a different mechanism for `asUser(req)`: instead of swapping the `WorkspaceClient` via AsyncLocalStorage, it creates a **separate `pg.Pool` per user**, each with its own OAuth token refresh. This is necessary because PostgreSQL connections are authenticated at connection time — the pool itself is the authentication boundary.
+
+See [Lakebase plugin — per-user connections](./docs/plugins/lakebase.md#on-behalf-of-obo--per-user-connections) for details.
+
 ## Development mode behavior[​](#development-mode-behavior "Direct link to Development mode behavior")
 
 In local development (`NODE_ENV=development`), if `asUser(req)` is called without a user token, it logs a warning and skips user impersonation — the operation runs with the default credentials configured for the app instead. The telemetry span will show `execution.context: "service"` with `execution.obo_dev_fallback: true` to distinguish these from regular service principal calls.