@dbx-tools/appkit 0.1.9

package/src/pgaddress.ts

@@ -0,0 +1,194 @@
+/**
+ * Flexible address parser for Lakebase Postgres connection inputs.
+ *
+ * Accepts whatever shape a user is likely to paste into
+ * `LAKEBASE_ENDPOINT` (or the matching config field) and extracts
+ * every recognizable piece. Whatever it can't recover is left for the
+ * Lakebase resolver to discover.
+ *
+ * Recognized formats:
+ *
+ * - **Postgres URI** -
+ *   `postgresql://user@host:port/db?sslmode=require` (also `postgres://`).
+ *   Yields `user`, `host`, `port`, `database`, `sslMode`.
+ *
+ * - **Canonical endpoint resource path** -
+ *   `projects/{p}/branches/{b}/endpoints/{e}` -
+ *   yields `project`, `branch`, `endpointId`, and the original string as
+ *   `endpoint` (already in lakebase's expected form).
+ *
+ * - **Database resource path** -
+ *   `projects/{p}/branches/{b}/databases/{d}` -
+ *   yields `project`, `branch`, and `databaseResourceId` (the UC
+ *   resource leaf, not `PGDATABASE`; the resolver looks up the real
+ *   Postgres name via REST).
+ *
+ * - **Branch resource path** -
+ *   `projects/{p}/branches/{b}` - yields `project`, `branch`.
+ *
+ * - **Project resource path** -
+ *   `projects/{p}` - yields `project`.
+ *
+ * - **Bare hostname** -
+ *   `ep-steep-forest-e199v43w.database.eastus2.azuredatabricks.net` -
+ *   yields `host` only; the resolver reverse-looks up the owning
+ *   endpoint to recover the resource path.
+ *
+ * - **Bare project id** -
+ *   `dbx-tools-demo` (1-63 chars, lowercase letters/digits/hyphens) -
+ *   yields `project`.
+ *
+ * Returns an empty object for inputs it doesn't recognize.
+ */
+
+/** Postgres TLS mode passed through to `pg`. */
+export type SslMode = "require" | "disable" | "prefer";
+
+/**
+ * Optional Lakebase Postgres connection fields shared by parsed addresses,
+ * resolver/env inputs, and resolved connections.
+ */
+export interface LakebaseConnectionInputs {
+  /** Lakebase project id. */
+  project?: string;
+  /** Branch id within the project. */
+  branch?: string;
+  /** Canonical endpoint resource path (`projects/.../endpoints/...`). */
+  endpoint?: string;
+  /** Postgres database name (`PGDATABASE`). */
+  database?: string;
+  /** Postgres hostname (`PGHOST`). */
+  host?: string;
+  /** Postgres port (`PGPORT`). */
+  port?: number;
+  /** Postgres TLS mode (`PGSSLMODE`). */
+  sslMode?: SslMode;
+}
+
+/** Pieces recovered from parsing a single address or resource-path input. */
+export interface ParsedAddress extends LakebaseConnectionInputs {
+  /** Endpoint leaf id (last segment of an endpoint resource path). */
+  endpointId?: string;
+  /**
+   * Database resource id leaf from a `.../databases/{id}` path. Not the
+   * Postgres database name.
+   */
+  databaseResourceId?: string;
+  /** Postgres user (URI-decoded if encoded). */
+  user?: string;
+}
+
+const URL_SCHEME_RE = /^(postgres|postgresql):\/\//i;
+const PROJECT_ID_RE = /^[a-z][a-z0-9-]{0,61}[a-z0-9]$|^[a-z]$/;
+const HOSTNAME_HINT_RE = /^[a-z0-9][a-z0-9-]*(\.[a-z0-9][a-z0-9-]*)+$/i;
+
+/**
+ * Parse a Lakebase connection input into whatever pieces it carries.
+ * See module docstring for the supported formats. Returns `{}` for
+ * `undefined`, empty strings, and unrecognized inputs.
+ */
+export function parseAddress(input: string | undefined | null): ParsedAddress {
+  if (!input) return {};
+  const s = input.trim();
+  if (!s) return {};
+
+  if (URL_SCHEME_RE.test(s)) return parseUri(s);
+  if (s.startsWith("projects/")) return parseResourcePathSegments(s);
+  // Resource ids never contain dots; a dotted input must be a hostname.
+  if (HOSTNAME_HINT_RE.test(s) && s.includes(".")) return { host: s };
+  if (PROJECT_ID_RE.test(s)) return { project: s };
+  return {};
+}
+
+/**
+ * Parse a Lakebase `projects/...` resource path. Returns `{}` when the
+ * input is not a resource path (so bare branch ids are not mistaken for
+ * project ids).
+ */
+export function parseResourcePath(input: string | undefined | null): ParsedAddress {
+  if (!input) return {};
+  const s = input.trim();
+  if (!s.startsWith("projects/")) return {};
+  return parseResourcePathSegments(s);
+}
+
+function parseUri(s: string): ParsedAddress {
+  let url: URL;
+  try {
+    url = new URL(s);
+  } catch {
+    return {};
+  }
+  const result: ParsedAddress = {};
+  if (url.hostname) result.host = url.hostname;
+  if (url.port) {
+    const port = Number.parseInt(url.port, 10);
+    if (!Number.isNaN(port)) result.port = port;
+  }
+  if (url.username) {
+    try {
+      result.user = decodeURIComponent(url.username);
+    } catch {
+      result.user = url.username;
+    }
+  }
+  const db = url.pathname.replace(/^\//, "");
+  if (db) result.database = decodeURIComponent(db);
+  const sslmodeRaw = url.searchParams.get("sslmode") ?? url.searchParams.get("sslMode");
+  const sslmode = sslmodeRaw?.toLowerCase();
+  if (sslmode === "require" || sslmode === "disable" || sslmode === "prefer") {
+    result.sslMode = sslmode;
+  }
+  return result;
+}
+
+function parseResourcePathSegments(s: string): ParsedAddress {
+  const parts = s.split("/");
+  if (parts[0] !== "projects" || parts.length < 2) {
+    return {};
+  }
+
+  const project = parts[1];
+  if (!project) {
+    return {};
+  }
+
+  if (parts.length === 2) {
+    return { project };
+  }
+
+  if (parts.length === 4 && parts[2] === "branches" && parts[3]) {
+    return { project, branch: parts[3] };
+  }
+
+  if (
+    parts.length === 6 &&
+    parts[2] === "branches" &&
+    parts[4] === "endpoints" &&
+    parts[3] &&
+    parts[5]
+  ) {
+    return {
+      project,
+      branch: parts[3],
+      endpointId: parts[5],
+      endpoint: s,
+    };
+  }
+
+  if (
+    parts.length === 6 &&
+    parts[2] === "branches" &&
+    parts[4] === "databases" &&
+    parts[3] &&
+    parts[5]
+  ) {
+    return {
+      project,
+      branch: parts[3],
+      databaseResourceId: parts[5],
+    };
+  }
+
+  return {};
+}

package/src/plugin.ts

@@ -0,0 +1,129 @@
+/**
+ * AppKit plugin lookup: typed access to sibling plugins registered on the
+ * AppKit plugin context (`this.context` on any class that extends `Plugin`).
+ *
+ * Why these live here instead of in `@databricks/appkit`: AppKit exposes
+ * `this.context.getPlugins()`, which returns `ReadonlyMap<string, BasePlugin>`,
+ * but provides no typed lookup helper. Every caller ends up writing the same
+ * `as InstanceType<ReturnType<typeof someFactory>["plugin"]>` cast. These
+ * wrappers absorb that boilerplate.
+ *
+ * API shape: pass the plugin's factory (`lakebase`, `serving`, `genie`, or any
+ * `toPlugin(...)` result) directly. TypeScript infers both the instance type
+ * (so `.exports()` resolves) and the registered name (so the runtime lookup
+ * works) from that single value. No `<T>` annotation or string literal needed
+ * at the call site.
+ *
+ * Generic AppKit runtime (execution context, `WorkspaceClientLike`) lives in
+ * `./appkit`; this module is plugin-lookup only.
+ */
+
+import { type NameLike } from "@dbx-tools/shared-core";
+
+/**
+ * Minimal structural shape of `this.context`. We mirror only the method we
+ * touch instead of depending on AppKit's `PluginContext` type, which is not
+ * part of the package's `exports` map and therefore cannot be imported. Any
+ * compatible object (real `PluginContext`, mocks, tests) satisfies this shape.
+ */
+export interface PluginContextLike {
+  getPlugins(): ReadonlyMap<string, unknown>;
+}
+
+type PluginData = {
+  plugin: abstract new (...args: never[]) => unknown;
+  name: string;
+};
+
+/**
+ * Structural shape of an AppKit plugin factory (the result of
+ * `toPlugin(SomePluginClass)`). Calling it returns a `PluginData` tuple whose
+ * `plugin` field is the *class constructor* and whose `name` field carries the
+ * registered plugin name as a literal string.
+ *
+ * Defined structurally so we don't pull `@databricks/appkit` into this as a
+ * type dependency for the bound. Any function returning the same shape (e.g.
+ * `lakebase`, `serving`, `genie`, or a user-defined `toPlugin(MyPlugin)`)
+ * satisfies it.
+ */
+type PluginDataFactory = (...args: never[]) => PluginData;
+
+/**
+ * Maps a plugin factory back to the *instance* type of its plugin class.
+ * Mirrors the inline pattern users would otherwise write:
+ * `InstanceType<ReturnType<typeof factory>["plugin"]>`.
+ */
+type PluginInstanceOf<F extends PluginDataFactory> = InstanceType<ReturnType<F>["plugin"]>;
+
+/**
+ * Registry name returned by `factory().name`, keyed by the factory function.
+ * Typical AppKit factories return stable metadata; caching avoids invoking
+ * `factory()` on every sibling lookup (which would allocate a fresh descriptor
+ * tuple each time).
+ */
+const dataCache = new WeakMap<PluginDataFactory, PluginData>();
+
+/**
+ * Returns the static `{ plugin, name }` descriptor for an AppKit plugin
+ * factory, caching per factory so repeated lookups do not allocate.
+ */
+export function data<F extends PluginDataFactory, D extends ReturnType<F>>(factory: F): D {
+  const cached = dataCache.get(factory);
+  if (cached !== undefined) {
+    return cached as D;
+  }
+  const result = factory();
+  dataCache.set(factory, result);
+  return result as D;
+}
+
+/**
+ * Look up a sibling plugin instance from the AppKit plugin context, keyed off
+ * the factory's registered name and typed via its plugin class.
+ *
+ * Returns `undefined` when the context is missing or the plugin is not
+ * registered. For required siblings prefer {@link require}.
+ *
+ * @example
+ * import { lakebase } from "@databricks/appkit";
+ * import { plugin } from "@dbx-tools/appkit";
+ *
+ * const lake = plugin.instance(this.context, lakebase);
+ * //    ^^ inferred as LakebasePlugin | undefined
+ * lake?.exports().pool;
+ */
+export function instance<F extends PluginDataFactory>(
+  ctx: PluginContextLike | undefined,
+  factory: F,
+): PluginInstanceOf<F> | undefined {
+  if (!ctx) return undefined;
+  const name = data(factory).name;
+  return ctx.getPlugins().get(name) as PluginInstanceOf<F> | undefined;
+}
+
+/**
+ * Like {@link instance} but throws when the plugin is not registered. Use for
+ * siblings whose absence is a wiring bug rather than a runtime condition (e.g.
+ * requiring `lakebase` when the caller has `storage` / `memory` enabled).
+ *
+ * `caller` is prepended to the error message so cross-plugin failures are easy
+ * to attribute in logs.
+ *
+ * @example
+ * import { lakebase } from "@databricks/appkit";
+ * import { plugin } from "@dbx-tools/appkit";
+ *
+ * const pool = plugin.require(this.context, lakebase, "mastra").exports().pool;
+ */
+export function require<F extends PluginDataFactory>(
+  ctx: PluginContextLike | undefined,
+  factory: F,
+  caller?: NameLike | string,
+): PluginInstanceOf<F> {
+  const found = instance(ctx, factory);
+  if (found) return found;
+  const prefix =
+    typeof caller === "string" ? `${caller}: ` : caller?.name ? `${caller.name}: ` : "";
+  const registeredName = data(factory).name;
+  throw new Error(`${prefix}required plugin not registered: ${registeredName}`);
+}

package/src/provision.ts

@@ -0,0 +1,111 @@
+/**
+ * Lakebase cache-schema grant fix-up.
+ *
+ * AppKit's persistent cache (`CacheManager` -> `PersistentStorage`) uses a
+ * schema `appkit` and a table `appkit.appkit_cache_entries`. When the connecting
+ * Databricks identity lacks privileges on that schema, the cache migration
+ * throws, and because AppKit is typically run with `cache.strictPersistence:
+ * true`, the cache is silently switched to a disabled in-memory stub - so every
+ * persistent read (chart long-poll, history, etc.) misses and 404s.
+ *
+ * This grants the connecting role full rights on the cache schema, but only when
+ * the schema ALREADY EXISTS - it never creates the schema itself. Run from a
+ * LOCAL identity that owns the schema (or holds grant option on it); it is
+ * skipped inside a Databricks App, where the app SP cannot grant and its
+ * Postgres role does not exist until its first connection.
+ *
+ * Must run AFTER {@link applyLakebaseToEnv} has written the resolved connection
+ * to `process.env` (so `createLakebasePool` picks up host / database / endpoint)
+ * and BEFORE `createApp` initializes the cache.
+ */
+
+import { error, type log } from "@dbx-tools/shared-core";
+import { createLakebasePool, getWorkspaceClient } from "@databricks/appkit";
+
+import { isAppEnv } from "./databricks";
+
+/** AppKit persistent-cache schema (see AppKit's `PersistentStorage`). */
+const CACHE_SCHEMA = "appkit";
+
+/**
+ * Quote a Postgres identifier: wrap in double quotes and double any embedded
+ * quote. Needed because Lakebase role names are often emails (`user@host`) that
+ * must be quoted to be valid identifiers.
+ */
+function quoteIdent(ident: string): string {
+  return `"${ident.replace(/"/g, '""')}"`;
+}
+
+/**
+ * Idempotent grants that make the (already-existing) AppKit cache schema fully
+ * usable by `role`. The `ALTER DEFAULT PRIVILEGES` lines cover the cache table
+ * whenever the schema owner creates it later.
+ */
+function cacheGrantStatements(role: string): readonly string[] {
+  const schema = quoteIdent(CACHE_SCHEMA);
+  const target = quoteIdent(role);
+  return [
+    `GRANT USAGE, CREATE ON SCHEMA ${schema} TO ${target}`,
+    `GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA ${schema} TO ${target}`,
+    `GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA ${schema} TO ${target}`,
+    `ALTER DEFAULT PRIVILEGES IN SCHEMA ${schema} GRANT ALL ON TABLES TO ${target}`,
+    `ALTER DEFAULT PRIVILEGES IN SCHEMA ${schema} GRANT ALL ON SEQUENCES TO ${target}`,
+  ];
+}
+
+/**
+ * Grant `role` rights on the AppKit cache schema, but only when that schema
+ * already exists.
+ *
+ * No-ops inside a Databricks App env, when `role` is undefined, or when the
+ * schema is absent. Best-effort otherwise: any failure (e.g. the local identity
+ * doesn't own the schema) is logged and swallowed so it never blocks startup -
+ * a disabled cache is degraded, not fatal.
+ *
+ * @param role - Postgres role to grant to and connect as (the resolved
+ *   workspace-client identity); skips when undefined.
+ */
+export async function provisionCacheSchema(
+  logger: log.Logger,
+  role: string | undefined,
+): Promise<void> {
+  if (isAppEnv()) {
+    logger.debug("autopg: skip cache provisioning (inside a Databricks App)");
+    return;
+  }
+  if (!role) {
+    logger.warn("autopg: skip cache provisioning (could not resolve connecting role)");
+    return;
+  }
+  // Pass `user` explicitly: autopg resolves the connection target
+  // (host/database/endpoint) but not the identity, so `createLakebasePool`'s
+  // synchronous username lookup would otherwise throw. `getWorkspaceClient({})`
+  // returns a fresh default-auth client - literally `new WorkspaceClient({})`,
+  // but built against the SDK version AppKit's `createLakebasePool` expects
+  // (importing the SDK class directly here pulls a newer, type-incompatible copy).
+  const pool = createLakebasePool({
+    user: role,
+    workspaceClient: getWorkspaceClient({}),
+  });
+  try {
+    const found = await pool.query(
+      "SELECT 1 FROM information_schema.schemata WHERE schema_name = $1",
+      [CACHE_SCHEMA],
+    );
+    if (found.rowCount === 0) {
+      logger.debug("autopg: skip cache provisioning (schema absent)", { schema: CACHE_SCHEMA });
+      return;
+    }
+    for (const sql of cacheGrantStatements(role)) {
+      await pool.query(sql);
+    }
+    logger.info("autopg: granted cache schema", { schema: CACHE_SCHEMA, role });
+  } catch (err) {
+    logger.warn("autopg: cache provisioning failed (continuing)", {
+      schema: CACHE_SCHEMA,
+      error: error.errorMessage(err),
+    });
+  } finally {
+    await pool.end().catch(() => {});
+  }
+}

package/tsconfig.json

@@ -0,0 +1,41 @@
+// ~~ Generated by projen. To modify, edit .projenrc.js and run "pnpm exec projen".
+{
+  "compilerOptions": {
+    "rootDir": "src",
+    "outDir": "lib",
+    "alwaysStrict": true,
+    "declaration": true,
+    "esModuleInterop": true,
+    "experimentalDecorators": true,
+    "inlineSourceMap": true,
+    "inlineSources": true,
+    "lib": [
+      "ES2022"
+    ],
+    "module": "ESNext",
+    "noEmitOnError": false,
+    "noFallthroughCasesInSwitch": true,
+    "noImplicitAny": true,
+    "noImplicitReturns": true,
+    "noImplicitThis": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "resolveJsonModule": true,
+    "strict": true,
+    "strictNullChecks": true,
+    "strictPropertyInitialization": true,
+    "stripInternal": true,
+    "target": "ES2022",
+    "types": [
+      "node"
+    ],
+    "moduleResolution": "bundler",
+    "skipLibCheck": true
+  },
+  "include": [
+    "src/**/*.ts"
+  ],
+  "exclude": [
+    "node_modules"
+  ]
+}