@dbx-tools/appkit 0.3.29 → 0.3.30

package/src/pgaddress.ts

@@ -43,27 +43,33 @@
  * @module
  */
 
+/** Postgres TLS modes accepted by {@link SslMode}, in `PGSSLMODE` spelling. */
+export const SSL_MODES = ["require", "disable", "prefer"] as const;
+
 /** Postgres TLS mode passed through to `pg`. */
-export type SslMode = "require" | "disable" | "prefer";
+export type SslMode = (typeof SSL_MODES)[number];
 
 /**
  * Optional Lakebase Postgres connection fields shared by parsed addresses,
  * resolver/env inputs, and resolved connections.
  */
 export interface LakebaseConnectionInputs {
-  /** Lakebase project id. */
+  /** Lakebase project id. Resolved from the workspace when unset. */
   project?: string;
-  /** Branch id within the project. */
+  /** Branch id within the project. Defaults to the project's default branch. */
   branch?: string;
-  /** Canonical endpoint resource path (`projects/.../endpoints/...`). */
+  /**
+   * Canonical endpoint resource path (`projects/.../endpoints/...`), from
+   * `LAKEBASE_ENDPOINT`. Defaults to the branch's read-write endpoint.
+   */
   endpoint?: string;
-  /** Postgres database name (`PGDATABASE`). */
+  /** Postgres database name (`PGDATABASE`). Defaults to `databricks_postgres`. */
   database?: string;
-  /** Postgres hostname (`PGHOST`). */
+  /** Postgres hostname (`PGHOST`). Defaults to the resolved endpoint's host. */
   host?: string;
-  /** Postgres port (`PGPORT`). */
+  /** Postgres port (`PGPORT`). Defaults to 5432. */
   port?: number;
-  /** Postgres TLS mode (`PGSSLMODE`). */
+  /** Postgres TLS mode (`PGSSLMODE`). Defaults to `require`. */
   sslMode?: SslMode;
 }
 
@@ -88,6 +94,15 @@ 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.
+ *
+ * @example
+ * import { pgaddress } from "@dbx-tools/appkit";
+ *
+ * pgaddress.parseAddress("projects/demo/branches/production/endpoints/ep-1");
+ * // { project: "demo", branch: "production", endpointId: "ep-1", endpoint: "projects/..." }
+ *
+ * pgaddress.parseAddress("postgresql://me@ep-1.database.azuredatabricks.net/app?sslmode=require");
+ * // { host: "ep-1.database.azuredatabricks.net", user: "me", database: "app", sslMode: "require" }
  */
 export function parseAddress(input: string | undefined | null): ParsedAddress {
   if (!input) return {};
@@ -138,12 +153,16 @@ function parseUri(s: string): ParsedAddress {
   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") {
+  if (isSslMode(sslmode)) {
     result.sslMode = sslmode;
   }
   return result;
 }
 
+function isSslMode(value: string | undefined): value is SslMode {
+  return SSL_MODES.some((mode) => mode === value);
+}
+
 function parseResourcePathSegments(s: string): ParsedAddress {
   const parts = s.split("/");
   if (parts[0] !== "projects" || parts.length < 2) {

package/src/plugin.ts

@@ -20,7 +20,10 @@
  * @module
  */
 
-import { type NameLike } from "@dbx-tools/shared-core";
+import { ConfigurationError } from "@databricks/appkit";
+import { log, type NameLike } from "@dbx-tools/shared-core";
+
+const logger = log.logger("plugin");
 
 /**
  * Minimal structural shape of `this.context`. We mirror only the method we
@@ -71,6 +74,9 @@ const dataCache = new WeakMap<PluginDataFactory, PluginData>();
  */
 export function data<F extends PluginDataFactory, D extends ReturnType<F>>(factory: F): D {
   const cached = dataCache.get(factory);
+  // The cache is keyed by the erased `PluginDataFactory` bound, so `WeakMap`
+  // hands back the widened `PluginData`; only the caller's `F` knows the exact
+  // descriptor type.
   if (cached !== undefined) {
     return cached as D;
   }
@@ -100,6 +106,8 @@ export function instance<F extends PluginDataFactory>(
 ): PluginInstanceOf<F> | undefined {
   if (!ctx) return undefined;
   const name = data(factory).name;
+  // AppKit's registry is a `Map<string, unknown>`, so the instance type is only
+  // recoverable from the factory the caller passed.
   return ctx.getPlugins().get(name) as PluginInstanceOf<F> | undefined;
 }
 
@@ -127,5 +135,12 @@ export function require<F extends PluginDataFactory>(
   const prefix =
     typeof caller === "string" ? `${caller}: ` : caller?.name ? `${caller.name}: ` : "";
   const registeredName = data(factory).name;
-  throw new Error(`${prefix}required plugin not registered: ${registeredName}`);
+  logger.debug("required plugin not registered", {
+    plugin: registeredName,
+    registered: [...(ctx?.getPlugins().keys() ?? [])],
+  });
+  throw ConfigurationError.resourceNotFound(
+    `${prefix}plugin '${registeredName}'`,
+    `Add ${registeredName}() to the plugins passed to createApp.`,
+  );
 }

package/src/provision.ts

@@ -21,18 +21,36 @@
  * @module
  */
 
-import { createLakebasePool, getWorkspaceClient } from "@databricks/appkit";
-import { error, type log } from "@dbx-tools/shared-core";
+import { createLakebasePool, getWorkspaceClient, ValidationError } from "@databricks/appkit";
+import { error, log } from "@dbx-tools/shared-core";
 
 import { isAppEnv } from "./databricks";
 
+const defaultLogger = log.logger("provision");
+
 /** AppKit persistent-cache schema (see AppKit's `PersistentStorage`). */
 const CACHE_SCHEMA = "appkit";
 
+/**
+ * Characters a Lakebase Postgres role can legitimately contain: workspace
+ * identities are emails, service principals are UUIDs.
+ */
+const ROLE_PATTERN = /^[A-Za-z0-9._@+-]{1,255}$/;
+
+/**
+ * Provisioning runs at boot as the local developer identity, before any plugin
+ * exists, so it is outside AppKit's interceptor chain and inherits no timeout
+ * from it. Both bounds are explicit for that reason.
+ */
+const CONNECT_TIMEOUT_MS = 10_000;
+const STATEMENT_TIMEOUT_MS = 15_000;
+
 /**
  * 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.
+ * quote. Schema and role names are identifiers, and Postgres does not accept a
+ * bind parameter in an identifier position, so quoting is the only defense
+ * available for these statements. Lakebase role names are usually emails
+ * (`user@host`), which must be quoted to be a valid identifier at all.
  */
 function quoteIdent(ident: string): string {
   return `"${ident.replace(/"/g, '""')}"`;
@@ -42,8 +60,14 @@ function quoteIdent(ident: string): string {
  * 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.
+ *
+ * Throws a {@link ValidationError} when `role` is not a plausible Postgres role
+ * name, since the value lands in an identifier position.
  */
-function cacheGrantStatements(role: string): readonly string[] {
+export function cacheGrantStatements(role: string): readonly string[] {
+  if (!ROLE_PATTERN.test(role)) {
+    throw ValidationError.invalidValue("role", role, "a Postgres role name");
+  }
   const schema = quoteIdent(CACHE_SCHEMA);
   const target = quoteIdent(role);
   return [
@@ -66,10 +90,16 @@ function cacheGrantStatements(role: string): readonly string[] {
  *
  * @param role - Postgres role to grant to and connect as (the resolved
  *   workspace-client identity); skips when undefined.
+ * @param logger - Logger to report progress on. Defaults to this module's.
+ *
+ * @example
+ * import { provision } from "@dbx-tools/appkit";
+ *
+ * await provision.provisionCacheSchema("app-service-principal@databricks.com");
  */
 export async function provisionCacheSchema(
-  logger: log.Logger,
   role: string | undefined,
+  logger: log.Logger = defaultLogger,
 ): Promise<void> {
   if (isAppEnv()) {
     logger.debug("autopg: skip cache provisioning (inside a Databricks App)");
@@ -88,6 +118,8 @@ export async function provisionCacheSchema(
   const pool = createLakebasePool({
     user: role,
     workspaceClient: getWorkspaceClient({}),
+    connectionTimeoutMillis: CONNECT_TIMEOUT_MS,
+    statement_timeout: STATEMENT_TIMEOUT_MS,
   });
   try {
     const found = await pool.query(

package/test/config.test.ts

@@ -0,0 +1,63 @@
+import assert from "node:assert/strict";
+import { afterEach, describe, it } from "node:test";
+import { ValidationError } from "@databricks/appkit";
+import { resolveConfigValue, withCliSources, type ConfigSource } from "../src/config";
+
+const ENV_KEY = "DBX_TOOLS_APPKIT_TEST_VALUE";
+
+afterEach(() => {
+  delete process.env[ENV_KEY];
+});
+
+describe("config precedence", () => {
+  it("prefers explicit config over the environment", async () => {
+    process.env[ENV_KEY] = "from-env";
+    const value = await resolveConfigValue(ENV_KEY, {
+      explicit: { [ENV_KEY]: "from-explicit" },
+    });
+    assert.equal(value, "from-explicit");
+  });
+
+  it("falls back to the environment when explicit config omits the key", async () => {
+    process.env[ENV_KEY] = "from-env";
+    const value = await resolveConfigValue(ENV_KEY, { explicit: { OTHER_KEY: "x" } });
+    assert.equal(value, "from-env");
+  });
+
+  it("prepends explicit even when the caller passes its own source list", async () => {
+    process.env[ENV_KEY] = "from-env";
+    const value = await resolveConfigValue(ENV_KEY, {
+      sources: ["env"],
+      explicit: { [ENV_KEY]: "from-explicit" },
+    });
+    assert.equal(value, "from-explicit");
+  });
+
+  it("lets a cli flag win when cli sources are requested", async () => {
+    process.env[ENV_KEY] = "from-env";
+    const value = await resolveConfigValue(ENV_KEY, {
+      sources: withCliSources(),
+      cli: { [ENV_KEY]: "from-cli" },
+      explicit: { [ENV_KEY]: "from-explicit" },
+    });
+    assert.equal(value, "from-cli");
+  });
+
+  it("orders cli sources ahead of explicit and the default sources", () => {
+    assert.deepEqual(withCliSources(), ["cli", "explicit", "env", "bundle"]);
+    assert.deepEqual(withCliSources(["explicit", "env"]), ["cli", "explicit", "env"]);
+  });
+
+  it("skips blank values instead of treating them as resolved", async () => {
+    process.env[ENV_KEY] = "from-env";
+    const value = await resolveConfigValue(ENV_KEY, { explicit: { [ENV_KEY]: "   " } });
+    assert.equal(value, "from-env");
+  });
+
+  it("rejects an unknown source", async () => {
+    await assert.rejects(
+      () => resolveConfigValue(ENV_KEY, { sources: ["nope" as ConfigSource] }),
+      ValidationError,
+    );
+  });
+});

package/test/lakebase-resolver.test.ts

@@ -0,0 +1,82 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import { ValidationError } from "@databricks/appkit";
+import { nextPollDelay, parsePort, parseSslMode, pollDelay } from "../src/lakebase-resolver";
+
+describe("PGPORT validation", () => {
+  it("accepts a numeric string and a number", () => {
+    assert.equal(parsePort("5433"), 5433);
+    assert.equal(parsePort(5432), 5432);
+  });
+
+  it("treats absent and empty values as unset", () => {
+    assert.equal(parsePort(undefined), undefined);
+    assert.equal(parsePort(""), undefined);
+  });
+
+  it("rejects anything that is not a TCP port instead of yielding NaN", () => {
+    for (const bad of ["abc", "0", "70000", "5432.5"]) {
+      assert.throws(
+        () => parsePort(bad),
+        (err) => {
+          assert.ok(err instanceof ValidationError);
+          assert.match(err.message, /PGPORT/);
+          return true;
+        },
+      );
+    }
+  });
+});
+
+describe("PGSSLMODE validation", () => {
+  it("normalizes case and surrounding space", () => {
+    assert.equal(parseSslMode(" Require "), "require");
+    assert.equal(parseSslMode("disable"), "disable");
+  });
+
+  it("treats absent and empty values as unset", () => {
+    assert.equal(parseSslMode(undefined), undefined);
+    assert.equal(parseSslMode(""), undefined);
+  });
+
+  it("rejects a mode pg does not accept", () => {
+    assert.throws(
+      () => parseSslMode("verify-full"),
+      (err) => {
+        assert.ok(err instanceof ValidationError);
+        assert.match(err.message, /PGSSLMODE/);
+        return true;
+      },
+    );
+  });
+});
+
+describe("poll backoff", () => {
+  it("grows with the attempt and stays inside the jitter band", () => {
+    const base = 2_000;
+    for (let attempt = 0; attempt < 10; attempt++) {
+      const delay = nextPollDelay(attempt, base);
+      assert.ok(delay >= 0);
+      assert.ok(delay <= 15_000 * 1.2 + 1, `attempt ${attempt} produced ${delay}`);
+    }
+    const first = Array.from({ length: 20 }, () => nextPollDelay(0, base));
+    const later = Array.from({ length: 20 }, () => nextPollDelay(4, base));
+    assert.ok(Math.max(...first) < Math.min(...later));
+  });
+
+  it("caps the delay so a long wait keeps polling", () => {
+    assert.ok(nextPollDelay(50, 2_000) <= 15_000 * 1.2 + 1);
+  });
+
+  it("rejects with the abort reason when cancelled mid-wait", async () => {
+    const controller = new AbortController();
+    const waiting = pollDelay(5, 10_000, controller.signal);
+    const reason = new Error("boot cancelled");
+    controller.abort(reason);
+    await assert.rejects(waiting, reason);
+  });
+
+  it("rejects immediately for an already-aborted signal", async () => {
+    await assert.rejects(pollDelay(0, 10_000, AbortSignal.abort(new Error("gone"))), /gone/);
+  });
+});

package/test/pgaddress.test.ts

@@ -0,0 +1,76 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import { parseAddress, parseResourcePath, SSL_MODES } from "../src/pgaddress";
+
+describe("pgaddress parseAddress", () => {
+  it("returns nothing for empty and unrecognized input", () => {
+    assert.deepEqual(parseAddress(undefined), {});
+    assert.deepEqual(parseAddress(""), {});
+    assert.deepEqual(parseAddress("   "), {});
+    assert.deepEqual(parseAddress("Not An Address"), {});
+  });
+
+  it("splits a Postgres URI into host, user, database, port and ssl mode", () => {
+    const parsed = parseAddress(
+      "postgresql://me%40acme.com@ep-1.database.eastus2.azuredatabricks.net:5433/app?sslmode=disable",
+    );
+    assert.equal(parsed.host, "ep-1.database.eastus2.azuredatabricks.net");
+    assert.equal(parsed.user, "me@acme.com");
+    assert.equal(parsed.database, "app");
+    assert.equal(parsed.port, 5433);
+    assert.equal(parsed.sslMode, "disable");
+  });
+
+  it("ignores an ssl mode the driver does not accept", () => {
+    assert.equal(
+      parseAddress("postgres://h.example.com/db?sslmode=verify-full").sslMode,
+      undefined,
+    );
+    assert.deepEqual([...SSL_MODES], ["require", "disable", "prefer"]);
+  });
+
+  it("recovers project, branch and endpoint from a canonical endpoint path", () => {
+    const path = "projects/demo/branches/production/endpoints/ep-1";
+    assert.deepEqual(parseAddress(path), {
+      project: "demo",
+      branch: "production",
+      endpointId: "ep-1",
+      endpoint: path,
+    });
+  });
+
+  it("keeps a database resource id separate from PGDATABASE", () => {
+    assert.deepEqual(
+      parseAddress("projects/demo/branches/production/databases/databricks-postgres"),
+      {
+        project: "demo",
+        branch: "production",
+        databaseResourceId: "databricks-postgres",
+      },
+    );
+  });
+
+  it("reads shorter resource paths", () => {
+    assert.deepEqual(parseAddress("projects/demo"), { project: "demo" });
+    assert.deepEqual(parseAddress("projects/demo/branches/main"), {
+      project: "demo",
+      branch: "main",
+    });
+    assert.deepEqual(parseAddress("projects/demo/branches"), {});
+  });
+
+  it("treats a dotted value as a hostname and a bare slug as a project id", () => {
+    assert.deepEqual(parseAddress("ep-1.database.azuredatabricks.net"), {
+      host: "ep-1.database.azuredatabricks.net",
+    });
+    assert.deepEqual(parseAddress("dbx-tools-demo"), { project: "dbx-tools-demo" });
+  });
+});
+
+describe("pgaddress parseResourcePath", () => {
+  it("only accepts `projects/` paths so a bare branch id is not read as a project", () => {
+    assert.deepEqual(parseResourcePath("production"), {});
+    assert.deepEqual(parseResourcePath(undefined), {});
+    assert.equal(parseResourcePath("projects/demo/branches/main").branch, "main");
+  });
+});

package/test/plugin.test.ts

@@ -0,0 +1,64 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import { ConfigurationError } from "@databricks/appkit";
+import { data, instance, require as requirePlugin, type PluginContextLike } from "../src/plugin";
+
+class FakeLakebasePlugin {
+  exports() {
+    return { pool: "pool" };
+  }
+}
+
+function fakeFactory(name: string, calls: { count: number }) {
+  return () => {
+    calls.count += 1;
+    return { plugin: FakeLakebasePlugin, name };
+  };
+}
+
+function fakeContext(entries: Record<string, unknown>): PluginContextLike {
+  const plugins = new Map(Object.entries(entries));
+  return { getPlugins: () => plugins };
+}
+
+describe("plugin lookup", () => {
+  it("caches the factory descriptor per factory", () => {
+    const calls = { count: 0 };
+    const factory = fakeFactory("lakebase", calls);
+    assert.equal(data(factory).name, "lakebase");
+    assert.equal(data(factory).name, "lakebase");
+    assert.equal(calls.count, 1);
+  });
+
+  it("returns the registered instance, or undefined without a context", () => {
+    const factory = fakeFactory("lakebase", { count: 0 });
+    const plugin = new FakeLakebasePlugin();
+    assert.equal(instance(fakeContext({ lakebase: plugin }), factory), plugin);
+    assert.equal(instance(fakeContext({}), factory), undefined);
+    assert.equal(instance(undefined, factory), undefined);
+  });
+
+  it("require returns the instance when registered", () => {
+    const factory = fakeFactory("lakebase", { count: 0 });
+    const plugin = new FakeLakebasePlugin();
+    assert.equal(requirePlugin(fakeContext({ lakebase: plugin }), factory), plugin);
+  });
+
+  it("require throws a ConfigurationError naming the plugin and the caller", () => {
+    const factory = fakeFactory("lakebase", { count: 0 });
+    assert.throws(
+      () => requirePlugin(fakeContext({ server: {} }), factory, "mastra"),
+      (err) => {
+        assert.ok(err instanceof ConfigurationError);
+        assert.match(err.message, /mastra/);
+        assert.match(err.message, /lakebase/);
+        return true;
+      },
+    );
+  });
+
+  it("require throws without a context at all", () => {
+    const factory = fakeFactory("lakebase", { count: 0 });
+    assert.throws(() => requirePlugin(undefined, factory), ConfigurationError);
+  });
+});

package/test/provision.test.ts

@@ -0,0 +1,35 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import { ValidationError } from "@databricks/appkit";
+import { cacheGrantStatements } from "../src/provision";
+
+describe("cache schema grants", () => {
+  it("quotes the schema and an email-shaped role", () => {
+    const statements = cacheGrantStatements("me@acme.com");
+    assert.equal(statements.length, 5);
+    for (const sql of statements) {
+      assert.match(sql, /"appkit"/);
+      assert.match(sql, /"me@acme\.com"/);
+    }
+    assert.equal(statements[0], 'GRANT USAGE, CREATE ON SCHEMA "appkit" TO "me@acme.com"');
+  });
+
+  it("rejects a role name that could not be an identifier", () => {
+    for (const bad of ['ev"il', "role; DROP SCHEMA appkit", "", "a b"]) {
+      assert.throws(
+        () => cacheGrantStatements(bad),
+        (err) => {
+          assert.ok(err instanceof ValidationError);
+          assert.match(err.message, /role/);
+          return true;
+        },
+      );
+    }
+  });
+
+  it("keeps the grants idempotent and scoped to the cache schema", () => {
+    const statements = cacheGrantStatements("sp-1234").join("\n");
+    assert.match(statements, /ALTER DEFAULT PRIVILEGES IN SCHEMA "appkit"/);
+    assert.doesNotMatch(statements, /public/);
+  });
+});