feat(appkit): add appkit db migrate

.oxlintrc.json

@@ -61,6 +61,56 @@
       "rules": {
         "no-restricted-imports": "off"
       }
+    },
+    {
+      "files": ["packages/appkit/src/**"],
+      "rules": {
+        "no-restricted-imports": [
+          "error",
+          {
+            "patterns": [
+              {
+                "group": [
+                  "@databricks/sdk-experimental",
+                  "@databricks/sdk-experimental/**"
+                ],
+                "message": "Import the Databricks SDK only through the wrapper in packages/shared/src/workspace-client. Add a re-export there if you need a new symbol."
+              },
+              {
+                "group": ["drizzle-kit", "drizzle-kit/**"],
+                "message": "drizzle-kit is an optional dev dependency that deployed apps do not install. Only packages/appkit/src/database/tooling may import it."
+              },
+              {
+                "group": [
+                  "**/tooling",
+                  "**/tooling/**",
+                  "@databricks/appkit/database-tooling"
+                ],
+                "message": "The database tooling may import drizzle-kit, so only the CLI loads it. Runtime code must not import it."
+              }
+            ]
+          }
+        ]
+      }
+    },
+    {
+      "files": ["packages/appkit/src/database/tooling/**"],
+      "rules": {
+        "no-restricted-imports": [
+          "error",
+          {
+            "patterns": [
+              {
+                "group": [
+                  "@databricks/sdk-experimental",
+                  "@databricks/sdk-experimental/**"
+                ],
+                "message": "Import the Databricks SDK only through the wrapper in packages/shared/src/workspace-client. Add a re-export there if you need a new symbol."
+              }
+            ]
+          }
+        ]
+      }
     }
   ],
   "ignorePatterns": [

knip.json

@@ -12,7 +12,8 @@
         "vitest",
         "@databricks/sdk-experimental",
         "@mlflow/core",
-        "autoevals"
+        "autoevals",
+        "drizzle-kit"
       ]
     },
     "packages/appkit-ui": {

packages/appkit/package.json

@@ -42,6 +42,11 @@
       "development": "./src/type-generator/index.ts",
       "default": "./dist/type-generator/index.js"
     },
+    "./database-tooling": {
+      "types": "./dist/database/tooling/index.d.ts",
+      "development": "./src/database/tooling/index.ts",
+      "default": "./dist/database/tooling/index.js"
+    },
     "./testing": {
       "types": "./dist/testing/index.d.ts",
       "development": "./src/testing/index.ts",
@@ -108,6 +113,7 @@
   "peerDependencies": {
     "@mlflow/core": "0.4.0",
     "autoevals": "0.3.0",
+    "drizzle-kit": "0.31.10",
     "vitest": ">=3"
   },
   "peerDependenciesMeta": {
@@ -117,6 +123,9 @@
     "autoevals": {
       "optional": true
     },
+    "drizzle-kit": {
+      "optional": true
+    },
     "vitest": {
       "optional": true
     }
@@ -131,6 +140,7 @@
     "@types/ws": "8.18.1",
     "@vitejs/plugin-react": "5.1.1",
     "autoevals": "0.3.0",
+    "drizzle-kit": "0.31.10",
     "vitest": "3.2.4"
   },
   "overrides": {
@@ -143,6 +153,7 @@
       "./beta": "./dist/beta.js",
       "./dist/shared/src/plugin": "./dist/shared/src/plugin.d.ts",
       "./type-generator": "./dist/type-generator/index.js",
+      "./database-tooling": "./dist/database/tooling/index.js",
       "./testing": {
         "types": "./dist/testing/index.d.ts",
         "default": "./dist/testing/index.js"

packages/appkit/src/connectors/lakebase/endpoint-host.ts

@@ -134,7 +134,7 @@ async function compareHosts(
 }
 
 /** Read every host the endpoint serves, read-write and read-only. */
-function endpointHosts(response: unknown): string[] {
+export function endpointHosts(response: unknown): string[] {
   if (!response || typeof response !== "object") return [];
   const status = Reflect.get(response, "status");
   if (!status || typeof status !== "object") return [];

packages/appkit/src/database/migrations/apply.ts

@@ -1,6 +1,5 @@
 import { createHash } from "node:crypto";
 import fs from "node:fs";
-import path from "node:path";
 import { setTimeout as sleep } from "node:timers/promises";
 
 import { readMigrationFiles } from "drizzle-orm/migrator";
@@ -22,8 +21,10 @@ import {
   type Journal,
   MIGRATIONS_TABLE,
   MigrationFilesError,
+  migrationFile,
   readJournal,
   readJsonFile,
+  snapshotFile,
 } from "./paths";
 
 const logger = createLogger("database");
@@ -228,7 +229,7 @@ type QueryClient = Pick<PoolClient, "query">;
 /** Every committed migration, in journal order. Fails on a missing file before connecting. */
 function readMigrations(folder: string, journal: Journal): MigrationFile[] {
   for (const entry of journal.entries) {
-    const file = path.join(folder, `${entry.tag}.sql`);
+    const file = migrationFile(folder, entry.tag);
     if (!fs.existsSync(file)) {
       throw new MigrationFilesError(
         `Missing ${file}, listed in meta/_journal.json. Restore it from version control.`,
@@ -454,11 +455,7 @@ function snapshotTables(
   idx: number,
   schemaName: string,
 ): Map<string, string[]> {
-  const file = path.join(
-    folder,
-    "meta",
-    `${String(idx).padStart(4, "0")}_snapshot.json`,
-  );
+  const file = snapshotFile(folder, idx);
   if (!fs.existsSync(file)) {
     throw new MigrationFilesError(
       `Missing ${file} for ${BASELINE_TAG}. Restore it from version control; it lists the tables the baseline adopts.`,

packages/appkit/src/database/migrations/paths.ts

@@ -33,9 +33,32 @@ export function migrationsFolder(root: string): string {
   return path.resolve(root, MIGRATIONS_DIR);
 }
 
+// drizzle-kit 0.31's v7 folder layout. drizzle-kit v1 replaces it with one
+// folder per migration, so every path into the folder is built here.
+
+/** Folder of the journal and the snapshots. */
+export function metaFolder(folder: string): string {
+  return path.join(folder, "meta");
+}
+
 /** Path of the journal inside a migrations folder. */
 export function journalFile(folder: string): string {
-  return path.join(folder, "meta", "_journal.json");
+  return path.join(metaFolder(folder), "_journal.json");
+}
+
+/** The `NNNN` that starts the tag of the migration at journal index idx. */
+export function migrationPrefix(idx: number): string {
+  return String(idx).padStart(4, "0");
+}
+
+/** Path of a migration's SQL. */
+export function migrationFile(folder: string, tag: string): string {
+  return path.join(folder, `${tag}.sql`);
+}
+
+/** Path of the snapshot written with the migration at journal index idx. */
+export function snapshotFile(folder: string, idx: number): string {
+  return path.join(metaFolder(folder), `${migrationPrefix(idx)}_snapshot.json`);
 }
 
 /** A committed migration file is missing or unreadable; the message names the file. */

packages/appkit/src/database/tooling/context.ts

@@ -0,0 +1,152 @@
+import type { LakebasePoolConfig } from "@databricks/lakebase";
+import type { Pool } from "pg";
+
+import { createLakebasePool } from "../../connectors/lakebase";
+import { endpointHosts } from "../../connectors/lakebase/endpoint-host";
+import { getClientOptions } from "../../context/client-options";
+import { createWorkspaceClient } from "../../workspace-client";
+import { currentIdentity } from "./identity";
+import { LakebaseApi, type LakebaseEndpoint } from "./lakebase-api";
+
+/** One REST call through the SDK's unified-auth client. */
+export interface ToolingRequest {
+  path: string;
+  method: "GET" | "POST" | "PATCH" | "DELETE";
+  query?: Record<string, string | number | boolean | undefined>;
+  headers: Headers;
+  raw: boolean;
+  payload?: unknown;
+}
+
+/**
+ * The slice of the SDK client the tooling uses. It is also everything the
+ * Lakebase pool needs: token refresh calls `apiClient.request` and the user
+ * lookup calls `currentUser.me`.
+ */
+export interface ToolingWorkspaceClient {
+  readonly apiClient: { request(options: ToolingRequest): Promise<unknown> };
+  readonly currentUser: { me(): Promise<{ userName?: string }> };
+}
+
+/** Where a command runs: the app root, its environment and its Databricks identity. */
+export interface DbContext {
+  readonly root: string;
+  readonly env: NodeJS.ProcessEnv;
+  /** Lazy; the profile comes from DATABRICKS_CONFIG_PROFILE. */
+  workspace(): ToolingWorkspaceClient;
+  /** Lazy Lakebase REST client over workspace(). */
+  api(): LakebaseApi;
+}
+
+export function createDbContext(
+  options: { root?: string; env?: NodeJS.ProcessEnv } = {},
+): DbContext {
+  const root = options.root ?? process.cwd();
+  const env = options.env ?? process.env;
+  let client: ToolingWorkspaceClient | undefined;
+  let api: LakebaseApi | undefined;
+  const ctx: DbContext = {
+    root,
+    env,
+    workspace() {
+      if (!client) {
+        const profile = env.DATABRICKS_CONFIG_PROFILE;
+        client = createWorkspaceClient({
+          ...(profile ? { profile } : {}),
+          clientOptions: getClientOptions(),
+        }).toLegacyWorkspaceClient();
+      }
+      return client;
+    },
+    api() {
+      api ??= new LakebaseApi(ctx.workspace());
+      return api;
+    },
+  };
+  return ctx;
+}
+
+/**
+ * The guard checks the branch named by LAKEBASE_ENDPOINT, but the pool connects
+ * to PGHOST. The runtime only warns when it cannot verify that PGHOST serves that
+ * endpoint; a command that writes refuses instead.
+ */
+async function assertHostServesEndpoint(
+  ctx: DbContext,
+  endpoint: string,
+  host: string,
+): Promise<void> {
+  let found: LakebaseEndpoint | null;
+  try {
+    found = await ctx.api().getEndpoint(endpoint);
+  } catch (error) {
+    throw new Error(
+      `Could not verify that PGHOST ${host} serves LAKEBASE_ENDPOINT ${endpoint}, so it did not connect.`,
+      { cause: error },
+    );
+  }
+  if (!found) {
+    throw new Error(
+      `LAKEBASE_ENDPOINT ${endpoint} does not exist, so it did not connect. Run \`appkit db init\` to point .env at your branch.`,
+    );
+  }
+  const hosts = endpointHosts(found);
+  if (!hosts.includes(host.toLowerCase())) {
+    throw new Error(
+      `PGHOST ${host} is not a host of LAKEBASE_ENDPOINT ${endpoint} (${hosts.length > 0 ? `expected ${hosts.join(" or ")}` : "it has no hosts"}), so it did not connect. Run \`appkit db init\` to rewrite .env.`,
+    );
+  }
+}
+
+export const NO_ENV_BRANCH =
+  "no branch is configured in .env (PGHOST, PGDATABASE and LAKEBASE_ENDPOINT); run `appkit db init`";
+
+/** Whether .env names a branch to connect to. */
+export function hasEnvBranch(env: NodeJS.ProcessEnv): boolean {
+  return Boolean(env.PGHOST && env.PGDATABASE && env.LAKEBASE_ENDPOINT);
+}
+
+/**
+ * The pool is created directly: initializeLakebasePool would look the endpoint
+ * and the user up a second time.
+ */
+async function openBranchPool(ctx: DbContext): Promise<Pool> {
+  const { PGHOST, PGDATABASE, LAKEBASE_ENDPOINT } = ctx.env;
+  if (!PGHOST || !PGDATABASE || !LAKEBASE_ENDPOINT) {
+    throw new Error(`Did not connect: ${NO_ENV_BRANCH}.`);
+  }
+  const [user] = await Promise.all([
+    ctx.env.PGUSER || currentIdentity(ctx).then(({ userName }) => userName),
+    assertHostServesEndpoint(ctx, LAKEBASE_ENDPOINT, PGHOST),
+  ]);
+  const port = ctx.env.PGPORT ? Number(ctx.env.PGPORT) : undefined;
+  const sslMode = ctx.env.PGSSLMODE as LakebasePoolConfig["sslMode"];
+  return createLakebasePool({
+    host: PGHOST,
+    database: PGDATABASE,
+    endpoint: LAKEBASE_ENDPOINT,
+    user,
+    ...(port ? { port } : {}),
+    ...(sslMode ? { sslMode } : {}),
+    // Only apiClient.request and currentUser.me are used (see ToolingWorkspaceClient).
+    workspaceClient: ctx.workspace() as unknown as NonNullable<
+      LakebasePoolConfig["workspaceClient"]
+    >,
+    // One-shot CLI pools: a small cap, and no lingering idle connections.
+    max: 2,
+    idleTimeoutMillis: 1_000,
+  });
+}
+
+/** Run fn with a pool on the .env branch, and close the pool after. */
+export async function withBranchPool<T>(
+  ctx: DbContext,
+  fn: (pool: Pool) => Promise<T>,
+): Promise<T> {
+  const pool = await openBranchPool(ctx);
+  try {
+    return await fn(pool);
+  } finally {
+    await pool.end().catch(() => undefined);
+  }
+}

packages/appkit/src/database/tooling/drizzle-exports.ts

@@ -0,0 +1,24 @@
+import { getTableColumns, is, type Table } from "drizzle-orm";
+import { PgEnumColumn } from "drizzle-orm/pg-core";
+
+import type { Schema } from "../schema-builder";
+
+/**
+ * The "module exports" drizzle-kit snapshots: every engine table plus every
+ * enum a column references. The pgSchema object is never included, so
+ * generated SQL contains no CREATE SCHEMA; the migrator creates the schema as
+ * the identity that applies the migrations, which then owns it.
+ */
+export function drizzleKitImports(schema: Schema): Record<string, unknown> {
+  const imports: Record<string, unknown> = Object.create(null);
+  for (const [name, table] of Object.entries(schema.$engine)) {
+    imports[`table:${name}`] = table;
+    const columns = getTableColumns(table as unknown as Table);
+    for (const column of Object.values(columns)) {
+      if (is(column, PgEnumColumn)) {
+        imports[`enum:${column.enum.enumName}`] = column.enum;
+      }
+    }
+  }
+  return imports;
+}