@dbx-tools/appkit-autopg 0.1.0

package/src/address.ts

@@ -0,0 +1,148 @@
+/**
+ * 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
+ * REST 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`. The database leaf isn't surfaced because
+ *   it's a resource id, not the Postgres database name; the resolver
+ *   will look up the real `status.postgres_database` value 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.
+ */
+
+import type { SslMode } from "./resolver.js";
+
+export interface ParsedAddress {
+  /** Lakebase project id. */
+  project?: string;
+  /** Branch id within the project. */
+  branch?: string;
+  /** Endpoint leaf id (last segment of an endpoint resource path). */
+  endpointId?: string;
+  /** Canonical endpoint resource path; only set for matching inputs. */
+  endpoint?: string;
+  /** Postgres database name (PGDATABASE) when parsed from a URI path. */
+  database?: string;
+  /** Postgres hostname. */
+  host?: string;
+  /** Postgres port. */
+  port?: number;
+  /** Postgres user (URI-decoded if encoded). */
+  user?: string;
+  /** Postgres TLS mode. */
+  sslMode?: SslMode;
+}
+
+const URL_SCHEME_RE = /^(postgres|postgresql):\/\//i;
+const RESOURCE_ENDPOINT_RE =
+  /^projects\/([^/]+)\/branches\/([^/]+)\/endpoints\/([^/]+)$/;
+const RESOURCE_DATABASE_RE =
+  /^projects\/([^/]+)\/branches\/([^/]+)\/databases\/([^/]+)$/;
+const RESOURCE_BRANCH_RE = /^projects\/([^/]+)\/branches\/([^/]+)$/;
+const RESOURCE_PROJECT_RE = /^projects\/([^/]+)$/;
+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 parseResourcePath(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 {};
+}
+
+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 {
+      // Malformed percent-escape; keep the raw form.
+      result.user = url.username;
+    }
+  }
+  const db = url.pathname.replace(/^\//, "");
+  if (db) result.database = decodeURIComponent(db);
+  // Postgres clients accept `sslmode`; URL params are case-sensitive but
+  // we tolerate either since users paste both.
+  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 parseResourcePath(s: string): ParsedAddress {
+  const ep = RESOURCE_ENDPOINT_RE.exec(s);
+  if (ep) {
+    return {
+      project: ep[1],
+      branch: ep[2],
+      endpointId: ep[3],
+      endpoint: s,
+    };
+  }
+  // databases/{d} is a resource id (often kebab-case), not the actual
+  // Postgres database name. Surface project + branch only; the resolver
+  // will fetch the real `postgres_database` value.
+  const db = RESOURCE_DATABASE_RE.exec(s);
+  if (db) return { project: db[1], branch: db[2] };
+  const br = RESOURCE_BRANCH_RE.exec(s);
+  if (br) return { project: br[1], branch: br[2] };
+  const pr = RESOURCE_PROJECT_RE.exec(s);
+  if (pr) return { project: pr[1] };
+  return {};
+}

package/src/autopg.ts

@@ -0,0 +1,93 @@
+/**
+ * Top-level Lakebase auto-discovery helper.
+ *
+ * Read this once at process startup BEFORE `createApp(...)` so the
+ * `lakebase` plugin (and anyone else who reads `process.env` during
+ * `setup()`) sees a fully-populated environment:
+ *
+ * ```ts
+ * import { autopg } from "@dbx-tools/appkit-autopg";
+ * import { createApp, lakebase, server } from "@databricks/appkit";
+ *
+ * await autopg();              // resolves + writes process.env
+ * await createApp({ plugins: [lakebase(), server()] });
+ * ```
+ *
+ * `autopg` is intentionally NOT an AppKit plugin. AppKit's `static phase`
+ * field only orders plugin `setup()` invocation, not async completion -
+ * `lakebase.setup()` calls `parsePoolConfig` synchronously after its
+ * first `await` and would throw on `PGHOST` before any sibling plugin's
+ * REST resolution could finish. Awaiting `autopg()` upfront sidesteps
+ * the race entirely.
+ *
+ * Inputs flow in this priority order:
+ *   1. Explicit `config.<field>` argument
+ *   2. Matching env var (`LAKEBASE_PROJECT`, `LAKEBASE_BRANCH`,
+ *      `LAKEBASE_ENDPOINT`, `PGHOST`, `PGDATABASE`, `PGPORT`, `PGSSLMODE`)
+ *   3. Derived from the Lakebase Autoscaling REST API under
+ *      `/api/2.0/postgres/` via the Databricks workspace client
+ *
+ * Resolved values are written back to `process.env` (only filling gaps;
+ * existing values are preserved) so the downstream `lakebase` plugin
+ * picks them up. Pass `{ exportEnv: false }` to keep `process.env`
+ * untouched and just inspect the returned record.
+ */
+
+import { logUtils } from "@dbx-tools/appkit-shared";
+
+import {
+  applyToEnv,
+  resolveConnection,
+  type Resolved,
+  type ResolverInputs,
+} from "./resolver.js";
+
+/** Options accepted by {@link autopg}. */
+export interface AutopgOptions extends ResolverInputs {
+  /**
+   * When `true` (the default), resolved values are written to
+   * `process.env` so the `lakebase` plugin sees them at startup.
+   * Set to `false` to leave `process.env` untouched and just receive
+   * the resolved record back.
+   */
+  exportEnv?: boolean;
+}
+
+/**
+ * Resolve Lakebase Postgres connection info from config + env (and the
+ * Databricks REST API when needed), write the resolved values to
+ * `process.env`, and return the fully-populated record.
+ *
+ * Always safe to call: when env already provides every field, it
+ * returns immediately without any network traffic.
+ *
+ * @throws when a `project` is set (directly or via env) but the
+ *   Databricks API returns no branches / endpoints / databases to
+ *   choose from. The error message lists the available candidates so
+ *   the caller can pin the right one via env or config.
+ */
+export async function autopg(opts: AutopgOptions = {}): Promise<Resolved> {
+  const { exportEnv = true, ...inputs } = opts;
+  const log = logUtils.logger("autopg");
+  const resolved = await resolveConnection(inputs, log);
+  if (exportEnv) {
+    applyToEnv(resolved);
+    log.info("env updated", redactForLog(resolved));
+  } else {
+    log.info("resolved (env untouched)", redactForLog(resolved));
+  }
+  return resolved;
+}
+
+/** Strip resolved record to log-safe primitive fields. */
+function redactForLog(resolved: Resolved): Record<string, unknown> {
+  return {
+    project: resolved.project,
+    branch: resolved.branch,
+    endpoint: resolved.endpoint,
+    database: resolved.database,
+    host: resolved.host,
+    port: resolved.port,
+    sslMode: resolved.sslMode,
+  };
+}