From 98ef27bcf47035fc1de3d59fec38a200cd4062a3 Mon Sep 17 00:00:00 2001 From: ditadi Date: Tue, 13 Oct 2026 00:05:20 +0100 Subject: [PATCH] feat(appkit): add appkit db init, clean and seed `appkit db init` gives a checkout its own Lakebase branch, forked from the project's default branch: your dev branch, or with --session a branch for this checkout that expires after 7 days (or --ttl). It asks nothing. It refuses the default branch, protected branches and branches databricks.yml deploys to before any change, and when the project is at its branch limit it lists your branches to delete. It makes sure the branch has a read-write endpoint that suspends after 5 idle minutes and a Postgres role for you, points .env at it with an atomic 0600 write, hands you the app schema when the fork copied it owned by another Lakebase role, installs drizzle-kit when the app lacks it, then starts the schema by migrating, introspecting or cleaning. `appkit db clean` drops the app schema (only your objects in public) and lists what outside it went along; --delete-branch deletes this checkout's session branch instead. `appkit db seed` runs config/database/seed.sql in one transaction with the app schema on the search_path. Both run only on a branch created for you or this checkout. Signed-off-by: ditadi --- .../appkit/src/database/migrations/apply.ts | 11 +- .../src/database/migrations/ownership.ts | 3 +- .../migrations/tests/ownership.test.ts | 3 + .../appkit/src/database/tooling/branch-sql.ts | 24 + packages/appkit/src/database/tooling/clean.ts | 219 +++++ .../appkit/src/database/tooling/context.ts | 2 + .../database/tooling/drizzle-kit-install.ts | 91 ++ .../appkit/src/database/tooling/env-file.ts | 124 +++ packages/appkit/src/database/tooling/guard.ts | 24 +- .../appkit/src/database/tooling/hand-off.ts | 104 +++ .../appkit/src/database/tooling/identity.ts | 13 +- packages/appkit/src/database/tooling/index.ts | 15 + packages/appkit/src/database/tooling/init.ts | 572 ++++++++++++ .../src/database/tooling/lakebase-api.ts | 277 +++++- .../appkit/src/database/tooling/migrate.ts | 31 +- .../appkit/src/database/tooling/migrations.ts | 16 + .../appkit/src/database/tooling/naming.ts | 76 +- packages/appkit/src/database/tooling/seed.ts | 47 + .../appkit/src/database/tooling/targets.ts | 57 +- .../database/tooling/tests/clean-seed.test.ts | 262 ++++++ .../database/tooling/tests/clean.pg.test.ts | 164 ++++ .../database/tooling/tests/contract.test.ts | 5 + .../tooling/tests/drizzle-kit-install.test.ts | 74 ++ .../database/tooling/tests/env-file.test.ts | 181 ++++ .../src/database/tooling/tests/fakes.ts | 141 ++- .../src/database/tooling/tests/guard.test.ts | 82 +- .../tooling/tests/hand-off.pg.test.ts | 75 ++ .../database/tooling/tests/identity.test.ts | 10 +- .../src/database/tooling/tests/init.test.ts | 838 ++++++++++++++++++ .../tooling/tests/lakebase-api.test.ts | 278 +++++- .../database/tooling/tests/migrate.test.ts | 23 + .../src/database/tooling/tests/naming.test.ts | 59 +- .../database/tooling/tests/targets.test.ts | 71 +- packages/shared/src/cli/commands/db/clean.ts | 81 ++ packages/shared/src/cli/commands/db/index.ts | 8 +- .../cli/commands/db/init-clean-seed.test.ts | 377 ++++++++ packages/shared/src/cli/commands/db/init.ts | 143 +++ packages/shared/src/cli/commands/db/output.ts | 21 +- packages/shared/src/cli/commands/db/seed.ts | 27 + .../src/cli/commands/db/tooling-contract.ts | 58 ++ 40 files changed, 4617 insertions(+), 70 deletions(-) create mode 100644 packages/appkit/src/database/tooling/branch-sql.ts create mode 100644 packages/appkit/src/database/tooling/clean.ts create mode 100644 packages/appkit/src/database/tooling/drizzle-kit-install.ts create mode 100644 packages/appkit/src/database/tooling/env-file.ts create mode 100644 packages/appkit/src/database/tooling/hand-off.ts create mode 100644 packages/appkit/src/database/tooling/init.ts create mode 100644 packages/appkit/src/database/tooling/seed.ts create mode 100644 packages/appkit/src/database/tooling/tests/clean-seed.test.ts create mode 100644 packages/appkit/src/database/tooling/tests/clean.pg.test.ts create mode 100644 packages/appkit/src/database/tooling/tests/drizzle-kit-install.test.ts create mode 100644 packages/appkit/src/database/tooling/tests/env-file.test.ts create mode 100644 packages/appkit/src/database/tooling/tests/hand-off.pg.test.ts create mode 100644 packages/appkit/src/database/tooling/tests/init.test.ts create mode 100644 packages/shared/src/cli/commands/db/clean.ts create mode 100644 packages/shared/src/cli/commands/db/init-clean-seed.test.ts create mode 100644 packages/shared/src/cli/commands/db/init.ts create mode 100644 packages/shared/src/cli/commands/db/seed.ts diff --git a/packages/appkit/src/database/migrations/apply.ts b/packages/appkit/src/database/migrations/apply.ts index b9f528640..14d08800c 100644 --- a/packages/appkit/src/database/migrations/apply.ts +++ b/packages/appkit/src/database/migrations/apply.ts @@ -207,7 +207,7 @@ function systemCodeOf(error: unknown): string | undefined { return findCode(error, (code) => SYSTEM_CODE.test(code)); } -function quoteIdentifier(name: string): string { +export function quoteIdentifier(name: string): string { return `"${name.replace(/"/g, '""')}"`; } @@ -367,14 +367,15 @@ async function acquireMigrationLock( } } -async function inTransaction( +export async function inTransaction( client: QueryClient, - work: () => Promise, -): Promise { + work: () => Promise, +): Promise { await client.query("begin"); try { - await work(); + const result = await work(); await client.query("commit"); + return result; } catch (error) { await client.query("rollback").catch(() => undefined); throw error; diff --git a/packages/appkit/src/database/migrations/ownership.ts b/packages/appkit/src/database/migrations/ownership.ts index 3136b0daa..59a7be8d4 100644 --- a/packages/appkit/src/database/migrations/ownership.ts +++ b/packages/appkit/src/database/migrations/ownership.ts @@ -64,7 +64,8 @@ export function ownershipHint( `Schema "${ownership.schemaName}" is owned by "${ownership.owner}", but the app connects as "${ownership.currentUser}". ` + "This happens when the app's database resource was removed and added again, or the app was recreated with a new service principal: " + `Lakebase hands the objects to another role. Have "${ownership.owner}" transfer the schema and its tables to "${ownership.currentUser}" ` + - "(ALTER … OWNER TO), or set another schemaName in schema.ts." + "(ALTER … OWNER TO), or set another schemaName in schema.ts. " + + "On your dev or session branch, rerun `appkit db init`, which hands you the schema." ); } return ( diff --git a/packages/appkit/src/database/migrations/tests/ownership.test.ts b/packages/appkit/src/database/migrations/tests/ownership.test.ts index 27acc56b5..d1dda3838 100644 --- a/packages/appkit/src/database/migrations/tests/ownership.test.ts +++ b/packages/appkit/src/database/migrations/tests/ownership.test.ts @@ -26,6 +26,9 @@ describe("ownership: one owner per branch", () => { "the app's database resource was removed and added again, or the app was recreated", ); expect(hint).toContain("set another schemaName in schema.ts"); + expect(hint).toContain( + "On your dev or session branch, rerun `appkit db init`, which hands you the schema.", + ); }); test("a schema that does not exist yet points at CREATE on the database", () => { diff --git a/packages/appkit/src/database/tooling/branch-sql.ts b/packages/appkit/src/database/tooling/branch-sql.ts new file mode 100644 index 000000000..8f2de812e --- /dev/null +++ b/packages/appkit/src/database/tooling/branch-sql.ts @@ -0,0 +1,24 @@ +import type { PoolClient } from "pg"; + +import { inTransaction } from "../migrations/apply"; +import { type DbContext, withBranchPool } from "./context"; + +/** Run work in one transaction on the `.env` branch. */ +export function runInTransaction( + ctx: DbContext, + work: (client: PoolClient) => Promise, +): Promise { + return withBranchPool(ctx, async (pool) => { + const client = await pool.connect(); + let failed: Error | undefined; + try { + return await inTransaction(client, () => work(client)); + } catch (error) { + failed = error instanceof Error ? error : new Error(String(error)); + throw error; + } finally { + // A connection whose transaction state is unknown is not reused. + client.release(failed); + } + }); +} diff --git a/packages/appkit/src/database/tooling/clean.ts b/packages/appkit/src/database/tooling/clean.ts new file mode 100644 index 000000000..ffba81e41 --- /dev/null +++ b/packages/appkit/src/database/tooling/clean.ts @@ -0,0 +1,219 @@ +import type { PoolClient } from "pg"; + +import { loadDefaultDatabaseSchema } from "../../plugins/database/load-schema"; +import { quoteIdentifier } from "../migrations/apply"; +import { runInTransaction } from "./branch-sql"; +import { type DbContext, NO_ENV_BRANCH, withBranchPool } from "./context"; +import { notYourBranch, resolveEnvBranch, writableEnvBranch } from "./guard"; +import { requireSchemaFile } from "./migrate"; +import { resourceId } from "./targets"; + +/** + * Drop the app schema with everything in it, the migration journal included. + * Schema public is never dropped: only the relations, types and routines the + * connected role owns go, and objects an extension installed stay. + */ +export function cleanStatements(schemaName: string): readonly string[] { + if (schemaName !== "public") { + return [`DROP SCHEMA IF EXISTS ${quoteIdentifier(schemaName)} CASCADE`]; + } + return [ + `DO $appkit$ +DECLARE + me oid := (SELECT oid FROM pg_catalog.pg_roles WHERE rolname = current_user); + ns oid := (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = 'public'); + r record; +BEGIN + FOR r IN + SELECT c.relname, c.relkind FROM pg_catalog.pg_class c + WHERE c.relnamespace = ns + AND c.relkind IN ('r', 'p', 'v', 'm', 'f', 'S') + AND c.relowner = me + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_depend d + WHERE d.classid = 'pg_catalog.pg_class'::regclass AND d.objid = c.oid + AND d.deptype IN ('a', 'i', 'e')) + LOOP + EXECUTE pg_catalog.format('DROP %s IF EXISTS public.%I CASCADE', + CASE r.relkind WHEN 'v' THEN 'VIEW' WHEN 'm' THEN 'MATERIALIZED VIEW' + WHEN 'f' THEN 'FOREIGN TABLE' WHEN 'S' THEN 'SEQUENCE' ELSE 'TABLE' END, + r.relname); + END LOOP; + FOR r IN + SELECT t.typname FROM pg_catalog.pg_type t + WHERE t.typnamespace = ns + AND t.typowner = me + AND (t.typtype IN ('e', 'd') + OR (t.typtype = 'c' AND EXISTS ( + SELECT 1 FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid AND c.relkind = 'c'))) + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_depend d + WHERE d.deptype = 'e' + AND ((d.classid = 'pg_catalog.pg_type'::regclass AND d.objid = t.oid) + OR (d.classid = 'pg_catalog.pg_class'::regclass AND d.objid = t.typrelid))) + LOOP + EXECUTE pg_catalog.format('DROP TYPE IF EXISTS public.%I CASCADE', r.typname); + END LOOP; + FOR r IN + SELECT p.proname, p.prokind, pg_catalog.pg_get_function_identity_arguments(p.oid) AS args + FROM pg_catalog.pg_proc p + WHERE p.pronamespace = ns + AND p.prokind IN ('f', 'p', 'a', 'w') + AND p.proowner = me + AND NOT EXISTS ( + SELECT 1 FROM pg_catalog.pg_depend d + WHERE d.classid = 'pg_catalog.pg_proc'::regclass AND d.objid = p.oid AND d.deptype = 'e') + LOOP + -- args is null when an earlier CASCADE already dropped the routine. + CONTINUE WHEN r.args IS NULL; + EXECUTE pg_catalog.format('DROP %s IF EXISTS public.%I(%s) CASCADE', + CASE r.prokind WHEN 'a' THEN 'AGGREGATE' WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, + r.proname, + CASE WHEN r.prokind = 'a' AND r.args = '' THEN '*' ELSE r.args END); + END LOOP; +END +$appkit$`, + ]; +} + +/** + * Objects outside what clean drops that depend on it, so CASCADE removes them + * too: a view in another schema, a foreign key into an app table, a column of + * an app type. For public, "outside" also means owned by another role. + */ +const DEPENDENTS_SQL = ` +with target as ( + select oid from pg_catalog.pg_namespace where nspname = $1 +), me as ( + select oid from pg_catalog.pg_roles where rolname = current_user +), dropped(classid, objid) as ( + select 'pg_catalog.pg_class'::regclass, c.oid from pg_catalog.pg_class c, target, me + where c.relnamespace = target.oid and (not $2 or c.relowner = me.oid) + union all + select 'pg_catalog.pg_type'::regclass, t.oid from pg_catalog.pg_type t, target, me + where t.typnamespace = target.oid and (not $2 or t.typowner = me.oid) + union all + select 'pg_catalog.pg_proc'::regclass, p.oid from pg_catalog.pg_proc p, target, me + where p.pronamespace = target.oid and (not $2 or p.proowner = me.oid) +) +select distinct o.label + from pg_catalog.pg_depend d + join dropped on d.refclassid = dropped.classid and d.refobjid = dropped.objid + cross join lateral ( + select c.relnamespace as namespace, c.relowner as owner, + pg_catalog.pg_describe_object(d.classid, d.objid, d.objsubid) as label + from pg_catalog.pg_class c + where d.classid = 'pg_catalog.pg_class'::regclass and c.oid = d.objid + union all + select c.relnamespace, c.relowner, + pg_catalog.pg_describe_object('pg_catalog.pg_class'::regclass, c.oid, 0) + from pg_catalog.pg_rewrite r join pg_catalog.pg_class c on c.oid = r.ev_class + where d.classid = 'pg_catalog.pg_rewrite'::regclass and r.oid = d.objid + union all + select con.connamespace, coalesce(c.relowner, t.typowner), + pg_catalog.pg_describe_object(d.classid, d.objid, 0) + from pg_catalog.pg_constraint con + left join pg_catalog.pg_class c on c.oid = con.conrelid + left join pg_catalog.pg_type t on t.oid = con.contypid + where d.classid = 'pg_catalog.pg_constraint'::regclass and con.oid = d.objid + union all + select p.pronamespace, p.proowner, pg_catalog.pg_describe_object(d.classid, d.objid, 0) + from pg_catalog.pg_proc p + where d.classid = 'pg_catalog.pg_proc'::regclass and p.oid = d.objid + union all + select t.typnamespace, t.typowner, pg_catalog.pg_describe_object(d.classid, d.objid, 0) + from pg_catalog.pg_type t + where d.classid = 'pg_catalog.pg_type'::regclass and t.oid = d.objid + union all + select c.relnamespace, c.relowner, pg_catalog.pg_describe_object(d.classid, d.objid, 0) + from pg_catalog.pg_trigger tg join pg_catalog.pg_class c on c.oid = tg.tgrelid + where d.classid = 'pg_catalog.pg_trigger'::regclass and tg.oid = d.objid + union all + select c.relnamespace, c.relowner, pg_catalog.pg_describe_object(d.classid, d.objid, 0) + from pg_catalog.pg_attrdef a join pg_catalog.pg_class c on c.oid = a.adrelid + where d.classid = 'pg_catalog.pg_attrdef'::regclass and a.oid = d.objid + union all + select c.relnamespace, c.relowner, pg_catalog.pg_describe_object(d.classid, d.objid, 0) + from pg_catalog.pg_policy pol join pg_catalog.pg_class c on c.oid = pol.polrelid + where d.classid = 'pg_catalog.pg_policy'::regclass and pol.oid = d.objid + ) o + cross join target + cross join me + where d.deptype = 'n' + and (o.namespace <> target.oid or ($2 and o.owner <> me.oid)) + order by 1`; + +export async function readCascadeDependents( + client: Pick, + schemaName: string, +): Promise { + const { rows } = await client.query<{ label: string }>(DEPENDENTS_SQL, [ + schemaName, + schemaName === "public", + ]); + return rows.map((row) => row.label); +} + +export interface CleanReport { + readonly branch: string; + readonly schemaName: string; + /** Objects outside the app schema that depend on it, which the drop removes too. */ + readonly dependents: readonly string[]; +} + +async function cleanTarget(ctx: DbContext) { + requireSchemaFile(ctx.root, "clean"); + const branch = await writableEnvBranch(ctx, "clean"); + const schema = await loadDefaultDatabaseSchema(ctx.root, { + showImportError: true, + }); + return { branch: branch.name, schemaName: schema.$schemaName }; +} + +/** What `appkit db clean` would drop, read without changing anything. */ +export async function previewClean(ctx: DbContext): Promise { + const target = await cleanTarget(ctx); + const dependents = await withBranchPool(ctx, (pool) => + readCascadeDependents(pool, target.schemaName), + ); + return { ...target, dependents }; +} + +/** + * `appkit db clean`: remove every app-schema object, the migration journal + * included, on your dev or session branch. Refuses any other branch before + * connecting. + */ +export async function runClean(ctx: DbContext): Promise { + const target = await cleanTarget(ctx); + const dependents = await runInTransaction(ctx, async (client) => { + const found = await readCascadeDependents(client, target.schemaName); + for (const statement of cleanStatements(target.schemaName)) { + await client.query(statement); + } + return found; + }); + return { ...target, dependents }; +} + +/** + * `appkit db clean --delete-branch`: delete this checkout's session branch. + * The dev branch is kept, so the next `appkit db init` finds it. + */ +export async function deleteSessionBranch( + ctx: DbContext, +): Promise<{ readonly branch: string }> { + const env = await resolveEnvBranch(ctx); + if (!env) throw new Error(`Did not delete the branch: ${NO_ENV_BRANCH}.`); + const { name } = env.branch; + if (env.kind === "dev") { + throw new Error( + `Refusing to delete ${resourceId(name)}: --delete-branch deletes session branches only, and your dev branch stays for the next \`appkit db init\`. Delete it with \`databricks postgres delete-branch ${name}\`.`, + ); + } + if (env.kind !== "session") { + throw new Error(`Refusing to delete the branch: ${notYourBranch(env)}.`); + } + await ctx.api().deleteBranch(name); + return { branch: name }; +} diff --git a/packages/appkit/src/database/tooling/context.ts b/packages/appkit/src/database/tooling/context.ts index b417297da..1f479f942 100644 --- a/packages/appkit/src/database/tooling/context.ts +++ b/packages/appkit/src/database/tooling/context.ts @@ -28,6 +28,8 @@ export interface ToolingWorkspaceClient { readonly currentUser: { me(request: { attributes?: string }): Promise<{ userName?: string }>; }; + /** The resolved workspace host, which init writes to a new .env. */ + getHost(): Promise; } /** Where a command runs: the app root, its environment and its Databricks identity. */ diff --git a/packages/appkit/src/database/tooling/drizzle-kit-install.ts b/packages/appkit/src/database/tooling/drizzle-kit-install.ts new file mode 100644 index 000000000..9fa9e9c36 --- /dev/null +++ b/packages/appkit/src/database/tooling/drizzle-kit-install.ts @@ -0,0 +1,91 @@ +import { spawnSync } from "node:child_process"; +import fs from "node:fs"; +import path from "node:path"; + +import { DRIZZLE_KIT_VERSION, hasDrizzleKit } from "./migrations"; + +/** Runs a command in `cwd`; throws when it fails. */ +export type RunCommand = ( + command: string, + args: readonly string[], + cwd: string, +) => void; + +/** Exact dev-dependency installs, by the lockfile that names the package manager. */ +const MANAGERS: ReadonlyArray< + readonly [lockfile: string, command: string, args: readonly string[]] +> = [ + ["pnpm-lock.yaml", "pnpm", ["add", "-D", "-E"]], + ["yarn.lock", "yarn", ["add", "-D", "-E"]], + ["bun.lock", "bun", ["add", "-d", "--exact"]], + ["bun.lockb", "bun", ["add", "-d", "--exact"]], + ["package-lock.json", "npm", ["install", "--save-dev", "--save-exact"]], +]; + +interface Install { + readonly command: string; + readonly args: readonly string[]; +} + +/** The install with the package manager whose lockfile is nearest `root`; npm without one. */ +export function drizzleKitInstall(root: string): Install { + const spec = `drizzle-kit@${DRIZZLE_KIT_VERSION}`; + for (let dir = path.resolve(root); ; dir = path.dirname(dir)) { + for (const [lockfile, command, args] of MANAGERS) { + if (fs.existsSync(path.join(dir, lockfile))) { + return { command, args: [...args, spec] }; + } + } + if (dir === path.dirname(dir)) break; + } + return { + command: "npm", + args: ["install", "--save-dev", "--save-exact", spec], + }; +} + +export const runCommand: RunCommand = (command, args, cwd) => { + const result = spawnSync(command, args, { + cwd, + encoding: "utf-8", + stdio: "pipe", + // Windows runs package managers through their .cmd shims. + shell: process.platform === "win32", + }); + if (result.error) throw result.error; + if (result.status !== 0) { + const output = (result.stderr || result.stdout || "").trim(); + throw new Error( + `${[command, ...args].join(" ")} exited with code ${result.status}${output ? `: ${output}` : ""}`, + ); + } +}; + +/** + * Install drizzle-kit as an exact dev dependency when the app cannot resolve + * it. Returns the command it ran; null when drizzle-kit was there, or there is + * no package.json to add it to. + */ +export function ensureDrizzleKit( + root: string, + run: RunCommand, + onStep: (message: string) => void, +): string | null { + if (hasDrizzleKit(root) || !fs.existsSync(path.join(root, "package.json"))) { + return null; + } + const { command, args } = drizzleKitInstall(root); + const line = [command, ...args].join(" "); + onStep( + `Installing drizzle-kit ${DRIZZLE_KIT_VERSION}, which writes migrations (${line})`, + ); + try { + run(command, args, root); + } catch (error) { + throw new Error( + `Installing drizzle-kit failed. Run \`${line}\`, then rerun \`appkit db init\`.`, + { cause: error }, + ); + } + return line; +} diff --git a/packages/appkit/src/database/tooling/env-file.ts b/packages/appkit/src/database/tooling/env-file.ts new file mode 100644 index 000000000..63eccdb18 --- /dev/null +++ b/packages/appkit/src/database/tooling/env-file.ts @@ -0,0 +1,124 @@ +import { randomBytes } from "node:crypto"; +import fs from "node:fs"; +import path from "node:path"; + +/** An entry's first line: indentation and `export ` (kept on rewrite), the key, the value. */ +const ENTRY = /^(\s*(?:export\s+)?)([\w.-]+)\s*=\s*(.*)$/; +const QUOTES = new Set(['"', "'", "`"]); +const PLAIN_VALUE = /^[A-Za-z0-9_./:@+-]*$/; + +/** + * A value dotenv reads back unchanged, always on one line: plain; else + * double-quoted, where dotenv turns `\n` and `\r` back into line breaks; else + * single-quoted, which dotenv takes literally. + */ +function formatValue(key: string, value: string): string { + if (PLAIN_VALUE.test(value)) return value; + if (!/["\\]/.test(value)) { + return `"${value.replace(/\r/g, "\\r").replace(/\n/g, "\\n")}"`; + } + if (!/['\r\n]/.test(value)) return `'${value}'`; + throw new Error( + `Cannot write ${key} to .env: its value mixes quotes, backslashes and line breaks.`, + ); +} + +/** Like dotenv, a quote preceded by a backslash does not close the value. */ +function closes(text: string, quote: string, from: number): boolean { + for (let i = from; i < text.length; i++) { + if (text[i] === quote && text[i - 1] !== "\\") return true; + } + return false; +} + +interface Entry { + /** null for comments, blank lines and lines dotenv does not read as `KEY=`. */ + readonly key: string | null; + readonly prefix: string; + readonly lines: readonly string[]; +} + +/** A quoted value without its closing quote continues on the next lines, as dotenv reads it. */ +function entries(lines: readonly string[]): Entry[] { + const out: Entry[] = []; + for (let i = 0; i < lines.length; i++) { + const line = lines[i] as string; + const match = ENTRY.exec(line); + if (!match) { + out.push({ key: null, prefix: "", lines: [line] }); + continue; + } + const value = match[3] ?? ""; + const quote = value[0] ?? ""; + let end = i; + if (QUOTES.has(quote) && !closes(value, quote, 1)) { + const closing = lines.findIndex( + (next, j) => j > i && closes(next, quote, 0), + ); + if (closing !== -1) end = closing; + } + out.push({ + key: match[2] ?? null, + prefix: match[1] ?? "", + lines: lines.slice(i, end + 1), + }); + i = end; + } + return out; +} + +export function refuseSymlink(file: string): void { + if (!fs.lstatSync(file, { throwIfNoEntry: false })?.isSymbolicLink()) return; + throw new Error( + `Refusing to rewrite ${file}: it is a symbolic link to ${fs.readlinkSync(file)}, and writing through it would change that file. Replace it with a regular file, then rerun.`, + ); +} + +/** A crash leaves the old file or the new one, never a truncated one. */ +function writeAtomically(file: string, contents: string): void { + const temp = path.join( + path.dirname(file), + `.${path.basename(file)}.${process.pid}.${randomBytes(4).toString("hex")}.tmp`, + ); + try { + fs.writeFileSync(temp, contents, { mode: 0o600, flag: "wx" }); + fs.renameSync(temp, file); + } catch (error) { + fs.rmSync(temp, { force: true }); + throw error; + } +} + +/** + * Replace `KEY=…` entries and append missing keys. Comments, blank lines, + * other keys, their order, `export ` prefixes and CRLF line endings are kept. + * Written 0600 through a temporary file; a symbolic link is refused. + */ +export function updateEnvFile( + file: string, + values: Readonly>, +): void { + refuseSymlink(file); + const text = fs.existsSync(file) ? fs.readFileSync(file, "utf-8") : ""; + const eol = text.includes("\r\n") ? "\r\n" : "\n"; + const lines = + text.length > 0 ? text.replace(/\r?\n$/, "").split(/\r?\n/) : []; + + const written = new Set(); + const out: string[] = []; + for (const entry of entries(lines)) { + const { key } = entry; + if (key === null || !Object.hasOwn(values, key)) { + out.push(...entry.lines); + continue; + } + // Later duplicates of a managed key are dropped, so the value is unambiguous. + if (written.has(key)) continue; + written.add(key); + out.push(`${entry.prefix}${key}=${formatValue(key, values[key] ?? "")}`); + } + for (const [key, value] of Object.entries(values)) { + if (!written.has(key)) out.push(`${key}=${formatValue(key, value)}`); + } + writeAtomically(file, `${out.join(eol)}${eol}`); +} diff --git a/packages/appkit/src/database/tooling/guard.ts b/packages/appkit/src/database/tooling/guard.ts index 994fe1fac..fc55e2385 100644 --- a/packages/appkit/src/database/tooling/guard.ts +++ b/packages/appkit/src/database/tooling/guard.ts @@ -1,4 +1,4 @@ -import { type DbContext, hasEnvBranch } from "./context"; +import { type DbContext, hasEnvBranch, NO_ENV_BRANCH } from "./context"; import { currentIdentity } from "./identity"; import type { LakebaseBranch } from "./lakebase-api"; import { devBranchId, sessionBranchId } from "./naming"; @@ -46,15 +46,12 @@ export function classifyBranch(input: { } const id = resourceId(branch.name); if (id === devBranchId(userName)) return { kind: "dev" }; - if (sessionId && id === sessionBranchId(sessionId)) { - return { kind: "session" }; - } + const session = sessionId ? sessionBranchId(userName, sessionId) : null; + if (id === session) return { kind: "session" }; return { kind: "other", reason: `it is not your dev branch (${devBranchId(userName)})${ - sessionId - ? ` or this checkout's session branch (${sessionBranchId(sessionId)})` - : "" + session ? ` or this checkout's session branch (${session})` : "" }`, }; } @@ -106,3 +103,16 @@ export async function resolveEnvBranch( }), }; } + +/** The .env branch, for a command that only runs on your dev or session branch. */ +export async function writableEnvBranch( + ctx: DbContext, + action: string, +): Promise { + const env = await resolveEnvBranch(ctx); + if (!env) throw new Error(`Did not ${action}: ${NO_ENV_BRANCH}.`); + if (!takesWrites(env.kind)) { + throw new Error(`Refusing to ${action}: ${notYourBranch(env)}.`); + } + return env.branch; +} diff --git a/packages/appkit/src/database/tooling/hand-off.ts b/packages/appkit/src/database/tooling/hand-off.ts new file mode 100644 index 000000000..adb3d17ae --- /dev/null +++ b/packages/appkit/src/database/tooling/hand-off.ts @@ -0,0 +1,104 @@ +import type { PoolClient } from "pg"; + +import { schemaOwnership } from "../migrations/ownership"; +import { type DbContext, withBranchPool } from "./context"; +import type { LakebaseRole } from "./lakebase-api"; +import { resourceId } from "./targets"; + +/** + * What keeps a role from being dropped once its objects in the current + * database are reassigned: objects it owns in other databases, shared objects + * (databases), and privileges granted to it anywhere. + */ +const BLOCKERS_SQL = ` +select coalesce(d.datname, '') as database, count(*)::int as count + from pg_catalog.pg_shdepend s + left join pg_catalog.pg_database d on d.oid = s.dbid + where s.refclassid = 'pg_catalog.pg_authid'::regclass + and s.refobjid = (select oid from pg_catalog.pg_roles where rolname = $1) + and not (s.deptype = 'o' and s.dbid = (select oid from pg_catalog.pg_database + where datname = current_database())) + group by 1 + order by 1`; + +interface RoleBlocker { + /** Empty for shared objects. */ + readonly database: string; + readonly count: number; +} + +export async function readRoleBlockers( + client: Pick, + role: string, +): Promise { + const { rows } = await client.query(BLOCKERS_SQL, [role]); + return rows; +} + +function describeBlockers(blockers: readonly RoleBlocker[]): string { + return blockers + .map( + ({ database, count }) => + `${database ? `database ${database}` : "shared objects"} (${count})`, + ) + .join(", "); +} + +interface HandOff { + /** The branch .env points at. */ + readonly branch: string; + readonly database: string; + readonly schemaName: string; + /** Your role on the branch, which receives the objects. */ + readonly developer: LakebaseRole; + readonly roles: readonly LakebaseRole[]; + readonly onStep: (message: string) => void; +} + +/** + * A fork copies its parent's objects with their owners, so on a fork of the + * deploy branch the app's service principal owns the app schema. Hand it to + * you by deleting that owner's copy of its role on this branch, with its + * objects reassigned to you. Only the app schema's owner is deleted. Returns + * that owner; null when you already own the schema or it does not exist. + */ +export async function handOffAppSchema( + ctx: DbContext, + input: HandOff, +): Promise { + const found = await withBranchPool(ctx, async (pool) => { + const { owner, currentUser } = await schemaOwnership( + pool, + input.schemaName, + ); + if (!owner || owner === currentUser) return null; + return { owner, blockers: await readRoleBlockers(pool, owner) }; + }); + if (!found) return null; + + const { owner, blockers } = found; + const where = `schema ${input.schemaName} in database ${input.database} on ${resourceId(input.branch)}`; + const role = input.roles.find((r) => r.status?.postgres_role === owner); + if (!role) { + throw new Error( + `Cannot hand ${where} to you: it belongs to "${owner}", which is not a Lakebase-managed role. .env already points at the branch. Have "${owner}" transfer the schema and its objects to you (ALTER … OWNER TO), or set another schemaName in schema.ts, then rerun \`appkit db init\`.`, + ); + } + input.onStep( + `Handing ${where} to you: deleting this branch's copy of role ${owner}, which owns it, with its objects reassigned to you`, + ); + try { + await ctx + .api() + .deleteRole(role.name, { reassignOwnedTo: input.developer.name }); + } catch (error) { + const reason = error instanceof Error ? error.message : String(error); + throw new Error( + blockers.length > 0 + ? `Could not hand ${where} to you: deleting role "${owner}" failed (${reason}). These objects and privileges of "${owner}" block the delete: ${describeBlockers(blockers)}. .env already points at the branch. Drop them or reassign them to you, then rerun \`appkit db init\`.` + : `Could not hand ${where} to you: deleting role "${owner}" failed (${reason}). .env already points at the branch; rerun \`appkit db init\` to retry.`, + { cause: error }, + ); + } + return owner; +} diff --git a/packages/appkit/src/database/tooling/identity.ts b/packages/appkit/src/database/tooling/identity.ts index e457c79a2..c0976b29c 100644 --- a/packages/appkit/src/database/tooling/identity.ts +++ b/packages/appkit/src/database/tooling/identity.ts @@ -1,8 +1,17 @@ import type { DbContext } from "./context"; +import type { IdentityType } from "./lakebase-api"; /** Who runs the command: the Postgres role name is the user name (email or SP client ID). */ -interface Identity { +export interface Identity { readonly userName: string; + readonly identityType: IdentityType; +} + +const UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i; + +/** A service principal's user name is its application (client) ID. */ +export function identityTypeOf(userName: string): IdentityType { + return UUID.test(userName) ? "SERVICE_PRINCIPAL" : "USER"; } const identities = new WeakMap>(); @@ -20,7 +29,7 @@ export function currentIdentity(ctx: DbContext): Promise { "Could not determine the current Databricks user (currentUser.me returned no userName). Check DATABRICKS_CONFIG_PROFILE.", ); } - return { userName }; + return { userName, identityType: identityTypeOf(userName) }; }); identity.catch(() => identities.delete(ctx)); identities.set(ctx, identity); diff --git a/packages/appkit/src/database/tooling/index.ts b/packages/appkit/src/database/tooling/index.ts index 29adba529..c355b52c2 100644 --- a/packages/appkit/src/database/tooling/index.ts +++ b/packages/appkit/src/database/tooling/index.ts @@ -6,10 +6,25 @@ * Internal to the AppKit CLI: it has no semver guarantee, and only the CLI of * the same AppKit release should import it. */ +export { + type CleanReport, + deleteSessionBranch, + previewClean, + runClean, +} from "./clean"; export { createDbContext, type DbContext } from "./context"; +export { + type InitOptions, + type InitReport, + type InitStart, + runInit, + type SchemaStart, +} from "./init"; export { type IntrospectOptions, type IntrospectReport, runIntrospect, } from "./introspect"; export { type MigrateOptions, type MigrateReport, runMigrate } from "./migrate"; +export { runSeed } from "./seed"; +export type { BranchMode } from "./targets"; diff --git a/packages/appkit/src/database/tooling/init.ts b/packages/appkit/src/database/tooling/init.ts new file mode 100644 index 000000000..6a5118980 --- /dev/null +++ b/packages/appkit/src/database/tooling/init.ts @@ -0,0 +1,572 @@ +import fs from "node:fs"; +import path from "node:path"; + +import dotenv from "dotenv"; + +import { loadDefaultDatabaseSchema } from "../../plugins/database/load-schema"; +import { DEFAULT_SCHEMA_NAME } from "../schema-builder/define-schema"; +import { type CleanReport, runClean } from "./clean"; +import type { DbContext } from "./context"; +import { + ensureDrizzleKit, + type RunCommand, + runCommand, +} from "./drizzle-kit-install"; +import { refuseSymlink, updateEnvFile } from "./env-file"; +import { classifyBranch, takesWrites } from "./guard"; +import { handOffAppSchema } from "./hand-off"; +import { currentIdentity, type Identity } from "./identity"; +import { type IntrospectReport, runIntrospect } from "./introspect"; +import { + isAlreadyExists, + isBranchLimit, + isPermissionDenied, + type LakebaseBranch, + type LakebaseEndpoint, + type LakebaseRole, + READ_WRITE_ENDPOINT, + SUPERUSER_MEMBERSHIP, +} from "./lakebase-api"; +import { + hasSchemaFile, + type MigrateReport, + requireSchemaFile, + runMigrate, +} from "./migrate"; +import { + devBranchId, + isUsersBranchId, + newSessionId, + parseTtl, + SESSION_BRANCH_TTL, + sessionBranchId, +} from "./naming"; +import { + type BranchMode, + type DeployTarget, + type LocalTarget, + parseResourceName, + readDeployBranches, + readDeployTarget, + readLocalTarget, + resourceId, +} from "./targets"; + +/** How init leaves the schema; "none" leaves it as the fork has it. */ +export type SchemaStart = "migrate" | "introspect" | "clean" | "none"; + +export interface InitOptions { + /** Default: this checkout's session branch when .env points at it, else your dev branch. */ + readonly mode?: BranchMode; + /** Default: migrate when schema.ts exists, introspect otherwise. */ + readonly schema?: SchemaStart; + /** A new session branch's TTL as `--ttl` takes it (30m, 1h, 7d); default seven days. */ + readonly ttl?: string; + /** Passed to migrate and introspect: false makes an ambiguous rename fail instead of prompting. */ + readonly interactive?: boolean; + /** One line per step that calls the Lakebase API or changes a file. */ + readonly onStep?: (message: string) => void; + /** Runs the drizzle-kit install. */ + readonly runCommand?: RunCommand; +} + +type RoleAction = "kept" | "created" | "updated"; + +export type InitStart = + | { readonly kind: "none" } + | { readonly kind: "migrate"; readonly migrate: MigrateReport } + | { readonly kind: "introspect"; readonly introspect: IntrospectReport } + | { + readonly kind: "clean"; + readonly clean: CleanReport; + readonly migrate: MigrateReport; + }; + +export interface InitReport { + readonly mode: BranchMode; + readonly branch: string; + readonly created: boolean; + /** The branch a new fork was created from. */ + readonly sourceBranch?: string; + /** When the branch expires; null for a branch without expiry. */ + readonly expireTime: string | null; + readonly endpoint: string; + readonly endpointCreated: boolean; + readonly host: string; + readonly database: string; + readonly user: string; + readonly role: { readonly name: string; readonly action: RoleAction }; + /** The Postgres role whose copy of the app schema was handed to you; null when you owned it. */ + readonly handedOff: string | null; + readonly envFile: string; + /** The drizzle-kit install init ran; null when it was installed. */ + readonly installed: string | null; + readonly start: InitStart; +} + +const ENV_FILE = ".env"; +const DEFAULT_ENDPOINT_ID = "primary"; +/** An idle compute counts against the project's 20 active ones. */ +const SUSPEND_TIMEOUT = "300s"; +/** Every Lakebase branch has it; used when nothing names another database. */ +const DEFAULT_DATABASE = "databricks_postgres"; + +/** LAKEBASE_PROJECT, else the project of LAKEBASE_ENDPOINT, else of the bundle's postgres_branch. */ +function resolveProject( + local: LocalTarget, + deploy: DeployTarget | null, +): string { + const project = + local.project ?? + parseResourceName(local.endpoint ?? undefined)?.project ?? + parseResourceName(deploy?.branch)?.project; + if (!project) { + throw new Error( + "No Lakebase project configured. Set LAKEBASE_PROJECT=projects/ in .env, or bind postgres_branch in databricks.yml (`databricks apps init` does).", + ); + } + return project; +} + +/** Drop .env connection values and a deploy binding that belong to another project than LAKEBASE_PROJECT. */ +function scopeToProject( + project: string, + local: LocalTarget, + deploy: DeployTarget | null, +): { local: LocalTarget; deploy: DeployTarget | null } { + const inProject = (name: string | null | undefined) => + parseResourceName(name ?? undefined)?.project === project; + return { + local: inProject(local.endpoint) + ? local + : { ...local, endpoint: null, database: null }, + deploy: deploy && inProject(deploy.branch) ? deploy : null, + }; +} + +/** Session when .env points at this checkout's session branch; dev otherwise. */ +function defaultMode(local: LocalTarget, userName: string): BranchMode { + const branch = parseResourceName(local.endpoint ?? undefined)?.branch; + return local.session && + branch && + resourceId(branch) === sessionBranchId(userName, local.session) + ? "session" + : "dev"; +} + +function isGone(branch: LakebaseBranch, now = Date.now()): boolean { + if (branch.status?.current_state === "DELETED") return true; + const expires = branch.status?.expire_time; + return !!expires && Date.parse(expires) <= now; +} + +interface BranchRequest { + readonly project: string; + readonly mode: BranchMode; + readonly identity: Identity; + readonly session: string | null; + /** The API's duration ("3600s"), and how the developer wrote it. */ + readonly ttl: { readonly api: string; readonly label: string }; + readonly deployBranches: readonly string[]; + readonly envFile: string; + readonly onStep: (message: string) => void; +} + +interface FoundBranch { + readonly branch: LakebaseBranch; + readonly created: boolean; + readonly sourceBranch?: string; + readonly session: string | null; +} + +function branchLimitError( + project: string, + branches: readonly LakebaseBranch[], + userName: string, + cause: unknown, +): Error { + const yours = branches.filter( + (b) => isUsersBranchId(userName, resourceId(b.name)) && !isGone(b), + ); + const message = + yours.length > 0 + ? `Project ${project} is at its branch limit, so init created nothing. Delete one of your other branches, then rerun \`appkit db init\`:\n${yours + .map( + (b) => + ` databricks postgres delete-branch ${b.name}${b.status?.expire_time ? ` # expires ${b.status.expire_time}` : ""}`, + ) + .join("\n")}` + : `Project ${project} is at its branch limit, so init created nothing, and none of its branches are yours. Ask a project owner to delete unused branches, then rerun \`appkit db init\`.`; + return new Error(message, { cause }); +} + +/** + * Your dev branch or this checkout's session branch: found with one GET, or + * forked from the project's default branch. The branch id is resolved and + * checked against databricks.yml before any API call or .env write. An expired + * session branch gets a new session id; an expired dev branch is an error. + */ +async function findOrCreateBranch( + ctx: DbContext, + request: BranchRequest, +): Promise { + const { project, mode, identity, deployBranches, onStep } = request; + const api = ctx.api(); + let session = mode === "session" ? (request.session ?? newSessionId()) : null; + let rotated = false; + const branchName = () => + `${project}/branches/${ + session === null + ? devBranchId(identity.userName) + : sessionBranchId(identity.userName, session) + }`; + + const reuse = (branch: LakebaseBranch): FoundBranch => { + const verdict = classifyBranch({ + branch, + userName: identity.userName, + sessionId: session, + deployBranches, + }); + if (!takesWrites(verdict.kind)) { + throw new Error( + `Refusing to set up ${resourceId(branch.name)}: ${verdict.reason}. \`appkit db init\` only sets up your dev branch or this checkout's session branch.`, + ); + } + onStep(`Using ${resourceId(branch.name)}`); + return { branch, created: false, session }; + }; + const replaceGone = (name: string): void => { + if (session === null) { + throw new Error( + `Your dev branch ${name} was deleted or expired. Restore it, or delete it for good, then rerun \`appkit db init\`. \`appkit db init --session\` sets up a session branch meanwhile.`, + ); + } + if (rotated) { + throw new Error( + `Could not create a session branch: ${name} was deleted or expired too. Rerun \`appkit db init --session\`.`, + ); + } + rotated = true; + session = newSessionId(); + }; + + for (;;) { + const name = branchName(); + if (deployBranches.includes(name)) { + throw new Error( + `Refusing to set up ${resourceId(name)}: it is a branch databricks.yml deploys to, and the deployed app migrates it itself.`, + ); + } + const existing = await api.getBranch(name); + if (existing && !isGone(existing)) return reuse(existing); + if (existing) { + replaceGone(name); + continue; + } + + const branches = await api.listBranches(project); + const parent = branches.find((b) => b.status?.default); + if (!parent) { + throw new Error(`Project ${project} has no default branch to fork from.`); + } + if (session !== null && session !== request.session) { + // Written first, so a failed or interrupted run reuses this id instead of forking again. + updateEnvFile(request.envFile, { APPKIT_DB_SESSION: session }); + } + onStep( + `Creating ${resourceId(name)} from ${resourceId(parent.name)}${session === null ? "" : `, expiring after ${request.ttl.label}`}`, + ); + try { + const branch = await api.createBranch(project, resourceId(name), { + source_branch: parent.name, + ...(session === null + ? { no_expiry: true as const } + : { ttl: request.ttl.api }), + }); + return { + branch: branch?.name ? branch : { name }, + created: true, + sourceBranch: parent.name, + session, + }; + } catch (error) { + if (isPermissionDenied(error)) { + throw new Error( + `Creating ${resourceId(name)} needs CAN MANAGE on ${project}. Ask a project owner for it, then rerun \`appkit db init\`.`, + { cause: error }, + ); + } + if (isBranchLimit(error)) { + throw branchLimitError(project, branches, identity.userName, error); + } + if (!isAlreadyExists(error)) throw error; + // Another run created it first, or the name is held by a deleted branch. + const raced = await api.getBranch(name); + if (raced && !isGone(raced)) return reuse(raced); + replaceGone(name); + } + } +} + +async function ensureReadWriteEndpoint( + ctx: DbContext, + branch: string, + onStep: (message: string) => void, +): Promise<{ endpoint: LakebaseEndpoint; created: boolean; host: string }> { + const api = ctx.api(); + const readWrite = (endpoints: LakebaseEndpoint[]) => + endpoints.find((e) => e.status?.endpoint_type === READ_WRITE_ENDPOINT); + + let endpoint = readWrite(await api.listEndpoints(branch)); + let created = false; + if (!endpoint) { + onStep( + `Creating a read-write endpoint on ${resourceId(branch)} that suspends after 5 idle minutes`, + ); + endpoint = await api.createEndpoint(branch, DEFAULT_ENDPOINT_ID, { + suspendTimeout: SUSPEND_TIMEOUT, + }); + created = true; + } + let host = endpoint?.status?.hosts?.host; + if (!host) { + endpoint = readWrite(await api.listEndpoints(branch)) ?? endpoint; + host = endpoint?.status?.hosts?.host; + } + if (!endpoint?.name || !host) { + throw new Error( + `The read-write endpoint of ${branch} has no host yet. Rerun \`appkit db init\` in a minute.`, + ); + } + return { endpoint, created, host }; +} + +async function ensureDeveloperRole( + ctx: DbContext, + branch: string, + identity: Identity, + roles: readonly LakebaseRole[], + onStep: (message: string) => void, +): Promise<{ role: LakebaseRole; action: RoleAction }> { + const api = ctx.api(); + const existing = roles.find( + (r) => r.status?.postgres_role === identity.userName, + ); + if (!existing) { + onStep( + `Creating your role ${identity.userName} with ${SUPERUSER_MEMBERSHIP}`, + ); + const role = await api.createRole(branch, { + identity_type: identity.identityType, + postgres_role: identity.userName, + auth_method: "LAKEBASE_OAUTH_V1", + membership_roles: [SUPERUSER_MEMBERSHIP], + }); + if (role?.name) return { role, action: "created" }; + // The operation response normally carries the role; fall back to a lookup. + const found = (await api.listRoles(branch)).find( + (r) => r.status?.postgres_role === identity.userName, + ); + if (!found) { + throw new Error(`Creating your role on ${branch} did not return it.`); + } + return { role: found, action: "created" }; + } + if (existing.status?.membership_roles?.includes(SUPERUSER_MEMBERSHIP)) { + return { role: existing, action: "kept" }; + } + onStep(`Adding ${SUPERUSER_MEMBERSHIP} to your role ${identity.userName}`); + const role = await api.setRoleMemberships(existing.name, [ + SUPERUSER_MEMBERSHIP, + ]); + return { role: role?.name ? role : existing, action: "updated" }; +} + +async function resolveDatabase( + ctx: DbContext, + branch: string, + local: LocalTarget, + deploy: DeployTarget | null, +): Promise { + // A real .env endpoint came with its database; forks inherit the parent's databases. + if (local.database) return local.database; + if (deploy?.database) { + const database = await ctx + .api() + .getDatabase(`${branch}/databases/${resourceId(deploy.database)}`); + if (database?.status?.postgres_database) + return database.status.postgres_database; + } + return DEFAULT_DATABASE; +} + +/** The profile or host init authenticated with, for a .env that names neither. */ +async function authValues( + ctx: DbContext, + current: Record, +): Promise> { + if (current.DATABRICKS_CONFIG_PROFILE || current.DATABRICKS_HOST) return {}; + const profile = ctx.env.DATABRICKS_CONFIG_PROFILE; + if (profile) return { DATABRICKS_CONFIG_PROFILE: profile }; + const host = ctx.env.DATABRICKS_HOST || (await ctx.workspace().getHost()); + return host ? { DATABRICKS_HOST: host } : {}; +} + +function readEnvValues(file: string): Record { + return fs.existsSync(file) + ? dotenv.parse(fs.readFileSync(file, "utf-8")) + : {}; +} + +/** + * `appkit db init`: find or create your dev branch or this checkout's session + * branch, ensure a read-write endpoint and your DATABRICKS_SUPERUSER role, + * point `.env` at the branch, hand the app schema the fork copied to you, then + * start the schema. No prompts: every choice has a default. It refuses the + * project's default branch, protected branches and branches databricks.yml + * deploys to before changing anything, and never touches the deploy bundle. + */ +export async function runInit( + ctx: DbContext, + options: InitOptions = {}, +): Promise { + const onStep = options.onStep ?? (() => undefined); + const interactive = options.interactive ?? false; + const ttl = + options.ttl === undefined + ? { api: SESSION_BRANCH_TTL, label: "7 days" } + : { api: parseTtl(options.ttl), label: options.ttl.trim() }; + + // Before any change: a .env init may write, what the start needs, and the schema the hand-off is for. + const envFile = path.join(ctx.root, ENV_FILE); + refuseSymlink(envFile); + const hasSchema = hasSchemaFile(ctx.root); + const startKind = options.schema ?? (hasSchema ? "migrate" : "introspect"); + if (startKind === "migrate" || startKind === "clean") { + requireSchemaFile(ctx.root, startKind); + } + const schemaName = + hasSchema && startKind !== "introspect" + ? (await loadDefaultDatabaseSchema(ctx.root, { showImportError: true })) + .$schemaName + : DEFAULT_SCHEMA_NAME; + + const configured = { + local: readLocalTarget(ctx.env), + deploy: readDeployTarget(ctx.root), + }; + const project = resolveProject(configured.local, configured.deploy); + const { local, deploy } = scopeToProject( + project, + configured.local, + configured.deploy, + ); + const identity = await currentIdentity(ctx); + const mode = options.mode ?? defaultMode(local, identity.userName); + if (options.ttl !== undefined && mode !== "session") { + throw new Error("--ttl applies to session branches only; add --session."); + } + const before = readEnvValues(envFile); + + const found = await findOrCreateBranch(ctx, { + project, + mode, + identity, + session: local.session, + ttl, + deployBranches: readDeployBranches(ctx.root), + envFile, + onStep, + }); + const { branch } = found; + + const { + endpoint, + created: endpointCreated, + host, + } = await ensureReadWriteEndpoint(ctx, branch.name, onStep); + const roles = await ctx.api().listRoles(branch.name); + const { role, action } = await ensureDeveloperRole( + ctx, + branch.name, + identity, + roles, + onStep, + ); + const database = await resolveDatabase(ctx, branch.name, local, deploy); + + const values: Record = { + LAKEBASE_PROJECT: project, + PGHOST: host, + LAKEBASE_ENDPOINT: endpoint.name, + PGDATABASE: database, + PGUSER: identity.userName, + ...(found.session === null ? {} : { APPKIT_DB_SESSION: found.session }), + ...(before.PGPORT && /^\d+$/.test(before.PGPORT) ? {} : { PGPORT: "5432" }), + ...(before.PGSSLMODE ? {} : { PGSSLMODE: "require" }), + ...(await authValues(ctx, before)), + }; + onStep(`Pointing .env at ${resourceId(branch.name)}, database ${database}`); + updateEnvFile(envFile, values); + const branchCtx: DbContext = { ...ctx, env: { ...ctx.env, ...values } }; + + const handedOff = await handOffAppSchema(branchCtx, { + branch: branch.name, + database, + schemaName, + developer: role, + roles, + onStep, + }); + + const installed = + startKind === "migrate" || startKind === "introspect" + ? ensureDrizzleKit(ctx.root, options.runCommand ?? runCommand, onStep) + : null; + + let start: InitStart; + const on = `on ${resourceId(branch.name)}`; + if (startKind === "migrate") { + onStep(`Migrating schema ${schemaName} ${on}`); + start = { + kind: "migrate", + migrate: await runMigrate(branchCtx, { interactive }), + }; + } else if (startKind === "introspect") { + onStep(`Writing schema.ts from schema ${schemaName} ${on}`); + start = { + kind: "introspect", + introspect: await runIntrospect(branchCtx, { interactive }), + }; + } else if (startKind === "clean") { + onStep(`Dropping schema ${schemaName} ${on}`); + const clean = await runClean(branchCtx); + // Recreate from the committed migrations, so the startup check passes. + start = { + kind: "clean", + clean, + migrate: await runMigrate(branchCtx, { interactive, generate: false }), + }; + } else { + start = { kind: "none" }; + } + + return { + mode, + branch: branch.name, + created: found.created, + ...(found.sourceBranch ? { sourceBranch: found.sourceBranch } : {}), + expireTime: branch.status?.expire_time ?? null, + endpoint: endpoint.name, + endpointCreated, + host, + database, + user: identity.userName, + role: { name: role.name, action }, + handedOff, + envFile, + installed, + start, + }; +} diff --git a/packages/appkit/src/database/tooling/lakebase-api.ts b/packages/appkit/src/database/tooling/lakebase-api.ts index e020e28ca..f4fc45fda 100644 --- a/packages/appkit/src/database/tooling/lakebase-api.ts +++ b/packages/appkit/src/database/tooling/lakebase-api.ts @@ -18,19 +18,61 @@ export interface LakebaseBranch { export interface LakebaseEndpoint { name: string; - status?: { hosts?: { host?: string } }; + status?: { + endpoint_type?: string; + hosts?: { host?: string }; + }; +} + +export interface LakebaseRole { + name: string; + status?: { + postgres_role?: string; + membership_roles?: string[]; + }; +} + +export interface LakebaseDatabase { + name: string; + status?: { postgres_database?: string }; +} + +export type IdentityType = "USER" | "SERVICE_PRINCIPAL"; + +export const SUPERUSER_MEMBERSHIP = "DATABRICKS_SUPERUSER"; +export const READ_WRITE_ENDPOINT = "ENDPOINT_TYPE_READ_WRITE"; + +export interface RoleSpec { + identity_type: IdentityType; + postgres_role: string; + auth_method: "LAKEBASE_OAUTH_V1"; + membership_roles: [typeof SUPERUSER_MEMBERSHIP]; +} + +export type BranchSpec = { source_branch: string } & ( + | { no_expiry: true } + | { ttl: string } +); + +interface Operation { + name?: string; + done?: boolean; + error?: { message?: string; error_code?: string }; + response?: unknown; } export class LakebaseApiError extends Error { - /** HTTP status of the failing call. */ + /** HTTP status of the failing call; absent for a failed long-running operation. */ readonly statusCode?: number; + /** The API's error code, e.g. ALREADY_EXISTS, from the call or the operation. */ + readonly errorCode?: string; /** What was being done, e.g. "create branch projects/p/branches/b". */ readonly operation: string; constructor( operation: string, message: string, - options: { statusCode?: number; cause?: unknown } = {}, + options: { statusCode?: number; errorCode?: string; cause?: unknown } = {}, ) { super(`Lakebase: ${operation} failed: ${message}`, { cause: options.cause, @@ -38,9 +80,37 @@ export class LakebaseApiError extends Error { this.name = "LakebaseApiError"; this.operation = operation; this.statusCode = options.statusCode; + this.errorCode = options.errorCode; } } +function failedWith( + error: unknown, + statusCode: number, + errorCodes: readonly string[], +): boolean { + return ( + error instanceof LakebaseApiError && + (error.statusCode === statusCode || + (error.errorCode !== undefined && errorCodes.includes(error.errorCode))) + ); +} + +export const isAlreadyExists = (error: unknown) => + failedWith(error, 409, ["ALREADY_EXISTS", "RESOURCE_ALREADY_EXISTS"]); +export const isPermissionDenied = (error: unknown) => + failedWith(error, 403, ["PERMISSION_DENIED"]); +const isNotFound = (error: unknown) => + failedWith(error, 404, ["NOT_FOUND", "RESOURCE_DOES_NOT_EXIST"]); + +/** The project holds as many branches as it may: "unarchived branches limit exceeded". */ +export function isBranchLimit(error: unknown): boolean { + return ( + error instanceof LakebaseApiError && + /branches limit exceeded/i.test(error.message) + ); +} + const API = "/api/2.0/postgres"; // The guard decides from these fields whether a command may write to a branch. @@ -70,6 +140,11 @@ function statusCodeOf(error: unknown): number | undefined { return typeof code === "number" ? code : undefined; } +function errorCodeOf(error: unknown): string | undefined { + const code = (error as { errorCode?: unknown } | undefined)?.errorCode; + return typeof code === "string" && code ? code : undefined; +} + function messageOf(error: unknown): string { return error instanceof Error ? error.message : String(error); } @@ -84,10 +159,27 @@ function encodeQuery( return entries.length > 0 ? Object.fromEntries(entries) : undefined; } +const sleep = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + +interface LakebaseApiOptions { + /** First poll delay for long-running operations; doubles up to 5 s. */ + readonly pollIntervalMs?: number; + /** Give up waiting for one long-running operation after this long. */ + readonly timeoutMs?: number; +} + export class LakebaseApi { + private readonly pollIntervalMs: number; + private readonly timeoutMs: number; + constructor( private readonly client: Pick, - ) {} + options: LakebaseApiOptions = {}, + ) { + this.pollIntervalMs = options.pollIntervalMs ?? 500; + this.timeoutMs = options.timeoutMs ?? 10 * 60_000; + } getBranch(name: string): Promise { return this.getParsedOrNull(`get branch ${name}`, name, branchResponse); @@ -112,6 +204,124 @@ export class LakebaseApi { throw new LakebaseApiError(operation, `unexpected response (${issues})`); } + listBranches(project: string): Promise { + return this.list( + `list branches of ${project}`, + `${project}/branches`, + "branches", + ); + } + + async createBranch( + project: string, + branchId: string, + spec: BranchSpec, + ): Promise { + const operation = `create branch ${project}/branches/${branchId}`; + const op = await this.call(operation, { + path: `${API}/${project}/branches`, + method: "POST", + query: { branch_id: branchId }, + payload: { spec }, + }); + return this.wait(operation, op); + } + + listEndpoints(branch: string): Promise { + return this.list( + `list endpoints of ${branch}`, + `${branch}/endpoints`, + "endpoints", + ); + } + + /** A read-write endpoint that suspends after `suspendTimeout` idle (e.g. "300s"). */ + async createEndpoint( + branch: string, + endpointId: string, + options: { suspendTimeout: string }, + ): Promise { + const operation = `create endpoint ${branch}/endpoints/${endpointId}`; + const op = await this.call(operation, { + path: `${API}/${branch}/endpoints`, + method: "POST", + query: { endpoint_id: endpointId }, + payload: { + spec: { + endpoint_type: READ_WRITE_ENDPOINT, + suspend_timeout_duration: options.suspendTimeout, + }, + }, + }); + return this.wait(operation, op); + } + + /** A branch that is already gone counts as deleted. */ + async deleteBranch(name: string): Promise { + await this.deleteResource(`delete branch ${name}`, { + path: `${API}/${name}`, + }); + } + + getDatabase(name: string): Promise { + return this.getOrNull(`get database ${name}`, name); + } + + listRoles(branch: string): Promise { + return this.list(`list roles of ${branch}`, `${branch}/roles`, "roles"); + } + + /** Without a role_id; the API derives one (sidesteps the USER role_id format). */ + async createRole(branch: string, spec: RoleSpec): Promise { + const operation = `create role ${spec.postgres_role} on ${branch}`; + const op = await this.call(operation, { + path: `${API}/${branch}/roles`, + method: "POST", + payload: { spec }, + }); + return this.wait(operation, op); + } + + async setRoleMemberships( + role: string, + memberships: RoleSpec["membership_roles"], + ): Promise { + const operation = `update role ${role}`; + const op = await this.call(operation, { + path: `${API}/${role}`, + method: "PATCH", + query: { update_mask: "spec.membership_roles" }, + payload: { spec: { membership_roles: memberships } }, + }); + return this.wait(operation, op); + } + + /** Every object the role owns is reassigned first (runs SQL on the branch's compute). */ + async deleteRole( + role: string, + options: { reassignOwnedTo: string }, + ): Promise { + await this.deleteResource(`delete role ${role}`, { + path: `${API}/${role}`, + query: { reassign_owned_to: options.reassignOwnedTo }, + }); + } + + private async deleteResource( + operation: string, + request: { path: string; query?: ToolingRequest["query"] }, + ): Promise { + try { + const op = await this.call(operation, { + ...request, + method: "DELETE", + }); + await this.wait(operation, op); + } catch (error) { + if (!isNotFound(error)) throw error; + } + } + private async call( operation: string, request: ToolingRequest, @@ -136,6 +346,7 @@ export class LakebaseApi { } catch (error) { throw new LakebaseApiError(operation, messageOf(error), { statusCode: statusCodeOf(error), + errorCode: errorCodeOf(error), cause: error, }); } @@ -151,10 +362,62 @@ export class LakebaseApi { method: "GET", }); } catch (error) { - if (error instanceof LakebaseApiError && error.statusCode === 404) { - return null; - } + if (isNotFound(error)) return null; throw error; } } + + private async list( + operation: string, + collection: string, + key: string, + ): Promise { + const items: T[] = []; + let pageToken: string | undefined; + do { + const page = await this.call>(operation, { + path: `${API}/${collection}`, + method: "GET", + ...(pageToken ? { query: { page_token: pageToken } } : {}), + }); + items.push(...((page?.[key] as T[] | undefined) ?? [])); + pageToken = (page?.next_page_token as string | undefined) || undefined; + } while (pageToken); + return items; + } + + /** Poll GET /api/2.0/postgres/{operation.name} until done. */ + private async wait(operation: string, first: Operation): Promise { + const deadline = Date.now() + this.timeoutMs; + let delay = this.pollIntervalMs; + let current = first; + while (!current?.done) { + if (!current?.name) { + throw new LakebaseApiError( + operation, + "the API returned an unfinished operation without a name", + ); + } + if (Date.now() >= deadline) { + throw new LakebaseApiError( + operation, + `timed out after ${Math.round(this.timeoutMs / 1000)}s waiting for ${current.name}`, + ); + } + await sleep(delay); + delay = Math.min(delay * 2, 5_000); + current = await this.call(operation, { + path: `${API}/${current.name}`, + method: "GET", + }); + } + if (current.error) { + throw new LakebaseApiError( + operation, + current.error.message ?? current.error.error_code ?? "unknown error", + { errorCode: current.error.error_code }, + ); + } + return current.response as T; + } } diff --git a/packages/appkit/src/database/tooling/migrate.ts b/packages/appkit/src/database/tooling/migrate.ts index 1613bb9d9..985e1b147 100644 --- a/packages/appkit/src/database/tooling/migrate.ts +++ b/packages/appkit/src/database/tooling/migrate.ts @@ -28,15 +28,23 @@ import { } from "./migrations"; /** Whether the checkout has a config/database/schema.ts yet. */ -function hasSchemaFile(root: string): boolean { +export function hasSchemaFile(root: string): boolean { return fs.existsSync(path.join(root, DEFAULT_SCHEMA_FILE)); } +const WITHOUT_SCHEMA = { + migrate: "there is nothing to migrate", + clean: "clean cannot tell which schema to drop", +} as const; + /** Without schema.ts, point to introspect instead of a load error. */ -function requireSchemaFile(root: string): void { +export function requireSchemaFile( + root: string, + command: keyof typeof WITHOUT_SCHEMA, +): void { if (hasSchemaFile(root)) return; throw new Error( - `No ${DEFAULT_SCHEMA_FILE} yet, so there is nothing to migrate. Run \`appkit db introspect\` to write it from the branch's tables, or create it.`, + `No ${DEFAULT_SCHEMA_FILE} yet, so ${WITHOUT_SCHEMA[command]}. Run \`appkit db introspect\` to write it from the branch's tables, or create it.`, ); } @@ -67,6 +75,8 @@ export interface MigrateReport { export interface MigrateOptions { /** false: an ambiguous rename throws instead of prompting. */ readonly interactive: boolean; + /** false: apply the committed migrations only, without writing one from schema.ts. */ + readonly generate?: boolean; } type ApplyOutcome = Pick; @@ -107,17 +117,20 @@ export async function runMigrate( ctx: DbContext, options: MigrateOptions, ): Promise { - requireSchemaFile(ctx.root); + requireSchemaFile(ctx.root, "migrate"); const schema = await loadDefaultDatabaseSchema(ctx.root, { showImportError: true, }); const folder = migrationsFolder(ctx.root); const journalBefore = readJournalText(folder); - const generated = await generateMigrationFiles({ - folder, - schema, - interactive: options.interactive, - }); + const generated = + options.generate === false + ? null + : await generateMigrationFiles({ + folder, + schema, + interactive: options.interactive, + }); // The types follow schema.ts alone, so they are current even if applying fails. const typesFile = path.join(ctx.root, DATABASE_TYPES_PATH); diff --git a/packages/appkit/src/database/tooling/migrations.ts b/packages/appkit/src/database/tooling/migrations.ts index 0983b5b74..2ed8488eb 100644 --- a/packages/appkit/src/database/tooling/migrations.ts +++ b/packages/appkit/src/database/tooling/migrations.ts @@ -218,6 +218,22 @@ function installedPackage(name: string, from: string): InstalledPackage | null { return null; } +/** + * Whether the app at `root` has drizzle-kit, at any version. The node_modules + * walk is what ESM resolution does; createRequire would also follow NODE_PATH, + * which `import("drizzle-kit/api")` ignores. + */ +export function hasDrizzleKit(root: string): boolean { + for (let dir = path.resolve(root); ; dir = path.dirname(dir)) { + if ( + fs.existsSync(path.join(dir, "node_modules/drizzle-kit/package.json")) + ) { + return true; + } + if (dir === path.dirname(dir)) return false; + } +} + /** * Why the drizzle-kit that resolves from `from` must not write migrations; * undefined when it may, or when it is not installed at all. drizzle-kit reads diff --git a/packages/appkit/src/database/tooling/naming.ts b/packages/appkit/src/database/tooling/naming.ts index aadaf374b..2310baf95 100644 --- a/packages/appkit/src/database/tooling/naming.ts +++ b/packages/appkit/src/database/tooling/naming.ts @@ -1,4 +1,4 @@ -import { createHash } from "node:crypto"; +import { createHash, randomInt } from "node:crypto"; // Branch names AppKit creates. Branch IDs: 1-63 chars, start with a // lowercase letter, lowercase letters, digits and hyphens only. @@ -8,7 +8,12 @@ const DEV_PREFIX = "dev-"; const SESSION_PREFIX = "session-"; const HASH_LENGTH = 6; -const SESSION_ID = /^[a-z0-9]{1,55}$/; +/** Seven days: a session branch outlives a working week, then expires on its own. */ +export const SESSION_BRANCH_TTL = "604800s"; + +const SESSION_ID_LENGTH = + MAX_BRANCH_ID - SESSION_PREFIX.length - HASH_LENGTH - 1; +const SESSION_ID = new RegExp(`^[a-z0-9]{1,${SESSION_ID_LENGTH}}$`); function slug(value: string): string { return value @@ -17,32 +22,75 @@ function slug(value: string): string { .replace(/^-+|-+$/g, ""); } +function normalized(userName: string): string { + return userName.trim().toLowerCase(); +} + +/** sha256(userName)[0:6]: the part of a branch name that says whose it is. */ +function userHash(userName: string): string { + return createHash("sha256") + .update(normalized(userName)) + .digest("hex") + .slice(0, HASH_LENGTH); +} + /** - * dev--. The hash keeps users who share - * a local part apart; the name is the same on every machine. + * dev--. The hash keeps users who share a local + * part apart; the name is the same on every machine. */ export function devBranchId(userName: string): string { - const normalized = userName.trim().toLowerCase(); - const hash = createHash("sha256") - .update(normalized) - .digest("hex") - .slice(0, HASH_LENGTH); const room = MAX_BRANCH_ID - DEV_PREFIX.length - 1 - HASH_LENGTH; - const local = slug(normalized.split("@")[0] ?? "") + const local = slug(normalized(userName).split("@")[0] ?? "") .slice(0, room) .replace(/-+$/, ""); - return `${DEV_PREFIX}${local || "user"}-${hash}`; + return `${DEV_PREFIX}${local || "user"}-${userHash(userName)}`; } export function isValidSessionId(sessionId: string): boolean { return SESSION_ID.test(sessionId); } -export function sessionBranchId(sessionId: string): string { +/** + * session--. The id lives in this checkout's .env; the + * hash makes a copied .env name a branch only its creator passes the guard on. + */ +export function sessionBranchId(userName: string, sessionId: string): string { if (!isValidSessionId(sessionId)) { throw new Error( - `Invalid APPKIT_DB_SESSION ${JSON.stringify(sessionId)}: use 1-55 lowercase letters or digits.`, + `Invalid APPKIT_DB_SESSION ${JSON.stringify(sessionId)}: use 1-${SESSION_ID_LENGTH} lowercase letters or digits.`, + ); + } + return `${SESSION_PREFIX}${userHash(userName)}-${sessionId}`; +} + +/** Whether a branch id is one of this user's dev or session branches. */ +export function isUsersBranchId(userName: string, branchId: string): boolean { + return ( + branchId === devBranchId(userName) || + branchId.startsWith(`${SESSION_PREFIX}${userHash(userName)}-`) + ); +} + +const BASE36 = "0123456789abcdefghijklmnopqrstuvwxyz"; + +/** 8 random base36 characters, generated once per checkout and kept in .env. */ +export function newSessionId(): string { + let id = ""; + for (let i = 0; i < 8; i++) id += BASE36[randomInt(BASE36.length)]; + return id; +} + +const TTL = /^(\d+)([smhd])$/; +const SECONDS = { s: 1, m: 60, h: 3_600, d: 86_400 } as const; + +/** `--ttl 30m`, `1h`, `7d` or `3600s` as the API's duration ("3600s"). */ +export function parseTtl(value: string): string { + const match = TTL.exec(value.trim()); + const amount = Number(match?.[1]); + if (!match || amount <= 0) { + throw new Error( + `Invalid --ttl ${JSON.stringify(value)}: use a positive number of seconds, minutes, hours or days, e.g. 30m, 1h or 7d.`, ); } - return `${SESSION_PREFIX}${sessionId}`; + return `${amount * SECONDS[match[2] as keyof typeof SECONDS]}s`; } diff --git a/packages/appkit/src/database/tooling/seed.ts b/packages/appkit/src/database/tooling/seed.ts new file mode 100644 index 000000000..6deaa2953 --- /dev/null +++ b/packages/appkit/src/database/tooling/seed.ts @@ -0,0 +1,47 @@ +import fs from "node:fs"; +import path from "node:path"; + +import { loadDefaultDatabaseSchema } from "../../plugins/database/load-schema"; +import { quoteIdentifier } from "../migrations/apply"; +import { DEFAULT_SCHEMA_NAME } from "../schema-builder/define-schema"; +import { runInTransaction } from "./branch-sql"; +import type { DbContext } from "./context"; +import { writableEnvBranch } from "./guard"; +import { hasSchemaFile } from "./migrate"; + +export const SEED_FILE = "config/database/seed.sql"; + +interface SeedReport { + readonly branch: string; + readonly file: string; +} + +/** + * `appkit db seed`: run config/database/seed.sql in one transaction on your + * dev or session branch, with the app schema first on the search_path so + * unqualified names resolve to it. Refuses any other branch before connecting. + */ +export async function runSeed(ctx: DbContext): Promise { + const file = path.join(ctx.root, SEED_FILE); + if (!fs.existsSync(file)) { + throw new Error( + `No seed file at ${SEED_FILE}. Create it with the SQL to run.`, + ); + } + const branch = await writableEnvBranch(ctx, "seed"); + const schemaName = hasSchemaFile(ctx.root) + ? (await loadDefaultDatabaseSchema(ctx.root, { showImportError: true })) + .$schemaName + : DEFAULT_SCHEMA_NAME; + const searchPath = + schemaName === "public" + ? "public" + : `${quoteIdentifier(schemaName)}, public`; + const sql = fs.readFileSync(file, "utf-8"); + await runInTransaction(ctx, async (client) => { + await client.query(`SET LOCAL search_path TO ${searchPath}`); + // One simple-protocol query: the file may hold several statements. + await client.query(sql); + }); + return { branch: branch.name, file }; +} diff --git a/packages/appkit/src/database/tooling/targets.ts b/packages/appkit/src/database/tooling/targets.ts index ff1cef3e8..b359e9364 100644 --- a/packages/appkit/src/database/tooling/targets.ts +++ b/packages/appkit/src/database/tooling/targets.ts @@ -5,6 +5,8 @@ import yaml from "js-yaml"; import { isValidSessionId } from "./naming"; +export type BranchMode = "dev" | "session"; + /** Full resource names, e.g. branch = "projects/p/branches/b". */ interface ResourceName { readonly project: string; @@ -39,9 +41,26 @@ export function resourceId(name: string): string { return name.slice(name.lastIndexOf("/") + 1); } +const PROJECT_ID = /^[a-z0-9][a-z0-9-]{0,62}$/; + +/** + * LAKEBASE_PROJECT as a project resource name: accepts "projects/" or a bare + * ""; null for anything else (e.g. the .env.example placeholder). + */ +export function parseProject(value: string | undefined): string | null { + const trimmed = value?.trim() ?? ""; + const id = trimmed.startsWith("projects/") + ? trimmed.slice("projects/".length) + : trimmed; + return PROJECT_ID.test(id) ? `projects/${id}` : null; +} + /** What `.env` says about the local branch. Unparseable values (template placeholders) read as null. */ -interface LocalTarget { +export interface LocalTarget { + /** LAKEBASE_PROJECT, normalized to "projects/". */ + readonly project: string | null; readonly endpoint: string | null; + readonly database: string | null; readonly session: string | null; } @@ -51,7 +70,19 @@ export function readLocalTarget(env: NodeJS.ProcessEnv): LocalTarget { env.APPKIT_DB_SESSION && isValidSessionId(env.APPKIT_DB_SESSION) ? env.APPKIT_DB_SESSION : null; - return { endpoint, session }; + return { + project: parseProject(env.LAKEBASE_PROJECT), + endpoint, + // Without a real endpoint PGDATABASE is a template placeholder too. + database: endpoint ? env.PGDATABASE || null : null, + session, + }; +} + +/** The branch the bundle's default target deploys to. */ +export interface DeployTarget { + readonly branch: string; + readonly database: string | null; } type Yaml = Record | undefined; @@ -92,6 +123,28 @@ function variableOf( return typeof value === "string" && value ? value : null; } +/** Default target's postgres_branch / postgres_database variables; null without databricks.yml or binding. */ +export function readDeployTarget(root: string): DeployTarget | null { + const bundle = readBundle(root); + if (!bundle) return null; + const targets = record(bundle.targets) ?? {}; + const target = + Object.values(targets) + .map(record) + .find((candidate) => candidate?.default === true) ?? + record(targets.default); + + const branch = parseResourceName( + variableOf(bundle, target, "postgres_branch") ?? undefined, + )?.branch; + if (!branch) return null; + const database = + parseResourceName( + variableOf(bundle, target, "postgres_database") ?? undefined, + )?.database ?? null; + return { branch, database }; +} + /** * Every branch databricks.yml can deploy to: each target's postgres_branch and * the variable's default. Empty without databricks.yml or a postgres binding. diff --git a/packages/appkit/src/database/tooling/tests/clean-seed.test.ts b/packages/appkit/src/database/tooling/tests/clean-seed.test.ts new file mode 100644 index 000000000..113e620db --- /dev/null +++ b/packages/appkit/src/database/tooling/tests/clean-seed.test.ts @@ -0,0 +1,262 @@ +import fs from "node:fs"; +import path from "node:path"; + +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +const { resolveEnvBranch, withBranchPool, writableEnvBranch } = vi.hoisted( + () => ({ + resolveEnvBranch: vi.fn(), + withBranchPool: vi.fn(), + writableEnvBranch: vi.fn(), + }), +); + +vi.mock("../guard", async (importOriginal) => ({ + ...(await importOriginal()), + resolveEnvBranch, + writableEnvBranch, +})); +vi.mock("../context", async (importOriginal) => ({ + ...(await importOriginal()), + withBranchPool, +})); + +import { deleteSessionBranch, previewClean, runClean } from "../clean"; +import { runSeed, SEED_FILE } from "../seed"; +import { cleanupTempApps, FakeLakebase, fakeContext, tempApp } from "./fakes"; + +const builder = path.resolve( + import.meta.dirname, + "../../schema-builder/index.ts", +); +const schemaFile = (schemaName: string) => + `import { defineSchema, id } from ${JSON.stringify(builder)}; + export const schema = defineSchema(({ table }) => ({ notes: table("notes", { id: id() }) }), { schemaName: ${JSON.stringify(schemaName)} });`; + +const DEPENDENTS = /pg_catalog\.pg_depend/; + +/** A pool whose one client records statements, can fail on one, and answers the dependents query. */ +function recordingPool( + options: { failOn?: RegExp; dependents?: string[] } = {}, +) { + const statements: string[] = []; + const client = { + released: undefined as { error: unknown } | undefined, + async query(text: string) { + statements.push(DEPENDENTS.test(text) ? "" : text); + if (options.failOn?.test(text)) throw new Error("boom"); + return { + rows: DEPENDENTS.test(text) + ? (options.dependents ?? []).map((label) => ({ label })) + : [], + }; + }, + release(error?: unknown) { + client.released = { error }; + }, + }; + const pool = { + connect: vi.fn(async () => client), + query: (text: string) => client.query(text), + }; + withBranchPool.mockImplementation((_ctx, fn: (p: unknown) => unknown) => + fn(pool), + ); + return { pool, client, statements }; +} + +const devBranch = { name: "projects/p/branches/dev-alice-3f9a2c" }; +const refusal = (action: string) => + new Error( + `Refusing to ${action}: production is not your dev or session branch (it is the project's default branch).`, + ); + +beforeEach(() => { + resolveEnvBranch.mockReset(); + withBranchPool.mockReset(); + writableEnvBranch.mockReset(); +}); +afterEach(cleanupTempApps); + +describe("runClean", () => { + test("drops the app schema in one transaction on your branch and reports what CASCADE removed outside it", async () => { + const ctx = fakeContext({ + root: tempApp({ "config/database/schema.ts": schemaFile("app") }), + }); + writableEnvBranch.mockResolvedValue(devBranch); + const { client, statements } = recordingPool({ + dependents: ["view reporting.recent"], + }); + + await expect(runClean(ctx)).resolves.toEqual({ + branch: devBranch.name, + schemaName: "app", + dependents: ["view reporting.recent"], + }); + expect(writableEnvBranch).toHaveBeenCalledWith(ctx, "clean"); + expect(statements).toEqual([ + "begin", + "", + `DROP SCHEMA IF EXISTS "app" CASCADE`, + "commit", + ]); + expect(client.released).toEqual({ error: undefined }); + }); + + test("without schema.ts it points to introspect before checking the branch", async () => { + const ctx = fakeContext({ root: tempApp() }); + + await expect(runClean(ctx)).rejects.toThrow( + "No config/database/schema.ts yet, so clean cannot tell which schema to drop. Run `appkit db introspect` to write it from the branch's tables, or create it.", + ); + expect(writableEnvBranch).not.toHaveBeenCalled(); + expect(withBranchPool).not.toHaveBeenCalled(); + }); + + test("refuses a branch that is not yours before connecting", async () => { + const ctx = fakeContext({ + root: tempApp({ "config/database/schema.ts": schemaFile("app") }), + }); + writableEnvBranch.mockRejectedValue(refusal("clean")); + + await expect(runClean(ctx)).rejects.toThrow(/^Refusing to clean: /); + expect(withBranchPool).not.toHaveBeenCalled(); + }); + + test("a failing statement rolls back and discards the connection", async () => { + const ctx = fakeContext({ + root: tempApp({ "config/database/schema.ts": schemaFile("app") }), + }); + writableEnvBranch.mockResolvedValue(devBranch); + const { client, statements } = recordingPool({ failOn: /DROP SCHEMA/ }); + + await expect(runClean(ctx)).rejects.toThrow("boom"); + expect(statements.slice(-1)).toEqual(["rollback"]); + expect(client.released?.error).toBeInstanceOf(Error); + }); +}); + +describe("previewClean", () => { + test("reads the dependents without opening a transaction or dropping anything", async () => { + const ctx = fakeContext({ + root: tempApp({ "config/database/schema.ts": schemaFile("app") }), + }); + writableEnvBranch.mockResolvedValue(devBranch); + const { statements } = recordingPool({ + dependents: ["constraint refs_fk on table reporting.refs"], + }); + + await expect(previewClean(ctx)).resolves.toEqual({ + branch: devBranch.name, + schemaName: "app", + dependents: ["constraint refs_fk on table reporting.refs"], + }); + expect(statements).toEqual([""]); + }); +}); + +describe("deleteSessionBranch", () => { + const session = { + branch: { name: "projects/p/branches/session-abc123-x1" }, + kind: "session", + }; + + test("deletes this checkout's session branch", async () => { + const lakebase = new FakeLakebase(); + lakebase.addBranch(session.branch.name); + resolveEnvBranch.mockResolvedValue(session); + + await expect( + deleteSessionBranch( + fakeContext({ root: tempApp(), api: lakebase.api() }), + ), + ).resolves.toEqual({ branch: session.branch.name }); + expect(lakebase.calls).toEqual([`deleteBranch ${session.branch.name}`]); + }); + + test("keeps your dev branch and says how to delete it", async () => { + resolveEnvBranch.mockResolvedValue({ branch: devBranch, kind: "dev" }); + + await expect( + deleteSessionBranch(fakeContext({ root: tempApp() })), + ).rejects.toThrow( + `Refusing to delete dev-alice-3f9a2c: --delete-branch deletes session branches only, and your dev branch stays for the next \`appkit db init\`. Delete it with \`databricks postgres delete-branch ${devBranch.name}\`.`, + ); + }); + + test("refuses any other branch, and a .env without one", async () => { + resolveEnvBranch.mockResolvedValueOnce({ + branch: { name: "projects/p/branches/production" }, + kind: "default", + reason: "it is the project's default branch", + }); + const ctx = fakeContext({ root: tempApp() }); + + await expect(deleteSessionBranch(ctx)).rejects.toThrow( + "Refusing to delete the branch: production is not your dev or session branch (it is the project's default branch).", + ); + resolveEnvBranch.mockResolvedValueOnce(null); + await expect(deleteSessionBranch(ctx)).rejects.toThrow( + /^Did not delete the branch: no branch is configured in \.env/, + ); + }); +}); + +describe("runSeed", () => { + test("runs config/database/seed.sql in one transaction on your branch, app schema first", async () => { + const sql = + "insert into app.notes default values;\ninsert into app.notes default values;\n"; + const root = tempApp({ [SEED_FILE]: sql }); + writableEnvBranch.mockResolvedValue(devBranch); + const { statements } = recordingPool(); + + await expect(runSeed(fakeContext({ root }))).resolves.toEqual({ + branch: devBranch.name, + file: path.join(root, SEED_FILE), + }); + expect(writableEnvBranch).toHaveBeenCalledWith(expect.anything(), "seed"); + expect(statements).toEqual([ + "begin", + 'SET LOCAL search_path TO "app", public', + sql, + "commit", + ]); + }); + + test.each([ + ["sales", 'SET LOCAL search_path TO "sales", public'], + ["public", "SET LOCAL search_path TO public"], + ])( + "with schema.ts on %s, unqualified names resolve there", + async (schemaName, searchPath) => { + const root = tempApp({ + [SEED_FILE]: "insert into notes default values;", + "config/database/schema.ts": schemaFile(schemaName), + }); + writableEnvBranch.mockResolvedValue(devBranch); + const { statements } = recordingPool(); + + await runSeed(fakeContext({ root })); + + expect(statements.slice(0, 2)).toEqual(["begin", searchPath]); + }, + ); + + test("a missing seed file is an error, before any API call", async () => { + await expect(runSeed(fakeContext({ root: tempApp() }))).rejects.toThrow( + "No seed file at config/database/seed.sql", + ); + expect(writableEnvBranch).not.toHaveBeenCalled(); + }); + + test("refuses a branch that is not yours before connecting", async () => { + const root = tempApp({ [SEED_FILE]: "select 1;" }); + writableEnvBranch.mockRejectedValue(refusal("seed")); + + await expect(runSeed(fakeContext({ root }))).rejects.toThrow( + /^Refusing to seed: /, + ); + expect(withBranchPool).not.toHaveBeenCalled(); + expect(fs.existsSync(path.join(root, SEED_FILE))).toBe(true); + }); +}); diff --git a/packages/appkit/src/database/tooling/tests/clean.pg.test.ts b/packages/appkit/src/database/tooling/tests/clean.pg.test.ts new file mode 100644 index 000000000..646502e92 --- /dev/null +++ b/packages/appkit/src/database/tooling/tests/clean.pg.test.ts @@ -0,0 +1,164 @@ +import { randomUUID } from "node:crypto"; + +import { Pool } from "pg"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; + +import { inTransaction } from "../../migrations/apply"; +import { cleanStatements, readCascadeDependents } from "../clean"; + +// The clean SQL against a real server, as a superuser that may create +// databases and roles, e.g. APPKIT_TEST_PG_URL=postgres://localhost:5432/postgres +const url = process.env.APPKIT_TEST_PG_URL; + +async function clean(pool: Pool, schemaName: string): Promise { + const client = await pool.connect(); + try { + return await inTransaction(client, async () => { + const dependents = await readCascadeDependents(client, schemaName); + for (const statement of cleanStatements(schemaName)) { + await client.query(statement); + } + return dependents; + }); + } finally { + client.release(); + } +} + +async function names(pool: Pool, sql: string): Promise { + const { rows } = await pool.query<{ name: string }>(sql); + return rows.map((row) => row.name); +} + +describe.skipIf(!url)("clean against Postgres", () => { + const suffix = randomUUID().slice(0, 8); + + describe("an app schema", () => { + const app = `appkit_clean_${suffix}`; + const other = `appkit_reporting_${suffix}`; + let pool: Pool; + + beforeAll(async () => { + pool = new Pool({ connectionString: url, max: 1 }); + await pool.query(` + create schema "${app}"; + create type "${app}".kind as enum ('a', 'b'); + create table "${app}".notes (id integer primary key, kind "${app}".kind); + create table "${app}".__drizzle_migrations (id serial primary key); + create schema "${other}"; + create view "${other}".recent as select id from "${app}".notes; + create table "${other}".refs (id integer primary key, note_id integer constraint refs_note_fk references "${app}".notes (id)); + create table "${other}".tagged (id integer primary key, kind "${app}".kind); + `); + }); + afterAll(async () => { + await pool.query(`drop schema if exists "${app}" cascade`); + await pool.query(`drop schema if exists "${other}" cascade`); + await pool.end(); + }); + + test("reports what CASCADE drops outside the schema, then drops the schema and only those", async () => { + const dependents = await clean(pool, app); + + expect(dependents).toEqual([ + `column kind of table ${other}.tagged`, + `constraint refs_note_fk on table ${other}.refs`, + `view ${other}.recent`, + ]); + expect( + await names( + pool, + `select nspname as name from pg_namespace where nspname = '${app}'`, + ), + ).toEqual([]); + expect( + await names( + pool, + `select relname as name from pg_class where relnamespace = '"${other}"'::regnamespace and relkind in ('r', 'v') order by 1`, + ), + ).toEqual(["refs", "tagged"]); + }); + }); + + describe("schema public", () => { + const database = `appkit_clean_${suffix}`; + const role = `appkit_other_${suffix}`; + let admin: Pool; + let pool: Pool; + + beforeAll(async () => { + admin = new Pool({ connectionString: url, max: 1 }); + await admin.query(`create database "${database}"`); + await admin.query(`create role "${role}"`); + const scratch = new URL(url as string); + scratch.pathname = `/${database}`; + pool = new Pool({ connectionString: scratch.toString(), max: 1 }); + await pool.query(` + create extension citext; + create extension tablefunc; + create type public.kind as enum ('a', 'b'); + create domain public.positive as integer check (value > 0); + create type public.pair as (a integer, b integer); + create table public.notes (id integer primary key, body citext, kind public.kind); + create table public.__drizzle_migrations (id serial primary key); + create function public.touch() returns integer language sql as 'select 1'; + create aggregate public.total(integer) (sfunc = int4pl, stype = integer); + create aggregate public.tally(*) (sfunc = int8inc, stype = bigint, initcond = '0'); + grant select on public.notes to "${role}"; + grant create on schema public to "${role}"; + set role "${role}"; + create view public.their_view as select id from public.notes; + reset role; + `); + }); + afterAll(async () => { + await pool.end(); + await admin.query(`drop database if exists "${database}"`); + await admin.query(`drop role if exists "${role}"`); + await admin.end(); + }); + + test("drops the role's own objects and keeps the extensions' and other roles' objects it does not depend on", async () => { + const extensionTypes = await names( + pool, + `select typname as name from pg_type t join pg_depend d on d.objid = t.oid and d.classid = 'pg_type'::regclass and d.deptype = 'e' order by 1`, + ); + expect(extensionTypes).toContain("citext"); + expect(extensionTypes).toContain("tablefunc_crosstab_2"); + + const dependents = await clean(pool, "public"); + + expect(dependents).toEqual(["view their_view"]); + expect( + await names( + pool, + `select relname as name from pg_class where relnamespace = 'public'::regnamespace and relkind in ('r', 'v', 'S')`, + ), + ).toEqual([]); + expect( + await names( + pool, + `select typname as name from pg_type where typnamespace = 'public'::regnamespace and typname in ('kind', 'positive', 'pair') `, + ), + ).toEqual([]); + expect( + await names( + pool, + `select proname as name from pg_proc where pronamespace = 'public'::regnamespace and proname in ('touch', 'total', 'tally')`, + ), + ).toEqual([]); + expect( + await names( + pool, + `select typname as name from pg_type where typnamespace = 'public'::regnamespace and typname in ('citext', 'tablefunc_crosstab_2') order by 1`, + ), + ).toEqual(["citext", "tablefunc_crosstab_2"]); + expect( + await names( + pool, + `select proname as name from pg_proc where pronamespace = 'public'::regnamespace and proname = 'crosstab' limit 1`, + ), + ).toEqual(["crosstab"]); + }); + }); +}); diff --git a/packages/appkit/src/database/tooling/tests/contract.test.ts b/packages/appkit/src/database/tooling/tests/contract.test.ts index 192305dcf..ed47a4f7a 100644 --- a/packages/appkit/src/database/tooling/tests/contract.test.ts +++ b/packages/appkit/src/database/tooling/tests/contract.test.ts @@ -10,8 +10,13 @@ const exports: DatabaseTooling = tooling; const contractKeys = { createDbContext: true, + deleteSessionBranch: true, + previewClean: true, + runClean: true, + runInit: true, runIntrospect: true, runMigrate: true, + runSeed: true, } satisfies Record; test("every function the CLI is typed with is exported", () => { diff --git a/packages/appkit/src/database/tooling/tests/drizzle-kit-install.test.ts b/packages/appkit/src/database/tooling/tests/drizzle-kit-install.test.ts new file mode 100644 index 000000000..6e3293f0c --- /dev/null +++ b/packages/appkit/src/database/tooling/tests/drizzle-kit-install.test.ts @@ -0,0 +1,74 @@ +import fs from "node:fs"; +import path from "node:path"; + +import { afterEach, describe, expect, test, vi } from "vitest"; + +import { drizzleKitInstall, ensureDrizzleKit } from "../drizzle-kit-install"; +import { DRIZZLE_KIT_VERSION } from "../migrations"; +import { cleanupTempApps, tempApp } from "./fakes"; + +afterEach(cleanupTempApps); + +const SPEC = `drizzle-kit@${DRIZZLE_KIT_VERSION}`; + +describe("drizzleKitInstall", () => { + test.each([ + ["pnpm-lock.yaml", "pnpm", ["add", "-D", "-E", SPEC]], + ["yarn.lock", "yarn", ["add", "-D", "-E", SPEC]], + ["bun.lock", "bun", ["add", "-d", "--exact", SPEC]], + [ + "package-lock.json", + "npm", + ["install", "--save-dev", "--save-exact", SPEC], + ], + ])("%s installs with %s, pinned", (lockfile, command, args) => { + expect(drizzleKitInstall(tempApp({ [lockfile]: "" }))).toEqual({ + command, + args, + }); + }); + + test("a workspace member uses the lockfile of the workspace root", () => { + const root = tempApp({ + "pnpm-lock.yaml": "", + "apps/web/package.json": "{}", + }); + expect(drizzleKitInstall(path.join(root, "apps/web")).command).toBe("pnpm"); + }); +}); + +describe("ensureDrizzleKit", () => { + test("installs into the app when it cannot resolve drizzle-kit, and reports the command", () => { + const root = tempApp({ "package.json": "{}", "yarn.lock": "" }); + const run = vi.fn(); + const steps: string[] = []; + + expect(ensureDrizzleKit(root, run, (step) => steps.push(step))).toBe( + `yarn add -D -E ${SPEC}`, + ); + expect(run).toHaveBeenCalledWith("yarn", ["add", "-D", "-E", SPEC], root); + expect(steps).toEqual([ + `Installing drizzle-kit ${DRIZZLE_KIT_VERSION}, which writes migrations (yarn add -D -E ${SPEC})`, + ]); + }); + + test("does nothing when drizzle-kit resolves, or there is no package.json", () => { + const run = vi.fn(); + const installed = tempApp({ + "package.json": "{}", + "node_modules/drizzle-kit/package.json": JSON.stringify({ + name: "drizzle-kit", + version: "0.30.0", + main: "index.js", + }), + "node_modules/drizzle-kit/index.js": "", + }); + + expect(ensureDrizzleKit(installed, run, () => undefined)).toBeNull(); + expect(ensureDrizzleKit(tempApp(), run, () => undefined)).toBeNull(); + expect(run).not.toHaveBeenCalled(); + expect(fs.existsSync(path.join(installed, "package-lock.json"))).toBe( + false, + ); + }); +}); diff --git a/packages/appkit/src/database/tooling/tests/env-file.test.ts b/packages/appkit/src/database/tooling/tests/env-file.test.ts new file mode 100644 index 000000000..35bfee92d --- /dev/null +++ b/packages/appkit/src/database/tooling/tests/env-file.test.ts @@ -0,0 +1,181 @@ +import fs from "node:fs"; +import path from "node:path"; + +import dotenv from "dotenv"; +import { afterEach, describe, expect, test, vi } from "vitest"; + +import { updateEnvFile } from "../env-file"; +import { cleanupTempApps, tempApp } from "./fakes"; + +afterEach(cleanupTempApps); + +describe("updateEnvFile", () => { + test("replaces values in place, keeps comments and order, appends new keys", () => { + const root = tempApp({ + ".env": [ + "# Databricks", + "DATABRICKS_CONFIG_PROFILE=DEFAULT", + "PGHOST='old-host'", + "", + "# keep me", + "LAKEBASE_ENDPOINT=projects/p/branches/old/endpoints/primary", + "APPKIT_DB_SESSION=abc", + "PGHOST=duplicate", + ].join("\n"), + }); + const file = path.join(root, ".env"); + + updateEnvFile(file, { + PGHOST: "new-host.example.com", + LAKEBASE_ENDPOINT: "projects/p/branches/dev-a/endpoints/primary", + OTHER_KEY: "value", + }); + + expect(fs.readFileSync(file, "utf-8")).toBe( + [ + "# Databricks", + "DATABRICKS_CONFIG_PROFILE=DEFAULT", + "PGHOST=new-host.example.com", + "", + "# keep me", + "LAKEBASE_ENDPOINT=projects/p/branches/dev-a/endpoints/primary", + "APPKIT_DB_SESSION=abc", + "OTHER_KEY=value", + "", + ].join("\n"), + ); + expect(fs.statSync(file).mode & 0o777).toBe(0o600); + }); + + test("a missing .env gets only the given keys, never .env.example's placeholders", () => { + const root = tempApp({ + ".env.example": "DATABRICKS_HOST=https://your-workspace\nPGPORT=5432\n", + }); + const file = path.join(root, ".env"); + + updateEnvFile(file, { PGHOST: "h.example.com" }); + + expect(fs.readFileSync(file, "utf-8")).toBe("PGHOST=h.example.com\n"); + expect(fs.statSync(file).mode & 0o777).toBe(0o600); + }); + + test("a symbolic link is refused and its target left alone", () => { + const root = tempApp({ "shared.env": "PGHOST=old\n" }); + const file = path.join(root, ".env"); + fs.symlinkSync(path.join(root, "shared.env"), file); + + expect(() => updateEnvFile(file, { PGHOST: "new" })).toThrow( + `Refusing to rewrite ${file}: it is a symbolic link to ${path.join(root, "shared.env")}, and writing through it would change that file. Replace it with a regular file, then rerun.`, + ); + expect(fs.readFileSync(path.join(root, "shared.env"), "utf-8")).toBe( + "PGHOST=old\n", + ); + expect(fs.lstatSync(file).isSymbolicLink()).toBe(true); + }); + + test("a multi-line quoted value is one entry: replaced whole, or kept whole", () => { + const root = tempApp({ + ".env": [ + 'PGHOST="first', + "PGHOST=not-a-key", + 'still the value"', + 'CERT="-----BEGIN-----', + "PGPORT=1", + '-----END-----"', + "PGPORT=5432", + ].join("\n"), + }); + const file = path.join(root, ".env"); + + updateEnvFile(file, { PGHOST: "new.example.com", PGPORT: "6543" }); + + const text = fs.readFileSync(file, "utf-8"); + expect(text).toBe( + [ + "PGHOST=new.example.com", + 'CERT="-----BEGIN-----', + "PGPORT=1", + '-----END-----"', + "PGPORT=6543", + "", + ].join("\n"), + ); + expect(dotenv.parse(text)).toEqual({ + PGHOST: "new.example.com", + CERT: "-----BEGIN-----\nPGPORT=1\n-----END-----", + PGPORT: "6543", + }); + }); + + test("a value whose quote never closes is a single line, as dotenv reads it", () => { + const root = tempApp({ ".env": 'NOTE="open\nPGHOST=old\n' }); + const file = path.join(root, ".env"); + + updateEnvFile(file, { PGHOST: "new" }); + + expect(fs.readFileSync(file, "utf-8")).toBe('NOTE="open\nPGHOST=new\n'); + }); + + test("a failed write leaves the old file and no temporary file", () => { + const root = tempApp({ ".env": "PGHOST=old\n" }); + const file = path.join(root, ".env"); + const rename = vi.spyOn(fs, "renameSync").mockImplementation(() => { + throw new Error("disk full"); + }); + try { + expect(() => updateEnvFile(file, { PGHOST: "new" })).toThrow("disk full"); + } finally { + rename.mockRestore(); + } + + expect(fs.readFileSync(file, "utf-8")).toBe("PGHOST=old\n"); + expect(fs.readdirSync(root)).toEqual([".env"]); + }); + + test("quotes values that need it", () => { + const root = tempApp(); + const file = path.join(root, ".env"); + updateEnvFile(file, { A: "has space", B: "plain@example.com" }); + expect(fs.readFileSync(file, "utf-8")).toBe( + 'A="has space"\nB=plain@example.com\n', + ); + }); + + test("keeps export prefixes and CRLF line endings", () => { + const root = tempApp({ + ".env": "# keep\r\nexport PGHOST=old\r\nPGPORT=5432\r\n", + }); + const file = path.join(root, ".env"); + + updateEnvFile(file, { PGHOST: "new.example.com", PGUSER: "a@example.com" }); + + expect(fs.readFileSync(file, "utf-8")).toBe( + "# keep\r\nexport PGHOST=new.example.com\r\nPGPORT=5432\r\nPGUSER=a@example.com\r\n", + ); + }); + + test("every value stays on one line and reads back unchanged", () => { + const root = tempApp({ ".env": "KEEP=1\n" }); + const file = path.join(root, ".env"); + const values = { + NEWLINE: "line1\nline2\r\nINJECTED=1", + QUOTES: 'say "hi"', + BACKSLASH: "back\\slash", + }; + + updateEnvFile(file, values); + + const text = fs.readFileSync(file, "utf-8"); + expect(text.split("\n").filter(Boolean)).toHaveLength(4); + expect(dotenv.parse(text)).toEqual({ KEEP: "1", ...values }); + }); + + test("a value dotenv cannot read back unchanged is refused", () => { + const root = tempApp(); + const file = path.join(root, ".env"); + expect(() => updateEnvFile(file, { BAD: `it's "x"\ny` })).toThrow( + "Cannot write BAD to .env: its value mixes quotes, backslashes and line breaks.", + ); + expect(fs.existsSync(file)).toBe(false); + }); +}); diff --git a/packages/appkit/src/database/tooling/tests/fakes.ts b/packages/appkit/src/database/tooling/tests/fakes.ts index ceaf6006a..d7e829fa4 100644 --- a/packages/appkit/src/database/tooling/tests/fakes.ts +++ b/packages/appkit/src/database/tooling/tests/fakes.ts @@ -4,16 +4,27 @@ import path from "node:path"; import type { DbContext } from "../context"; import type { + BranchSpec, LakebaseApi, LakebaseBranch, + LakebaseDatabase, LakebaseEndpoint, + LakebaseRole, + RoleSpec, } from "../lakebase-api"; +import { LakebaseApiError } from "../lakebase-api"; -/** In-memory Lakebase: the branches the guard looks up, and their endpoints. */ +/** In-memory Lakebase with the semantics init relies on: forks copy roles, endpoints come with new forks. */ export class FakeLakebase { readonly branches = new Map(); readonly endpoints = new Map(); + readonly roles = new Map(); + readonly databases = new Map(); readonly calls: string[] = []; + /** New forks get a read-write `primary` endpoint, like Lakebase forks do. */ + autoEndpoint = true; + createBranchError?: Error; + deleteRoleError?: Error; addBranch( name: string, @@ -22,13 +33,29 @@ export class FakeLakebase { const branch = { name, status: { current_state: "READY", ...status } }; this.branches.set(name, branch); this.endpoints.set(name, [this.endpoint(name, "primary")]); + this.roles.set(name, []); return branch; } + addRole( + branch: string, + roleId: string, + postgresRole: string, + memberships: string[] = [], + ) { + const role: LakebaseRole = { + name: `${branch}/roles/${roleId}`, + status: { postgres_role: postgresRole, membership_roles: memberships }, + }; + this.roles.get(branch)?.push(role); + return role; + } + endpoint(branch: string, id: string): LakebaseEndpoint { return { name: `${branch}/endpoints/${id}`, status: { + endpoint_type: "ENDPOINT_TYPE_READ_WRITE", hosts: { host: `${id}.${branch.split("/").pop()}.example.com` }, }, }; @@ -50,11 +77,120 @@ export class FakeLakebase { ?.find((endpoint) => endpoint.name === name) ?? null ); }, + async listBranches(project: string) { + fake.calls.push(`listBranches ${project}`); + return [...fake.branches.values()].filter((b) => + b.name.startsWith(`${project}/`), + ); + }, + async createBranch(project: string, branchId: string, spec: BranchSpec) { + fake.calls.push(`createBranch ${branchId} ${JSON.stringify(spec)}`); + if (fake.createBranchError) throw fake.createBranchError; + const name = `${project}/branches/${branchId}`; + if (fake.branches.has(name)) { + throw new LakebaseApiError( + `create branch ${name}`, + "ALREADY_EXISTS", + { + statusCode: 409, + errorCode: "ALREADY_EXISTS", + }, + ); + } + const ttlSeconds = "ttl" in spec ? Number.parseInt(spec.ttl, 10) : 0; + const branch: LakebaseBranch = { + name, + status: { + current_state: "READY", + source_branch: spec.source_branch, + ...(ttlSeconds + ? { + expire_time: new Date( + Date.now() + ttlSeconds * 1000, + ).toISOString(), + } + : {}), + }, + }; + fake.branches.set(name, branch); + fake.endpoints.set( + name, + fake.autoEndpoint ? [fake.endpoint(name, "primary")] : [], + ); + // A fork copies the parent's roles. + fake.roles.set( + name, + (fake.roles.get(spec.source_branch) ?? []).map((role) => ({ + ...role, + name: `${name}/roles/${role.name.split("/").pop()}`, + })), + ); + return branch; + }, + async listEndpoints(branch: string) { + fake.calls.push(`listEndpoints ${branch}`); + return fake.endpoints.get(branch) ?? []; + }, + async createEndpoint( + branch: string, + endpointId: string, + options: { suspendTimeout: string }, + ) { + fake.calls.push( + `createEndpoint ${branch} ${endpointId} suspend=${options.suspendTimeout}`, + ); + const endpoint = fake.endpoint(branch, endpointId); + fake.endpoints.set(branch, [ + ...(fake.endpoints.get(branch) ?? []), + endpoint, + ]); + return endpoint; + }, + async getDatabase(name: string) { + fake.calls.push(`getDatabase ${name}`); + return fake.databases.get(name) ?? null; + }, + async listRoles(branch: string) { + fake.calls.push(`listRoles ${branch}`); + return [...(fake.roles.get(branch) ?? [])]; + }, + async createRole(branch: string, spec: RoleSpec) { + fake.calls.push(`createRole ${branch} ${JSON.stringify(spec)}`); + return fake.addRole(branch, "rol-generated", spec.postgres_role, [ + ...spec.membership_roles, + ]); + }, + async setRoleMemberships(role: string, memberships: string[]) { + fake.calls.push(`setRoleMemberships ${role} ${memberships.join(",")}`); + const branch = role.slice(0, role.indexOf("/roles/")); + const found = fake.roles.get(branch)?.find((r) => r.name === role); + if (!found?.status) throw new Error(`no role ${role}`); + found.status.membership_roles = [...memberships]; + return found; + }, + async deleteBranch(name: string) { + fake.calls.push(`deleteBranch ${name}`); + fake.branches.delete(name); + }, + async deleteRole(role: string, options: { reassignOwnedTo: string }) { + fake.calls.push(`deleteRole ${role} → ${options.reassignOwnedTo}`); + if (fake.deleteRoleError) throw fake.deleteRoleError; + const branch = role.slice(0, role.indexOf("/roles/")); + fake.roles.set( + branch, + (fake.roles.get(branch) ?? []).filter((r) => r.name !== role), + ); + }, }; return api as unknown as LakebaseApi; } } +export const forbidden = () => + new LakebaseApiError("create branch", "PERMISSION_DENIED", { + statusCode: 403, + }); + const roots: string[] = []; /** A temp app root with optional files (relative path → contents). */ @@ -78,6 +214,8 @@ export function fakeContext(options: { env?: NodeJS.ProcessEnv; userName?: string; api?: LakebaseApi; + /** The host the client resolved; empty when a test does not care. */ + host?: string; }): DbContext { const userName = options.userName ?? "alice@example.com"; return { @@ -88,6 +226,7 @@ export function fakeContext(options: { throw new Error("no raw API calls expected"); }, currentUser: { me: async () => ({ userName }) }, + getHost: async () => options.host ?? "", }), api: () => { if (!options.api) throw new Error("no Lakebase API calls expected"); diff --git a/packages/appkit/src/database/tooling/tests/guard.test.ts b/packages/appkit/src/database/tooling/tests/guard.test.ts index ee79cbf59..146ad1333 100644 --- a/packages/appkit/src/database/tooling/tests/guard.test.ts +++ b/packages/appkit/src/database/tooling/tests/guard.test.ts @@ -1,7 +1,12 @@ import { afterEach, describe, expect, test } from "vitest"; -import { classifyBranch, resolveEnvBranch, takesWrites } from "../guard"; -import { devBranchId } from "../naming"; +import { + classifyBranch, + resolveEnvBranch, + takesWrites, + writableEnvBranch, +} from "../guard"; +import { devBranchId, sessionBranchId } from "../naming"; import { cleanupTempApps, databricksYml, @@ -15,6 +20,7 @@ afterEach(cleanupTempApps); const USER = "alice@example.com"; const P = "projects/p"; const dev = `${P}/branches/${devBranchId(USER)}`; +const session = `${P}/branches/${sessionBranchId(USER, "k2x9q0ab")}`; function verdict( name: string, @@ -33,9 +39,7 @@ function verdict( describe("classifyBranch", () => { test("only this developer's dev branch and this checkout's session branch take writes", () => { expect(verdict(dev)).toEqual({ kind: "dev" }); - expect(verdict(`${P}/branches/session-k2x9q0ab`)).toEqual({ - kind: "session", - }); + expect(verdict(session)).toEqual({ kind: "session" }); expect(takesWrites("dev")).toBe(true); expect(takesWrites("session")).toBe(true); }); @@ -57,11 +61,14 @@ describe("classifyBranch", () => { verdict(`${P}/branches/${devBranchId("bob@example.com")}`).kind, ).toBe("other"); expect(verdict(`${P}/branches/session-other`).kind).toBe("other"); - expect(verdict(`${P}/branches/session-k2x9q0ab`, {}, null).kind).toBe( - "other", - ); - expect(verdict(`${P}/branches/staging`).reason).toContain( - devBranchId(USER), + expect(verdict(session, {}, null).kind).toBe("other"); + // Bob's .env copied into Alice's checkout names Bob's session branch. + expect( + verdict(`${P}/branches/${sessionBranchId("bob@example.com", "k2x9q0ab")}`) + .kind, + ).toBe("other"); + expect(verdict(`${P}/branches/staging`).reason).toBe( + `it is not your dev branch (${devBranchId(USER)}) or this checkout's session branch (${sessionBranchId(USER, "k2x9q0ab")})`, ); for (const kind of ["default", "protected", "deploy-target", "other"]) { expect(takesWrites(kind as "other")).toBe(false); @@ -170,3 +177,58 @@ describe("resolveEnvBranch", () => { ).rejects.toThrow(/does not exist.*appkit db init/); }); }); + +describe("writableEnvBranch", () => { + const envFor = (branch: string, sessionId?: string) => ({ + PGHOST: "primary.example.com", + PGDATABASE: "databricks_postgres", + LAKEBASE_ENDPOINT: `${branch}/endpoints/primary`, + ...(sessionId ? { APPKIT_DB_SESSION: sessionId } : {}), + }); + + test("returns your dev or session branch", async () => { + const lakebase = new FakeLakebase(); + lakebase.addBranch(dev); + lakebase.addBranch(session); + + await expect( + writableEnvBranch( + fakeContext({ root: tempApp(), env: envFor(dev), api: lakebase.api() }), + "seed", + ), + ).resolves.toMatchObject({ name: dev }); + await expect( + writableEnvBranch( + fakeContext({ + root: tempApp(), + env: envFor(session, "k2x9q0ab"), + api: lakebase.api(), + }), + "seed", + ), + ).resolves.toMatchObject({ name: session }); + }); + + test("refuses another branch, and says what to run without one", async () => { + const lakebase = new FakeLakebase(); + lakebase.addBranch(`${P}/branches/production`, { default: true }); + + await expect( + writableEnvBranch( + fakeContext({ + root: tempApp(), + env: envFor(`${P}/branches/production`), + api: lakebase.api(), + }), + "clean", + ), + ).rejects.toThrow( + "Refusing to clean: production is not your dev or session branch (it is the project's default branch).", + ); + await expect( + writableEnvBranch(fakeContext({ root: tempApp() }), "seed"), + ).rejects.toThrow( + "Did not seed: no branch is configured in .env (PGHOST, PGDATABASE and LAKEBASE_ENDPOINT); run `appkit db init`.", + ); + }); +}); diff --git a/packages/appkit/src/database/tooling/tests/hand-off.pg.test.ts b/packages/appkit/src/database/tooling/tests/hand-off.pg.test.ts new file mode 100644 index 000000000..5466c5465 --- /dev/null +++ b/packages/appkit/src/database/tooling/tests/hand-off.pg.test.ts @@ -0,0 +1,75 @@ +import { randomUUID } from "node:crypto"; + +import { Pool } from "pg"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; + +import { readRoleBlockers } from "../hand-off"; + +// The pg_shdepend query against a real server, as a superuser that may create +// databases and roles, e.g. APPKIT_TEST_PG_URL=postgres://localhost:5432/postgres +const url = process.env.APPKIT_TEST_PG_URL; + +function databaseUrl(database: string): string { + const next = new URL(url as string); + next.pathname = `/${database}`; + return next.toString(); +} + +describe.skipIf(!url)("readRoleBlockers against Postgres", () => { + const suffix = randomUUID().slice(0, 8); + const role = `appkit_sp_${suffix}`; + const main = `appkit_main_${suffix}`; + const other = `appkit_other_${suffix}`; + const owned = `appkit_owned_${suffix}`; + let admin: Pool; + let mainPool: Pool; + let otherPool: Pool; + + beforeAll(async () => { + admin = new Pool({ connectionString: url, max: 1 }); + await admin.query(`create role "${role}"`); + await admin.query(`create database "${main}"`); + await admin.query(`create database "${other}"`); + await admin.query(`create database "${owned}" owner "${role}"`); + mainPool = new Pool({ connectionString: databaseUrl(main), max: 1 }); + otherPool = new Pool({ connectionString: databaseUrl(other), max: 1 }); + await mainPool.query(` + create schema app authorization "${role}"; + create table app.notes (id integer primary key); + alter table app.notes owner to "${role}"; + create table public.shared (id integer); + grant select on public.shared to "${role}"; + `); + await otherPool.query(` + create table public.reports (id integer); + alter table public.reports owner to "${role}"; + `); + }); + afterAll(async () => { + await mainPool.end(); + await otherPool.end(); + for (const database of [main, other, owned]) { + await admin.query(`drop database if exists "${database}"`); + } + await admin.query(`drop role if exists "${role}"`); + await admin.end(); + }); + + test("lists what outlives a reassign in this database: other databases' objects, shared objects, and grants", async () => { + const blockers = await readRoleBlockers(mainPool, role); + + expect(blockers).toEqual( + [ + { database: "", count: 1 }, + { database: main, count: 1 }, + { database: other, count: 1 }, + ].sort((a, b) => a.database.localeCompare(b.database)), + ); + }); + + test("a role without dependencies has no blockers", async () => { + await expect( + readRoleBlockers(mainPool, `appkit_missing_${suffix}`), + ).resolves.toEqual([]); + }); +}); diff --git a/packages/appkit/src/database/tooling/tests/identity.test.ts b/packages/appkit/src/database/tooling/tests/identity.test.ts index cc49fcc56..ef76b6f3e 100644 --- a/packages/appkit/src/database/tooling/tests/identity.test.ts +++ b/packages/appkit/src/database/tooling/tests/identity.test.ts @@ -1,11 +1,18 @@ import { afterEach, describe, expect, test, vi } from "vitest"; -import { currentIdentity } from "../identity"; +import { currentIdentity, identityTypeOf } from "../identity"; import { cleanupTempApps, fakeContext, tempApp } from "./fakes"; afterEach(cleanupTempApps); describe("identity", () => { + test("an email is a USER, a client ID is a SERVICE_PRINCIPAL", () => { + expect(identityTypeOf("alice@example.com")).toBe("USER"); + expect(identityTypeOf("11111111-2222-3333-4444-555555555555")).toBe( + "SERVICE_PRINCIPAL", + ); + }); + test("currentUser.me() is called once per context", async () => { const ctx = fakeContext({ root: tempApp() }); const me = vi.fn(async (_request: object) => ({ @@ -16,6 +23,7 @@ describe("identity", () => { await expect(currentIdentity(counted)).resolves.toEqual({ userName: "alice@example.com", + identityType: "USER", }); await currentIdentity(counted); expect(me).toHaveBeenCalledOnce(); diff --git a/packages/appkit/src/database/tooling/tests/init.test.ts b/packages/appkit/src/database/tooling/tests/init.test.ts new file mode 100644 index 000000000..0f0bb38ba --- /dev/null +++ b/packages/appkit/src/database/tooling/tests/init.test.ts @@ -0,0 +1,838 @@ +import fs from "node:fs"; +import path from "node:path"; + +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +const { runClean, runIntrospect, runMigrate, withBranchPool } = vi.hoisted( + () => ({ + runClean: vi.fn(), + runIntrospect: vi.fn(), + runMigrate: vi.fn(), + withBranchPool: vi.fn(), + }), +); + +vi.mock("../context", async (importOriginal) => ({ + ...(await importOriginal()), + withBranchPool, +})); +vi.mock("../migrate", async (importOriginal) => ({ + ...(await importOriginal()), + runMigrate, +})); +vi.mock("../clean", () => ({ runClean })); +vi.mock("../introspect", () => ({ runIntrospect })); + +import type { DbContext } from "../context"; +import { type InitOptions, runInit } from "../init"; +import { LakebaseApiError } from "../lakebase-api"; +import { DRIZZLE_KIT_VERSION } from "../migrations"; +import { devBranchId, sessionBranchId } from "../naming"; +import { + cleanupTempApps, + databricksYml, + FakeLakebase, + fakeContext, + forbidden, + tempApp, +} from "./fakes"; + +const USER = "alice@example.com"; +const SP = "11111111-2222-3333-4444-555555555555"; +const P = "projects/p"; +const PRODUCTION = `${P}/branches/production`; +const DEV = `${P}/branches/${devBranchId(USER)}`; +const session = (id: string) => `${P}/branches/${sessionBranchId(USER, id)}`; +const migrateReport = { + generated: null, + applied: [], + typesFile: "/app/shared/appkit-types/database.d.ts", +}; + +/** What the hand-off reads on the branch: the app schema's owner, and what blocks deleting that role. */ +let schemaOwner: string | null; +let blockers: { database: string; count: number }[]; +const queries: string[] = []; +const pool = { + query: vi.fn(async (text: string) => { + queries.push(text); + if (text.includes("nspowner")) { + return { rows: [{ owner: schemaOwner, current_user: USER }] }; + } + if (text.includes("pg_shdepend")) return { rows: blockers }; + throw new Error(`unexpected query ${text}`); + }), +}; + +function project(): FakeLakebase { + const lakebase = new FakeLakebase(); + lakebase.addBranch(PRODUCTION, { default: true }); + lakebase.addRole(PRODUCTION, "alice", USER, ["DATABRICKS_SUPERUSER"]); + lakebase.addRole(PRODUCTION, `dbrx-apps-${SP}`, SP); + lakebase.addRole(PRODUCTION, "dbrx-apps-user-bob", "bob@example.com"); + return lakebase; +} + +/** An app that deploys to `deploy` (the default branch unless set), with .env pointing at the default branch. */ +function appOn( + lakebase: FakeLakebase, + env: Record = {}, + { + schema = true, + deploy = PRODUCTION, + files = {}, + }: { schema?: boolean; deploy?: string; files?: Record } = {}, +): DbContext { + const dotenv = { + DATABRICKS_CONFIG_PROFILE: "DEFAULT", + PGHOST: "primary.production.example.com", + PGDATABASE: "databricks_postgres", + LAKEBASE_ENDPOINT: `${PRODUCTION}/endpoints/primary`, + ...env, + }; + const root = tempApp({ + ".env": `# Lakebase\n${Object.entries(dotenv) + .map(([key, value]) => `${key}=${value}`) + .join("\n")}\n`, + "databricks.yml": databricksYml( + deploy, + `${deploy}/databases/databricks-postgres`, + ), + ...(schema ? { "config/database/schema.ts": schemaSource("app") } : {}), + ...files, + }); + return fakeContext({ + root, + env: dotenv, + userName: USER, + api: lakebase.api(), + }); +} + +const builder = path.resolve( + import.meta.dirname, + "../../schema-builder/index.ts", +); +function schemaSource(schemaName: string): string { + return `import { defineSchema, id } from ${JSON.stringify(builder)}; +export const schema = defineSchema(({ table }) => ({ notes: table("notes", { id: id() }) }), { schemaName: ${JSON.stringify(schemaName)} }); +`; +} + +function envFile(root: string): Record { + return Object.fromEntries( + fs + .readFileSync(path.join(root, ".env"), "utf-8") + .split("\n") + .filter((line) => line.includes("=")) + .map((line) => [ + line.slice(0, line.indexOf("=")), + line.slice(line.indexOf("=") + 1), + ]), + ); +} + +/** Run init, collecting its progress lines. */ +async function init(ctx: DbContext, options: InitOptions = {}) { + const steps: string[] = []; + const report = await runInit(ctx, { + ...options, + onStep: (step) => steps.push(step), + }); + return { report, steps }; +} + +function conflict(name: string): LakebaseApiError { + return new LakebaseApiError(`create branch ${name}`, "ALREADY_EXISTS", { + statusCode: 409, + errorCode: "ALREADY_EXISTS", + }); +} + +beforeEach(() => { + schemaOwner = SP; + blockers = []; + queries.length = 0; + pool.query.mockClear(); + withBranchPool + .mockReset() + .mockImplementation((_ctx, fn: (p: unknown) => unknown) => fn(pool)); + runMigrate.mockReset().mockResolvedValue(migrateReport); + runClean.mockReset().mockResolvedValue({ + branch: DEV, + schemaName: "app", + dependents: [], + }); + runIntrospect.mockReset().mockResolvedValue({ branch: DEV, tables: [] }); +}); +afterEach(cleanupTempApps); + +describe("runInit — dev branch", () => { + test("forks the default branch with no expiry, hands the app schema to you, rewrites .env, migrates", async () => { + const lakebase = project(); + const ctx = appOn(lakebase); + + const { report, steps } = await init(ctx); + + expect(lakebase.calls).toContain( + `createBranch ${devBranchId(USER)} ${JSON.stringify({ source_branch: PRODUCTION, no_expiry: true })}`, + ); + // Only the owner of the app schema is deleted; Bob's copied role stays. + expect( + lakebase.calls.filter((call) => call.startsWith("deleteRole")), + ).toEqual([`deleteRole ${DEV}/roles/dbrx-apps-${SP} → ${DEV}/roles/alice`]); + expect(report).toMatchObject({ + mode: "dev", + branch: DEV, + created: true, + sourceBranch: PRODUCTION, + expireTime: null, + endpoint: `${DEV}/endpoints/primary`, + endpointCreated: false, + host: `primary.${devBranchId(USER)}.example.com`, + database: "databricks_postgres", + user: USER, + role: { name: `${DEV}/roles/alice`, action: "kept" }, + handedOff: SP, + installed: null, + start: { kind: "migrate", migrate: migrateReport }, + }); + expect(steps).toEqual([ + `Creating ${devBranchId(USER)} from production`, + `Pointing .env at ${devBranchId(USER)}, database databricks_postgres`, + `Handing schema app in database databricks_postgres on ${devBranchId(USER)} to you: deleting this branch's copy of role ${SP}, which owns it, with its objects reassigned to you`, + `Migrating schema app on ${devBranchId(USER)}`, + ]); + + const written = envFile(ctx.root); + expect(written).toEqual({ + DATABRICKS_CONFIG_PROFILE: "DEFAULT", + LAKEBASE_PROJECT: P, + PGHOST: `primary.${devBranchId(USER)}.example.com`, + LAKEBASE_ENDPOINT: `${DEV}/endpoints/primary`, + PGDATABASE: "databricks_postgres", + PGUSER: USER, + PGPORT: "5432", + PGSSLMODE: "require", + }); + expect(fs.readFileSync(path.join(ctx.root, ".env"), "utf-8")).toMatch( + /^# Lakebase\n/, + ); + // Migrate runs against the new branch; the caller's environment is left alone. + const [migrateCtx, migrateOptions] = runMigrate.mock.calls[0] ?? []; + expect(migrateCtx.env.LAKEBASE_ENDPOINT).toBe(`${DEV}/endpoints/primary`); + expect(migrateOptions).toEqual({ interactive: false }); + expect(ctx.env.LAKEBASE_ENDPOINT).toBe(`${PRODUCTION}/endpoints/primary`); + }); + + test("a re-run reuses the dev branch; when you own the app schema there is no hand-off", async () => { + const lakebase = project(); + lakebase.addBranch(DEV); + lakebase.addRole(DEV, "alice", USER, ["DATABRICKS_SUPERUSER"]); + lakebase.addRole(DEV, `dbrx-apps-${SP}`, SP); + schemaOwner = USER; + + const { report, steps } = await init(appOn(lakebase), { schema: "none" }); + + expect(report).toMatchObject({ + created: false, + handedOff: null, + start: { kind: "none" }, + }); + expect(steps[0]).toBe(`Using ${devBranchId(USER)}`); + expect( + lakebase.calls.some((call) => + /^(createBranch|deleteRole|createRole|createEndpoint)/.test(call), + ), + ).toBe(false); + expect(queries.some((text) => text.includes("pg_shdepend"))).toBe(false); + }); + + test("the hand-off reads the schema schema.ts names", async () => { + const lakebase = project(); + const ctx = appOn(lakebase, {}, { schema: false }); + fs.mkdirSync(path.join(ctx.root, "config/database"), { recursive: true }); + fs.writeFileSync( + path.join(ctx.root, "config/database/schema.ts"), + schemaSource("sales"), + ); + + const { steps } = await init(ctx, { schema: "none" }); + + expect(pool.query).toHaveBeenCalledWith( + expect.stringContaining("nspowner"), + ["sales"], + ); + expect(steps).toContain( + `Handing schema sales in database databricks_postgres on ${devBranchId(USER)} to you: deleting this branch's copy of role ${SP}, which owns it, with its objects reassigned to you`, + ); + }); + + test("creates a read-write endpoint that scales to zero after 5 idle minutes when the branch has none", async () => { + const lakebase = project(); + lakebase.autoEndpoint = false; + + const { report, steps } = await init(appOn(lakebase), { schema: "none" }); + + expect(lakebase.calls).toContain( + `createEndpoint ${DEV} primary suspend=300s`, + ); + expect(report.endpointCreated).toBe(true); + expect(steps).toContain( + `Creating a read-write endpoint on ${devBranchId(USER)} that suspends after 5 idle minutes`, + ); + }); + + test("creates your role when missing, and adds the superuser membership to one without it", async () => { + schemaOwner = USER; + const lakebase = new FakeLakebase(); + lakebase.addBranch(PRODUCTION, { default: true }); + + const created = await init(appOn(lakebase), { schema: "none" }); + expect(created.report.role.action).toBe("created"); + expect(lakebase.calls).toContain( + `createRole ${DEV} ${JSON.stringify({ identity_type: "USER", postgres_role: USER, auth_method: "LAKEBASE_OAUTH_V1", membership_roles: ["DATABRICKS_SUPERUSER"] })}`, + ); + + const other = new FakeLakebase(); + other.addBranch(PRODUCTION, { default: true }); + other.addBranch(DEV); + other.addRole(DEV, "alice", USER, []); + const updated = await init(appOn(other), { schema: "none" }); + expect(updated.report.role.action).toBe("updated"); + expect(other.calls).toContain( + `setRoleMemberships ${DEV}/roles/alice DATABRICKS_SUPERUSER`, + ); + }); + + test("a 403 on branch creation names the CAN MANAGE requirement", async () => { + const lakebase = project(); + lakebase.createBranchError = forbidden(); + + await expect(init(appOn(lakebase))).rejects.toThrow( + `Creating ${devBranchId(USER)} needs CAN MANAGE on ${P}. Ask a project owner for it, then rerun \`appkit db init\`.`, + ); + }); + + test("when another run creates the dev branch first, the 409 reuses it", async () => { + const lakebase = project(); + const api = lakebase.api(); + const create = api.createBranch.bind(api); + api.createBranch = async (...args) => { + await create(...args); + throw conflict(DEV); + }; + const app = appOn(lakebase); + const ctx = fakeContext({ + root: app.root, + env: app.env, + userName: USER, + api, + }); + + const { report } = await init(ctx, { schema: "none" }); + + expect(report).toMatchObject({ branch: DEV, created: false }); + }); + + test("a deleted dev branch, or a 409 for one that is gone, says so instead of forking again", async () => { + const deleted = project(); + deleted.addBranch(DEV, { current_state: "DELETED" }); + await expect(init(appOn(deleted))).rejects.toThrow( + `Your dev branch ${DEV} was deleted or expired. Restore it, or delete it for good, then rerun \`appkit db init\`. \`appkit db init --session\` sets up a session branch meanwhile.`, + ); + + const held = project(); + held.createBranchError = conflict(DEV); + await expect(init(appOn(held))).rejects.toThrow( + `Your dev branch ${DEV} was deleted or expired.`, + ); + }); +}); + +describe("runInit — branch limit", () => { + test("stops before creating anything and lists your other branches with the delete command", async () => { + const lakebase = project(); + const old = session("old1"); + lakebase.addBranch(old, { expire_time: "2099-01-01T00:00:00Z" }); + lakebase.addBranch( + `${P}/branches/${sessionBranchId("bob@example.com", "b1")}`, + ); + lakebase.createBranchError = new LakebaseApiError( + "create branch", + "unarchived branches limit exceeded [TraceId: abc]", + { statusCode: 400, errorCode: "BAD_REQUEST" }, + ); + const ctx = appOn(lakebase); + const before = fs.readFileSync(path.join(ctx.root, ".env"), "utf-8"); + + await expect(init(ctx)).rejects.toThrow( + `Project ${P} is at its branch limit, so init created nothing. Delete one of your other branches, then rerun \`appkit db init\`:\n databricks postgres delete-branch ${old} # expires 2099-01-01T00:00:00Z`, + ); + expect(fs.readFileSync(path.join(ctx.root, ".env"), "utf-8")).toBe(before); + }); + + test("with none of your own, it says to ask a project owner", async () => { + const lakebase = project(); + lakebase.createBranchError = new LakebaseApiError( + "create branch", + "unarchived branches limit exceeded", + { statusCode: 400 }, + ); + + await expect(init(appOn(lakebase))).rejects.toThrow( + "none of its branches are yours. Ask a project owner to delete unused branches", + ); + }); +}); + +describe("runInit — session branch", () => { + test("creates session-- with a 7-day TTL and stores APPKIT_DB_SESSION", async () => { + const lakebase = project(); + const ctx = appOn(lakebase); + + const { report, steps } = await init(ctx, { + mode: "session", + schema: "none", + }); + + const id = envFile(ctx.root).APPKIT_DB_SESSION as string; + expect(id).toMatch(/^[0-9a-z]{8}$/); + expect(report).toMatchObject({ + mode: "session", + branch: session(id), + created: true, + }); + expect(report.expireTime).not.toBeNull(); + expect(lakebase.calls).toContain( + `createBranch ${sessionBranchId(USER, id)} ${JSON.stringify({ source_branch: PRODUCTION, ttl: "604800s" })}`, + ); + expect(steps[0]).toBe( + `Creating ${sessionBranchId(USER, id)} from production, expiring after 7 days`, + ); + }); + + test("--ttl sets a new session branch's TTL; a dev branch takes none, and a bad value fails before any call", async () => { + const lakebase = project(); + const { report } = await init(appOn(lakebase), { + mode: "session", + ttl: "1h", + schema: "none", + }); + expect(lakebase.calls).toContain( + `createBranch ${report.branch.split("/").pop()} ${JSON.stringify({ source_branch: PRODUCTION, ttl: "3600s" })}`, + ); + + const untouched = project(); + await expect(init(appOn(untouched), { ttl: "1h" })).rejects.toThrow( + "--ttl applies to session branches only; add --session.", + ); + await expect( + init(appOn(untouched), { mode: "session", ttl: "soon" }), + ).rejects.toThrow(/^Invalid --ttl "soon"/); + expect(untouched.calls).toEqual([]); + }); + + test("the session id is in .env before the branch is created, so a failed run's rerun reuses it", async () => { + const lakebase = project(); + lakebase.createBranchError = new Error("network down"); + const ctx = appOn(lakebase); + + await expect(init(ctx, { mode: "session" })).rejects.toThrow( + "network down", + ); + const id = envFile(ctx.root).APPKIT_DB_SESSION as string; + expect(id).toMatch(/^[0-9a-z]{8}$/); + + lakebase.createBranchError = undefined; + const rerun = fakeContext({ + root: ctx.root, + env: { ...ctx.env, APPKIT_DB_SESSION: id }, + userName: USER, + api: lakebase.api(), + }); + const { report } = await init(rerun, { mode: "session", schema: "none" }); + expect(report.branch).toBe(session(id)); + }); + + test("a rerun without --session keeps the session branch .env points at", async () => { + schemaOwner = USER; + const lakebase = project(); + lakebase.addBranch(session("k2x9q0ab"), { + expire_time: "2099-01-01T00:00:00Z", + }); + lakebase.addRole(session("k2x9q0ab"), "alice", USER, [ + "DATABRICKS_SUPERUSER", + ]); + const ctx = appOn(lakebase, { + APPKIT_DB_SESSION: "k2x9q0ab", + LAKEBASE_ENDPOINT: `${session("k2x9q0ab")}/endpoints/primary`, + }); + + const { report } = await init(ctx, { schema: "none" }); + + expect(report).toMatchObject({ + mode: "session", + branch: session("k2x9q0ab"), + created: false, + expireTime: "2099-01-01T00:00:00Z", + }); + }); + + test("an expired session branch, or a 409 for one that is gone, gets a fresh session id", async () => { + const lakebase = project(); + lakebase.addBranch(session("k2x9q0ab"), { + expire_time: "2000-01-01T00:00:00Z", + }); + const ctx = appOn(lakebase, { APPKIT_DB_SESSION: "k2x9q0ab" }); + + const { report } = await init(ctx, { mode: "session", schema: "none" }); + + const id = envFile(ctx.root).APPKIT_DB_SESSION as string; + expect(id).not.toBe("k2x9q0ab"); + expect(report).toMatchObject({ branch: session(id), created: true }); + + const held = project(); + const api = held.api(); + const create = api.createBranch.bind(api); + let first = true; + api.createBranch = async (...args) => { + if (first) { + first = false; + throw conflict(args[1]); + } + return create(...args); + }; + const heldCtx = fakeContext({ + root: tempApp({ + "databricks.yml": databricksYml( + PRODUCTION, + `${PRODUCTION}/databases/db`, + ), + }), + env: { LAKEBASE_PROJECT: P, APPKIT_DB_SESSION: "k2x9q0ab" }, + userName: USER, + api, + }); + const retried = await init(heldCtx, { mode: "session", schema: "none" }); + expect(retried.report.branch).not.toBe(session("k2x9q0ab")); + expect(retried.report.created).toBe(true); + }); +}); + +describe("runInit — branches it never sets up", () => { + test("a dev branch databricks.yml deploys to is refused before any API call or .env write", async () => { + const lakebase = project(); + const ctx = appOn(lakebase, {}, { deploy: DEV }); + const before = fs.readFileSync(path.join(ctx.root, ".env"), "utf-8"); + + await expect(init(ctx)).rejects.toThrow( + `Refusing to set up ${devBranchId(USER)}: it is a branch databricks.yml deploys to, and the deployed app migrates it itself.`, + ); + expect(lakebase.calls).toEqual([]); + expect(fs.readFileSync(path.join(ctx.root, ".env"), "utf-8")).toBe(before); + }); + + test.each([ + [{ default: true }, "it is the project's default branch"], + [{ is_protected: true }, "it is a protected branch"], + ])( + "an existing dev branch that became %j is refused before any change", + async (status, reason) => { + const lakebase = project(); + lakebase.addBranch(DEV, status); + const ctx = appOn(lakebase); + const before = fs.readFileSync(path.join(ctx.root, ".env"), "utf-8"); + + await expect(init(ctx)).rejects.toThrow( + `Refusing to set up ${devBranchId(USER)}: ${reason}.`, + ); + expect(lakebase.calls).toEqual([`getBranch ${DEV}`]); + expect(fs.readFileSync(path.join(ctx.root, ".env"), "utf-8")).toBe( + before, + ); + }, + ); +}); + +describe("runInit — hand-off failures", () => { + test("an owner the role API does not manage is named, with the fix", async () => { + schemaOwner = "legacy_owner"; + + await expect(init(appOn(project()))).rejects.toThrow( + `Cannot hand schema app in database databricks_postgres on ${devBranchId(USER)} to you: it belongs to "legacy_owner", which is not a Lakebase-managed role. .env already points at the branch.`, + ); + }); + + test("a failed role delete lists what blocks it", async () => { + const lakebase = project(); + lakebase.deleteRoleError = new LakebaseApiError( + "delete role", + "INTERNAL_ERROR", + { statusCode: 500 }, + ); + blockers = [ + { database: "analytics", count: 3 }, + { database: "databricks_postgres", count: 1 }, + ]; + + await expect(init(appOn(lakebase))).rejects.toThrow( + `These objects and privileges of "${SP}" block the delete: database analytics (3), database databricks_postgres (1). .env already points at the branch. Drop them or reassign them to you, then rerun \`appkit db init\`.`, + ); + expect(pool.query).toHaveBeenCalledWith( + expect.stringContaining("pg_shdepend"), + [SP], + ); + }); + + test("without blockers it says to rerun", async () => { + const lakebase = project(); + lakebase.deleteRoleError = new Error("timeout"); + + await expect(init(appOn(lakebase))).rejects.toThrow( + `deleting role "${SP}" failed (timeout). .env already points at the branch; rerun \`appkit db init\` to retry.`, + ); + }); +}); + +describe("runInit — schema start", () => { + test("without schema.ts the default is introspect, which installs drizzle-kit first", async () => { + const lakebase = project(); + const ctx = appOn( + lakebase, + {}, + { + schema: false, + files: { "package.json": "{}", "pnpm-lock.yaml": "" }, + }, + ); + const runCommand = vi.fn(); + const steps: string[] = []; + + const report = await runInit(ctx, { + runCommand, + onStep: (step) => steps.push(step), + }); + + const install = ["add", "-D", "-E", `drizzle-kit@${DRIZZLE_KIT_VERSION}`]; + expect(runCommand).toHaveBeenCalledWith("pnpm", install, ctx.root); + expect(report.installed).toBe(`pnpm ${install.join(" ")}`); + expect(report.start).toEqual({ + kind: "introspect", + introspect: { branch: DEV, tables: [] }, + }); + expect(runIntrospect.mock.calls[0]?.[1]).toEqual({ interactive: false }); + expect( + steps.indexOf( + `Installing drizzle-kit ${DRIZZLE_KIT_VERSION}, which writes migrations (pnpm ${install.join(" ")})`, + ), + ).toBeLessThan( + steps.indexOf( + `Writing schema.ts from schema app on ${devBranchId(USER)}`, + ), + ); + }); + + test("a failed install says what to run", async () => { + const ctx = appOn( + project(), + {}, + { files: { "package.json": "{}", "package-lock.json": "{}" } }, + ); + + await expect( + runInit(ctx, { + runCommand: () => { + throw new Error("ENOTFOUND registry"); + }, + }), + ).rejects.toThrow( + `Installing drizzle-kit failed. Run \`npm install --save-dev --save-exact drizzle-kit@${DRIZZLE_KIT_VERSION}\`, then rerun \`appkit db init\`.`, + ); + }); + + test("clean drops the schema, then applies the committed migrations without generating; it needs no drizzle-kit", async () => { + const ctx = appOn(project(), {}, { files: { "package.json": "{}" } }); + const runCommand = vi.fn(); + + const { report } = await init(ctx, { schema: "clean", runCommand }); + + expect(runClean).toHaveBeenCalledOnce(); + expect(runMigrate.mock.calls[0]?.[1]).toEqual({ + interactive: false, + generate: false, + }); + expect(runClean.mock.invocationCallOrder[0]).toBeLessThan( + runMigrate.mock.invocationCallOrder[0] ?? 0, + ); + expect(runCommand).not.toHaveBeenCalled(); + expect(report.start).toEqual({ + kind: "clean", + clean: { branch: DEV, schemaName: "app", dependents: [] }, + migrate: migrateReport, + }); + }); + + test.each(["migrate", "clean"] as const)( + "%s without schema.ts fails before any change", + async (start) => { + const lakebase = project(); + const ctx = appOn(lakebase, {}, { schema: false }); + + await expect(init(ctx, { schema: start })).rejects.toThrow( + /^No config\/database\/schema\.ts yet/, + ); + expect(lakebase.calls).toEqual([]); + }, + ); +}); + +describe("runInit — .env", () => { + test("a missing .env gets the host init authenticated with, and never .env.example's placeholders", async () => { + const lakebase = project(); + const root = tempApp({ + ".env.example": + "DATABRICKS_HOST=https://your-workspace.cloud.databricks.com\nPGHOST=your_postgres_host\n", + "databricks.yml": databricksYml( + PRODUCTION, + `${PRODUCTION}/databases/databricks-postgres`, + ), + }); + lakebase.databases.set(`${DEV}/databases/databricks-postgres`, { + name: `${DEV}/databases/databricks-postgres`, + status: { postgres_database: "databricks_postgres" }, + }); + const ctx = fakeContext({ + root, + userName: USER, + api: lakebase.api(), + host: "https://real.cloud.databricks.com", + }); + + await init(ctx, { schema: "none" }); + + expect(envFile(root)).toEqual({ + LAKEBASE_PROJECT: P, + PGHOST: `primary.${devBranchId(USER)}.example.com`, + LAKEBASE_ENDPOINT: `${DEV}/endpoints/primary`, + PGDATABASE: "databricks_postgres", + PGUSER: USER, + PGPORT: "5432", + PGSSLMODE: "require", + DATABRICKS_HOST: "https://real.cloud.databricks.com", + }); + expect(lakebase.calls).toContain( + `getDatabase ${DEV}/databases/databricks-postgres`, + ); + expect(fs.statSync(path.join(root, ".env")).mode & 0o777).toBe(0o600); + }); + + test("the profile init authenticated with wins over the host", async () => { + const lakebase = project(); + const root = tempApp(); + const ctx = fakeContext({ + root, + env: { LAKEBASE_PROJECT: P, DATABRICKS_CONFIG_PROFILE: "staging" }, + userName: USER, + api: lakebase.api(), + host: "https://real.cloud.databricks.com", + }); + + await init(ctx, { schema: "none" }); + + expect(envFile(root).DATABRICKS_CONFIG_PROFILE).toBe("staging"); + expect(envFile(root).DATABRICKS_HOST).toBeUndefined(); + }); + + test("a symlinked .env is refused before any API call", async () => { + const lakebase = project(); + const ctx = appOn(lakebase, {}, { files: { "shared.env": "" } }); + fs.rmSync(path.join(ctx.root, ".env")); + fs.symlinkSync( + path.join(ctx.root, "shared.env"), + path.join(ctx.root, ".env"), + ); + + await expect(init(ctx)).rejects.toThrow( + /^Refusing to rewrite .*\.env: it is a symbolic link/, + ); + expect(lakebase.calls).toEqual([]); + }); +}); + +describe("runInit — project", () => { + test("no project anywhere is a clear error", async () => { + const ctx = fakeContext({ + root: tempApp(), + userName: USER, + api: project().api(), + }); + await expect(init(ctx, { schema: "none" })).rejects.toThrow( + /No Lakebase project configured\. Set LAKEBASE_PROJECT=projects\/ in \.env/, + ); + }); + + test.each([P, "p"])( + "a .env with only the profile and LAKEBASE_PROJECT=%s is enough", + async (value) => { + const lakebase = project(); + const root = tempApp({ + ".env": `DATABRICKS_CONFIG_PROFILE=DEFAULT\nLAKEBASE_PROJECT=${value}\n`, + }); + const ctx = fakeContext({ + root, + env: { DATABRICKS_CONFIG_PROFILE: "DEFAULT", LAKEBASE_PROJECT: value }, + userName: USER, + api: lakebase.api(), + }); + + const { report } = await init(ctx, { schema: "none" }); + + expect(report).toMatchObject({ + branch: DEV, + created: true, + sourceBranch: PRODUCTION, + database: "databricks_postgres", + }); + expect(envFile(root)).toEqual({ + DATABRICKS_CONFIG_PROFILE: "DEFAULT", + LAKEBASE_PROJECT: P, + PGHOST: `primary.${devBranchId(USER)}.example.com`, + LAKEBASE_ENDPOINT: `${DEV}/endpoints/primary`, + PGDATABASE: "databricks_postgres", + PGUSER: USER, + PGPORT: "5432", + PGSSLMODE: "require", + }); + }, + ); + + test("LAKEBASE_PROJECT wins over LAKEBASE_ENDPOINT and databricks.yml; their values from another project are ignored", async () => { + schemaOwner = USER; + const lakebase = project(); + const otherProduction = "projects/q/branches/production"; + const otherDev = `projects/q/branches/${devBranchId(USER)}`; + lakebase.addBranch(otherProduction, { default: true }); + lakebase.addRole(otherProduction, "alice", USER, ["DATABRICKS_SUPERUSER"]); + const ctx = appOn(lakebase, { + LAKEBASE_PROJECT: "q", + PGDATABASE: "p_only_database", + }); + + const { report } = await init(ctx, { schema: "none" }); + + expect(report).toMatchObject({ + branch: otherDev, + sourceBranch: otherProduction, + database: "databricks_postgres", + }); + expect(lakebase.calls.some((call) => call.includes("projects/p/"))).toBe( + false, + ); + expect(envFile(ctx.root)).toMatchObject({ + LAKEBASE_PROJECT: "projects/q", + LAKEBASE_ENDPOINT: `${otherDev}/endpoints/primary`, + PGDATABASE: "databricks_postgres", + }); + }); +}); diff --git a/packages/appkit/src/database/tooling/tests/lakebase-api.test.ts b/packages/appkit/src/database/tooling/tests/lakebase-api.test.ts index 139ee959a..3ce37d925 100644 --- a/packages/appkit/src/database/tooling/tests/lakebase-api.test.ts +++ b/packages/appkit/src/database/tooling/tests/lakebase-api.test.ts @@ -1,7 +1,13 @@ import { describe, expect, test } from "vitest"; import type { WorkspaceRequest } from "../../../workspace-client"; -import { LakebaseApi, LakebaseApiError } from "../lakebase-api"; +import { + isAlreadyExists, + isBranchLimit, + isPermissionDenied, + LakebaseApi, + LakebaseApiError, +} from "../lakebase-api"; /** A request as LakebaseApi sent it, with its JSON body decoded. */ interface Sent { @@ -47,6 +53,8 @@ function client(replies: Reply[]) { const apiError = (statusCode: number, message: string) => Object.assign(new Error(message), { statusCode, errorCode: "X" }); +const fast = { pollIntervalMs: 1 }; + describe("LakebaseApi", () => { test("getBranch: GET /api/2.0/postgres/{name}; 404 → null; other errors keep their status", async () => { const branch = { @@ -58,7 +66,7 @@ describe("LakebaseApi", () => { apiError(404, "NOT_FOUND"), apiError(403, "PERMISSION_DENIED"), ]); - const api = new LakebaseApi(fake); + const api = new LakebaseApi(fake, fast); await expect(api.getBranch("projects/p/branches/b")).resolves.toEqual( branch, @@ -153,4 +161,270 @@ describe("LakebaseApi", () => { api.getEndpoint("projects/p/branches/b/endpoints/primary"), ).rejects.toThrow(/unexpected response \(name: /); }); + + test("lists follow next_page_token", async () => { + const fake = client([ + { branches: [{ name: "a" }], next_page_token: "t2" }, + { branches: [{ name: "b" }] }, + ]); + const branches = await new LakebaseApi(fake, fast).listBranches( + "projects/p", + ); + + expect(branches.map((b) => b.name)).toEqual(["a", "b"]); + expect(fake.requests.map((r) => [r.path, r.query])).toEqual([ + ["/api/2.0/postgres/projects/p/branches", undefined], + ["/api/2.0/postgres/projects/p/branches", { page_token: "t2" }], + ]); + }); + + test("createBranch: POST with branch_id and {spec}; a finished operation returns its response", async () => { + const fake = client([ + { + name: "projects/p/branches/dev-a/operations/1", + done: true, + response: { name: "projects/p/branches/dev-a" }, + }, + ]); + const branch = await new LakebaseApi(fake, fast).createBranch( + "projects/p", + "dev-a", + { + source_branch: "projects/p/branches/production", + no_expiry: true, + }, + ); + + expect(branch).toEqual({ name: "projects/p/branches/dev-a" }); + expect(fake.requests[0]).toMatchObject({ + path: "/api/2.0/postgres/projects/p/branches", + method: "POST", + query: { branch_id: "dev-a" }, + payload: { + spec: { + source_branch: "projects/p/branches/production", + no_expiry: true, + }, + }, + }); + expect(fake.requests[0]?.headers.get("Content-Type")).toBe( + "application/json", + ); + }); + + test("polls GET /api/2.0/postgres/{operation.name} until the operation is done", async () => { + const op = "projects/p/branches/b/operations/7"; + const fake = client([ + { name: op, done: false }, + { name: op, done: false }, + { + name: op, + done: true, + response: { name: "projects/p/branches/b/endpoints/primary" }, + }, + ]); + const endpoint = await new LakebaseApi(fake, fast).createEndpoint( + "projects/p/branches/b", + "primary", + { suspendTimeout: "300s" }, + ); + + expect(endpoint.name).toBe("projects/p/branches/b/endpoints/primary"); + expect(fake.requests[0]).toMatchObject({ + path: "/api/2.0/postgres/projects/p/branches/b/endpoints", + method: "POST", + query: { endpoint_id: "primary" }, + payload: { + spec: { + endpoint_type: "ENDPOINT_TYPE_READ_WRITE", + suspend_timeout_duration: "300s", + }, + }, + }); + expect(fake.requests.slice(1).map((r) => [r.method, r.path])).toEqual([ + ["GET", `/api/2.0/postgres/${op}`], + ["GET", `/api/2.0/postgres/${op}`], + ]); + }); + + test("a failed operation becomes a LakebaseApiError with the operation's message", async () => { + const fake = client([ + { + name: "op", + done: true, + error: { + error_code: "RESOURCE_EXHAUSTED", + message: "branch limit reached", + }, + }, + ]); + await expect( + new LakebaseApi(fake, fast).createBranch("projects/p", "dev-a", { + source_branch: "projects/p/branches/production", + no_expiry: true, + }), + ).rejects.toThrow( + "Lakebase: create branch projects/p/branches/dev-a failed: branch limit reached", + ); + }); + + test("an operation that never finishes times out", async () => { + const fake = client( + Array.from({ length: 50 }, () => ({ name: "op", done: false })), + ); + await expect( + new LakebaseApi(fake, { pollIntervalMs: 1, timeoutMs: 5 }).createEndpoint( + "projects/p/branches/b", + "primary", + { suspendTimeout: "300s" }, + ), + ).rejects.toThrow(/timed out/); + }); + + test("error codes from the call and from a failed operation classify the failure", async () => { + const create = (api: LakebaseApi) => + api + .createBranch("projects/p", "dev-a", { + source_branch: "projects/p/branches/production", + no_expiry: true, + }) + .catch((e: unknown) => e); + const conflict = await create( + new LakebaseApi( + client([ + { + name: "op", + done: true, + error: { error_code: "ALREADY_EXISTS", message: "exists" }, + }, + ]), + fast, + ), + ); + const denied = await create( + new LakebaseApi(client([apiError(403, "PERMISSION_DENIED")]), fast), + ); + const limit = await create( + new LakebaseApi( + client([ + Object.assign( + new Error("unarchived branches limit exceeded [TraceId: abc]"), + { statusCode: 400, errorCode: "BAD_REQUEST" }, + ), + ]), + fast, + ), + ); + + expect(conflict).toMatchObject({ errorCode: "ALREADY_EXISTS" }); + expect(isAlreadyExists(conflict)).toBe(true); + expect(isPermissionDenied(conflict)).toBe(false); + expect(isPermissionDenied(denied)).toBe(true); + expect(isAlreadyExists(denied)).toBe(false); + expect(limit).toMatchObject({ statusCode: 400, errorCode: "BAD_REQUEST" }); + expect(isBranchLimit(limit)).toBe(true); + expect(isBranchLimit(denied)).toBe(false); + }); + + test("deleteBranch: DELETE the branch; one that is already gone counts as deleted", async () => { + const fake = client([ + { name: "op", done: true }, + apiError(404, "NOT_FOUND"), + { + name: "op", + done: true, + error: { error_code: "RESOURCE_DOES_NOT_EXIST", message: "gone" }, + }, + apiError(403, "PERMISSION_DENIED"), + ]); + const api = new LakebaseApi(fake, fast); + const name = "projects/p/branches/session-abc123-x"; + + await api.deleteBranch(name); + await api.deleteBranch(name); + await api.deleteBranch(name); + await expect(api.deleteBranch(name)).rejects.toMatchObject({ + statusCode: 403, + }); + + expect(fake.requests[0]).toMatchObject({ + path: `/api/2.0/postgres/${name}`, + method: "DELETE", + }); + expect(fake.requests[0]?.query).toBeUndefined(); + }); + + test("createRole posts {spec} without a role_id", async () => { + const fake = client([ + { + name: "op", + done: true, + response: { name: "projects/p/branches/b/roles/r" }, + }, + ]); + await new LakebaseApi(fake, fast).createRole("projects/p/branches/b", { + identity_type: "USER", + postgres_role: "alice@example.com", + auth_method: "LAKEBASE_OAUTH_V1", + membership_roles: ["DATABRICKS_SUPERUSER"], + }); + + expect(fake.requests[0]).toEqual( + expect.objectContaining({ + path: "/api/2.0/postgres/projects/p/branches/b/roles", + method: "POST", + payload: { + spec: { + identity_type: "USER", + postgres_role: "alice@example.com", + auth_method: "LAKEBASE_OAUTH_V1", + membership_roles: ["DATABRICKS_SUPERUSER"], + }, + }, + }), + ); + expect(fake.requests[0]?.query).toBeUndefined(); + }); + + test("setRoleMemberships: PATCH with update_mask=spec.membership_roles", async () => { + const fake = client([ + { + name: "op", + done: true, + response: { name: "projects/p/branches/b/roles/r" }, + }, + ]); + await new LakebaseApi(fake, fast).setRoleMemberships( + "projects/p/branches/b/roles/r", + ["DATABRICKS_SUPERUSER"], + ); + + expect(fake.requests[0]).toMatchObject({ + path: "/api/2.0/postgres/projects/p/branches/b/roles/r", + method: "PATCH", + query: { update_mask: "spec.membership_roles" }, + payload: { spec: { membership_roles: ["DATABRICKS_SUPERUSER"] } }, + }); + }); + + test("deleteRole: DELETE with reassign_owned_to = the target role's resource name", async () => { + const fake = client([ + { name: "op", done: false }, + { name: "op", done: true }, + ]); + await new LakebaseApi(fake, fast).deleteRole( + "projects/p/branches/b/roles/sp", + { + reassignOwnedTo: "projects/p/branches/b/roles/alice", + }, + ); + + expect(fake.requests[0]).toMatchObject({ + path: "/api/2.0/postgres/projects/p/branches/b/roles/sp", + method: "DELETE", + query: { reassign_owned_to: "projects/p/branches/b/roles/alice" }, + }); + expect(fake.requests[0]?.payload).toBeUndefined(); + expect(fake.requests).toHaveLength(2); + }); }); diff --git a/packages/appkit/src/database/tooling/tests/migrate.test.ts b/packages/appkit/src/database/tooling/tests/migrate.test.ts index 15056e199..4246a7964 100644 --- a/packages/appkit/src/database/tooling/tests/migrate.test.ts +++ b/packages/appkit/src/database/tooling/tests/migrate.test.ts @@ -182,6 +182,29 @@ describe("runMigrate", () => { expect(fs.existsSync(report.typesFile)).toBe(true); }); + test("generate: false applies the committed migrations without writing new ones", async () => { + const ctx = app(`notes: table("notes", { id: id() })`); + resolveEnvBranch.mockResolvedValue(null); + await runMigrate(ctx, { interactive: false }); + writeSchema(ctx.root, `notes: table("notes", { id: id(), body: text() })`); + + resolveEnvBranch.mockResolvedValue(devBranch); + applyPendingMigrations.mockResolvedValue({ + applied: ["0000_migration"], + adopted: [], + }); + const report = await runMigrate(ctx, { + interactive: false, + generate: false, + }); + + expect(report.generated).toBeNull(); + expect(report.applied).toEqual(["0000_migration"]); + expect( + fs.readdirSync(path.join(ctx.root, "config/database/migrations")), + ).toEqual(["0000_migration.sql", "meta"]); + }); + test("reports a baseline adopted without running it", async () => { const ctx = app(`notes: table("notes", { id: id() })`); resolveEnvBranch.mockResolvedValue(null); diff --git a/packages/appkit/src/database/tooling/tests/naming.test.ts b/packages/appkit/src/database/tooling/tests/naming.test.ts index 423f33ad0..b2ab9bb12 100644 --- a/packages/appkit/src/database/tooling/tests/naming.test.ts +++ b/packages/appkit/src/database/tooling/tests/naming.test.ts @@ -1,6 +1,14 @@ import { describe, expect, test } from "vitest"; -import { devBranchId, isValidSessionId, sessionBranchId } from "../naming"; +import { + devBranchId, + isUsersBranchId, + isValidSessionId, + newSessionId, + parseTtl, + SESSION_BRANCH_TTL, + sessionBranchId, +} from "../naming"; const BRANCH_ID = /^[a-z][a-z0-9-]{0,62}$/; @@ -31,9 +39,52 @@ describe("branch naming", () => { ); }); - test("session branches are session-; ids are lowercase base36", () => { - expect(sessionBranchId("k2x9q0ab")).toBe("session-k2x9q0ab"); + test("session branches are session--; ids are lowercase base36", () => { + const sessionId = newSessionId(); + expect(sessionId).toMatch(/^[0-9a-z]{8}$/); + const hash = devBranchId("alice@example.com").slice(-6); + expect(sessionBranchId("Alice@example.com", sessionId)).toBe( + `session-${hash}-${sessionId}`, + ); expect(isValidSessionId("Bad-Id")).toBe(false); - expect(() => sessionBranchId("Bad-Id")).toThrow(/APPKIT_DB_SESSION/); + expect(() => sessionBranchId("alice@example.com", "Bad-Id")).toThrow( + /APPKIT_DB_SESSION/, + ); + expect(SESSION_BRANCH_TTL).toBe("604800s"); + }); + + test("the longest session id still makes a valid branch ID", () => { + const id = sessionBranchId("alice@example.com", "a".repeat(48)); + expect(id).toHaveLength(63); + expect(id).toMatch(BRANCH_ID); + expect(isValidSessionId("a".repeat(49))).toBe(false); + }); + + test("a copied .env names a session branch only its creator owns", () => { + const alice = "alice@example.com"; + const bob = "bob@example.com"; + expect(sessionBranchId(alice, "abc")).not.toBe(sessionBranchId(bob, "abc")); + expect(isUsersBranchId(alice, sessionBranchId(alice, "abc"))).toBe(true); + expect(isUsersBranchId(alice, devBranchId(alice))).toBe(true); + expect(isUsersBranchId(alice, sessionBranchId(bob, "abc"))).toBe(false); + expect(isUsersBranchId(alice, devBranchId(bob))).toBe(false); + expect(isUsersBranchId(alice, "production")).toBe(false); + }); +}); + +describe("parseTtl", () => { + test.each([ + ["90s", "90s"], + ["30m", "1800s"], + [" 1h ", "3600s"], + ["7d", "604800s"], + ])("%s is %s", (value, expected) => { + expect(parseTtl(value)).toBe(expected); + }); + + test.each(["0h", "1w", "1.5h", "-1h", "h", ""])("refuses %j", (value) => { + expect(() => parseTtl(value)).toThrow( + `Invalid --ttl ${JSON.stringify(value)}: use a positive number of seconds, minutes, hours or days, e.g. 30m, 1h or 7d.`, + ); }); }); diff --git a/packages/appkit/src/database/tooling/tests/targets.test.ts b/packages/appkit/src/database/tooling/tests/targets.test.ts index ac159a7df..6be36190e 100644 --- a/packages/appkit/src/database/tooling/tests/targets.test.ts +++ b/packages/appkit/src/database/tooling/tests/targets.test.ts @@ -1,8 +1,10 @@ import { afterEach, describe, expect, test } from "vitest"; import { + parseProject, parseResourceName, readDeployBranches, + readDeployTarget, readLocalTarget, } from "../targets"; import { cleanupTempApps, databricksYml, tempApp } from "./fakes"; @@ -35,16 +37,35 @@ describe("parseResourceName", () => { }); }); +describe("parseProject", () => { + test("accepts projects/ or a bare , normalized to the resource name", () => { + expect(parseProject("projects/my-app")).toBe("projects/my-app"); + expect(parseProject(" my-app ")).toBe("projects/my-app"); + }); + + test("placeholders, branch paths and junk are null", () => { + expect(parseProject("your_postgres_project")).toBeNull(); + expect(parseProject("projects/p/branches/b")).toBeNull(); + expect(parseProject("")).toBeNull(); + expect(parseProject(undefined)).toBeNull(); + }); +}); + describe("readLocalTarget", () => { - test("reads the .env endpoint and this checkout's session", () => { + test("reads a real .env target", () => { expect( readLocalTarget({ + LAKEBASE_PROJECT: "p", LAKEBASE_ENDPOINT: "projects/p/branches/b/endpoints/primary", PGHOST: "ep.example.com", + PGDATABASE: "databricks_postgres", + PGUSER: "alice@example.com", APPKIT_DB_SESSION: "k2x9q0ab", }), ).toEqual({ + project: "projects/p", endpoint: "projects/p/branches/b/endpoints/primary", + database: "databricks_postgres", session: "k2x9q0ab", }); }); @@ -52,10 +73,18 @@ describe("readLocalTarget", () => { test("placeholders from .env.example read as not configured", () => { expect( readLocalTarget({ + LAKEBASE_PROJECT: "your_postgres_project", LAKEBASE_ENDPOINT: "your_postgres_endpointPath", + PGHOST: "your_postgres_host", + PGDATABASE: "your_postgres_databaseName", APPKIT_DB_SESSION: "NOT VALID", }), - ).toEqual({ endpoint: null, session: null }); + ).toEqual({ + project: null, + endpoint: null, + database: null, + session: null, + }); }); }); @@ -109,3 +138,41 @@ targets: ); }); }); + +describe("readDeployTarget", () => { + test("reads the default target's postgres variables", () => { + const root = tempApp({ + "databricks.yml": databricksYml( + "projects/p/branches/production", + "projects/p/branches/production/databases/databricks-postgres", + ), + }); + expect(readDeployTarget(root)).toEqual({ + branch: "projects/p/branches/production", + database: "projects/p/branches/production/databases/databricks-postgres", + }); + }); + + test("falls back to top-level variable defaults", () => { + const root = tempApp({ + "databricks.yml": `variables: + postgres_branch: + default: projects/p/branches/main +targets: + dev: + default: true +`, + }); + expect(readDeployTarget(root)).toEqual({ + branch: "projects/p/branches/main", + database: null, + }); + }); + + test("null without databricks.yml or without a postgres binding", () => { + expect(readDeployTarget(tempApp())).toBeNull(); + expect( + readDeployTarget(tempApp({ "databricks.yml": "bundle:\n name: x\n" })), + ).toBeNull(); + }); +}); diff --git a/packages/shared/src/cli/commands/db/clean.ts b/packages/shared/src/cli/commands/db/clean.ts new file mode 100644 index 000000000..79a56f09e --- /dev/null +++ b/packages/shared/src/cli/commands/db/clean.ts @@ -0,0 +1,81 @@ +import { Command } from "commander"; + +import { loadDatabaseTooling } from "./load-tooling"; +import { + fail, + info, + isInteractive, + ok, + printCleanReport, + shortName, + warn, +} from "./output"; + +interface CleanFlags { + readonly yes?: boolean; + readonly deleteBranch?: boolean; +} + +async function runDbClean(flags: CleanFlags): Promise { + try { + const tooling = await loadDatabaseTooling(); + const ctx = tooling.createDbContext(); + if (flags.deleteBranch) { + const { branch } = await tooling.deleteSessionBranch(ctx); + ok(`Deleted ${shortName(branch)}`); + info("`appkit db init --session` creates a new one"); + return; + } + if (!flags.yes) { + if (!isInteractive()) { + throw new Error( + "Refusing to clean without confirmation in a non-interactive shell. Pass --yes.", + ); + } + const preview = await tooling.previewClean(ctx); + if (preview.dependents.length > 0) { + warn( + `Dropping schema ${preview.schemaName} also removes these objects outside it, which depend on it: ${preview.dependents.join(", ")}`, + ); + } + const { confirm, isCancel } = await import("@clack/prompts"); + const answer = await confirm({ + message: `Remove every object in schema ${preview.schemaName} on ${shortName(preview.branch)}, including the migration journal?`, + initialValue: false, + }); + if (isCancel(answer) || !answer) { + info("Cancelled"); + process.exitCode = 1; + return; + } + } + printCleanReport(await tooling.runClean(ctx)); + info("`appkit db migrate` recreates it from the committed migrations"); + } catch (error) { + fail(error); + } +} + +export const dbCleanCommand = new Command("clean") + .description( + "Remove every object in the app schema, including the migration journal, on your dev or session branch", + ) + .option("-y, --yes", "skip the confirmation") + .option( + "--delete-branch", + "delete this checkout's session branch instead, e.g. at the end of a CI job", + ) + .addHelpText( + "after", + ` +Only branches \`appkit db init\` created for you (dev-… or session-…) can be cleaned; +the default, protected and deploy branches are always refused. Objects outside +the schema that depend on it, such as a view in another schema, are dropped +with it and listed. --delete-branch deletes session branches only; your dev +branch stays. + +Examples: + $ appkit db clean + $ appkit db init --session --ttl 1h --yes && npm test && appkit db clean --delete-branch`, + ) + .action(runDbClean); diff --git a/packages/shared/src/cli/commands/db/index.ts b/packages/shared/src/cli/commands/db/index.ts index 35ff2f721..ed811c91c 100644 --- a/packages/shared/src/cli/commands/db/index.ts +++ b/packages/shared/src/cli/commands/db/index.ts @@ -1,7 +1,10 @@ import { Command } from "commander"; +import { dbCleanCommand } from "./clean"; +import { dbInitCommand } from "./init"; import { dbIntrospectCommand } from "./introspect"; import { dbMigrateCommand } from "./migrate"; +import { dbSeedCommand } from "./seed"; /** * Lakebase database lifecycle for the DatabasePlugin. The logic lives in @@ -9,5 +12,8 @@ import { dbMigrateCommand } from "./migrate"; */ export const dbCommand = new Command("db") .description("Lakebase database lifecycle") + .addCommand(dbInitCommand) .addCommand(dbMigrateCommand) - .addCommand(dbIntrospectCommand); + .addCommand(dbIntrospectCommand) + .addCommand(dbCleanCommand) + .addCommand(dbSeedCommand); diff --git a/packages/shared/src/cli/commands/db/init-clean-seed.test.ts b/packages/shared/src/cli/commands/db/init-clean-seed.test.ts new file mode 100644 index 000000000..90dcdf53f --- /dev/null +++ b/packages/shared/src/cli/commands/db/init-clean-seed.test.ts @@ -0,0 +1,377 @@ +import { + afterEach, + beforeEach, + describe, + expect, + type MockInstance, + test, + vi, +} from "vitest"; + +const tooling = vi.hoisted(() => ({ + createDbContext: vi.fn(() => ({ root: "/app", env: {} })), + runInit: vi.fn(), + previewClean: vi.fn(), + runClean: vi.fn(), + deleteSessionBranch: vi.fn(), + runSeed: vi.fn(), +})); +const prompts = vi.hoisted(() => ({ + confirm: vi.fn(), + isCancel: vi.fn((value: unknown) => typeof value === "symbol"), +})); + +vi.mock("@databricks/appkit/database-tooling", () => tooling); +vi.mock("@clack/prompts", () => prompts); + +import { dbCommand } from "./index"; + +let log: MockInstance; +let error: MockInstance; +const output = (spy: MockInstance) => + spy.mock.calls.map((call) => call.join(" ")).join("\n"); + +function setTTY(value: boolean | undefined) { + Object.defineProperty(process.stdin, "isTTY", { value, configurable: true }); + Object.defineProperty(process.stdout, "isTTY", { value, configurable: true }); +} +const ttyBefore = { stdin: process.stdin.isTTY, stdout: process.stdout.isTTY }; + +const migrate = { + generated: null, + applied: ["0000_migration"], + typesFile: `${process.cwd()}/shared/appkit-types/database.d.ts`, +}; +const report = { + mode: "dev", + branch: "projects/p/branches/dev-alice-3f9a2c", + created: true, + sourceBranch: "projects/p/branches/production", + expireTime: null, + endpoint: "projects/p/branches/dev-alice-3f9a2c/endpoints/primary", + endpointCreated: false, + host: "ep-1.example.com", + database: "databricks_postgres", + user: "alice@example.com", + role: { + name: "projects/p/branches/dev-alice-3f9a2c/roles/alice", + action: "kept", + }, + handedOff: "11111111-2222-3333-4444-555555555555", + envFile: "/app/.env", + installed: null, + start: { kind: "migrate", migrate }, +}; + +/** Commander keeps parsed option values on the command; each test parses afresh. */ +function resetOptionValues() { + for (const command of dbCommand.commands) { + for (const option of command.options) { + command.setOptionValueWithSource( + option.attributeName(), + undefined, + "default", + ); + } + } +} + +const init = (...args: string[]) => + dbCommand.parseAsync(["init", ...args], { from: "user" }); +const clean = (...args: string[]) => + dbCommand.parseAsync(["clean", ...args], { from: "user" }); +const initOptions = () => tooling.runInit.mock.calls.at(-1)?.[1]; + +beforeEach(() => { + resetOptionValues(); + for (const mock of Object.values(tooling)) mock.mockReset(); + prompts.confirm.mockReset(); + tooling.createDbContext.mockReturnValue({ root: "/app", env: {} }); + tooling.runInit.mockResolvedValue(report); + log = vi.spyOn(console, "log").mockImplementation(() => undefined); + error = vi.spyOn(console, "error").mockImplementation(() => undefined); +}); + +afterEach(() => { + vi.restoreAllMocks(); + setTTY(undefined); + Object.defineProperty(process.stdin, "isTTY", { + value: ttyBefore.stdin, + configurable: true, + }); + Object.defineProperty(process.stdout, "isTTY", { + value: ttyBefore.stdout, + configurable: true, + }); + process.exitCode = undefined; +}); + +describe("appkit db init", () => { + test("registers init, migrate, introspect, clean and seed", () => { + expect(dbCommand.commands.map((c) => c.name())).toEqual([ + "init", + "migrate", + "introspect", + "clean", + "seed", + ]); + }); + + test("init only sets up a dev or session branch: there is no option to pick another one", () => { + const command = dbCommand.commands.find((c) => c.name() === "init"); + expect(command?.options.map((o) => o.long)).toEqual([ + "--session", + "--dev", + "--ttl", + "--schema", + "--yes", + ]); + }); + + test("asks nothing, even on a TTY: the tooling picks the mode and the start", async () => { + setTTY(true); + await init(); + + expect(prompts.confirm).not.toHaveBeenCalled(); + expect(tooling.runInit).toHaveBeenCalledWith( + { root: "/app", env: {} }, + { interactive: true, onStep: expect.any(Function) }, + ); + }); + + test("flags map to the options; --yes turns off the rename prompts", async () => { + await init("--session", "--ttl", "1h", "--schema", "none", "--yes"); + + expect(initOptions()).toEqual({ + mode: "session", + schema: "none", + ttl: "1h", + interactive: false, + onStep: expect.any(Function), + }); + + resetOptionValues(); + await init("--dev"); + expect(initOptions()).toMatchObject({ mode: "dev" }); + }); + + test("--session and --dev conflict", async () => { + vi.spyOn(process.stderr, "write").mockImplementation(() => true); + const exit = vi.spyOn(process, "exit").mockImplementation(() => { + throw new Error("exit"); + }); + await expect(init("--session", "--dev")).rejects.toThrow("exit"); + expect(exit).toHaveBeenCalled(); + expect(tooling.runInit).not.toHaveBeenCalled(); + }); + + test("prints each step as it happens", async () => { + tooling.runInit.mockImplementation(async (_ctx, options) => { + options.onStep("Creating dev-alice-3f9a2c from production"); + return report; + }); + await init(); + expect(output(log)).toContain("Creating dev-alice-3f9a2c from production"); + }); + + test("prints what it did", async () => { + await init(); + + const printed = output(log); + expect(printed).toContain( + "Branch dev-alice-3f9a2c (created from production)", + ); + expect(printed).toContain("Endpoint primary: ep-1.example.com"); + expect(printed).toContain("Role alice@example.com (DATABRICKS_SUPERUSER)"); + expect(printed).toContain( + "The app schema is yours: it belonged to 11111111-2222-3333-4444-555555555555", + ); + expect(printed).toContain( + "Pointed .env at dev-alice-3f9a2c (database databricks_postgres)", + ); + expect(printed).toContain("Applied 0000_migration to the .env branch"); + expect(printed).not.toContain("Installed drizzle-kit"); + }); + + test("a session branch reports when it expires, and an install is reported", async () => { + tooling.runInit.mockResolvedValue({ + ...report, + mode: "session", + branch: "projects/p/branches/session-3f9a2c-k2x9", + created: false, + expireTime: "2026-10-26T12:00:00Z", + handedOff: null, + installed: "npm install --save-dev --save-exact drizzle-kit@0.31.10", + start: { kind: "none" }, + }); + await init(); + + const printed = output(log); + expect(printed).toContain( + "Branch session-3f9a2c-k2x9 (reused, expires 2026-10-26T12:00:00Z)", + ); + expect(printed).toContain( + "Installed drizzle-kit (npm install --save-dev --save-exact drizzle-kit@0.31.10)", + ); + expect(printed).toContain("Left the schema as it is"); + expect(printed).not.toContain("The app schema is yours"); + }); + + test("a clean start reports the drop, what it took along, then the migrations", async () => { + tooling.runInit.mockResolvedValue({ + ...report, + start: { + kind: "clean", + clean: { + branch: report.branch, + schemaName: "app", + dependents: ["view reports.summary"], + }, + migrate, + }, + }); + await init("--schema", "clean", "--yes"); + + const printed = output(log); + expect(printed).toContain("Dropped schema app on dev-alice-3f9a2c"); + expect(printed).toContain( + "Also removed, because they depended on it: view reports.summary", + ); + expect(printed).toContain("Applied 0000_migration to the .env branch"); + }); + + test("--schema clean without a TTY needs --yes, and stops before any API call", async () => { + await init("--schema", "clean"); + + expect(output(error)).toContain( + "Refusing to run `--schema clean` without confirmation in a non-interactive shell. Pass --yes.", + ); + expect(tooling.runInit).not.toHaveBeenCalled(); + expect(process.exitCode).toBe(1); + }); + + test("--schema clean on a TTY asks first; a no cancels with exit code 1", async () => { + setTTY(true); + prompts.confirm.mockResolvedValueOnce(false); + await init("--schema", "clean"); + + expect(prompts.confirm).toHaveBeenCalledOnce(); + expect(tooling.runInit).not.toHaveBeenCalled(); + expect(output(log)).toContain("Cancelled"); + expect(process.exitCode).toBe(1); + + process.exitCode = undefined; + resetOptionValues(); + prompts.confirm.mockResolvedValueOnce(true); + await init("--schema", "clean"); + expect(initOptions()).toMatchObject({ schema: "clean" }); + }); + + test("an error is printed and exits 1", async () => { + tooling.runInit.mockRejectedValue( + new Error( + "Creating dev-a needs CAN MANAGE on projects/p. Ask a project owner for it, then rerun `appkit db init`.", + ), + ); + await init("--yes"); + expect(output(error)).toContain("needs CAN MANAGE on projects/p"); + expect(process.exitCode).toBe(1); + }); +}); + +describe("appkit db clean", () => { + const cleaned = { + branch: "projects/p/branches/dev-a", + schemaName: "app", + dependents: [], + }; + + test("refuses without --yes when there is no TTY, before any API call", async () => { + await clean(); + + expect(output(error)).toContain( + "Refusing to clean without confirmation in a non-interactive shell. Pass --yes.", + ); + expect(tooling.previewClean).not.toHaveBeenCalled(); + expect(tooling.runClean).not.toHaveBeenCalled(); + expect(process.exitCode).toBe(1); + }); + + test("on a TTY it names the branch and what else goes, then asks; no cancels with exit code 1", async () => { + setTTY(true); + tooling.previewClean.mockResolvedValue({ + ...cleaned, + dependents: ["view reports.summary"], + }); + prompts.confirm.mockResolvedValueOnce(false); + await clean(); + + expect(output(log)).toContain( + "Dropping schema app also removes these objects outside it, which depend on it: view reports.summary", + ); + expect(prompts.confirm.mock.calls[0]?.[0].message).toBe( + "Remove every object in schema app on dev-a, including the migration journal?", + ); + expect(tooling.runClean).not.toHaveBeenCalled(); + expect(process.exitCode).toBe(1); + + process.exitCode = undefined; + prompts.confirm.mockResolvedValueOnce(true); + tooling.runClean.mockResolvedValue(cleaned); + await clean(); + expect(tooling.runClean).toHaveBeenCalledOnce(); + expect(output(log)).toContain("Dropped schema app on dev-a"); + }); + + test("--yes cleans without asking; public says only your objects went", async () => { + tooling.runClean.mockResolvedValue({ ...cleaned, schemaName: "public" }); + await clean("--yes"); + + expect(tooling.previewClean).not.toHaveBeenCalled(); + expect(prompts.confirm).not.toHaveBeenCalled(); + expect(output(log)).toContain( + "Removed your objects in schema public on dev-a", + ); + }); + + test("--delete-branch deletes this checkout's session branch without a prompt", async () => { + tooling.deleteSessionBranch.mockResolvedValue({ + branch: "projects/p/branches/session-3f9a2c-k2x9", + }); + await clean("--delete-branch"); + + expect(tooling.runClean).not.toHaveBeenCalled(); + expect(output(log)).toContain("Deleted session-3f9a2c-k2x9"); + expect(process.exitCode).toBeUndefined(); + }); + + test("the guard's refusal is printed and exits 1", async () => { + tooling.runClean.mockRejectedValue( + new Error( + "Refusing to clean production: it is the project's default branch.", + ), + ); + await clean("--yes"); + expect(output(error)).toContain("Refusing to clean production"); + expect(process.exitCode).toBe(1); + }); +}); + +describe("appkit db seed", () => { + test("prints the seed file and branch", async () => { + tooling.runSeed.mockResolvedValue({ + branch: "projects/p/branches/dev-a", + file: `${process.cwd()}/config/database/seed.sql`, + }); + await dbCommand.parseAsync(["seed"], { from: "user" }); + expect(output(log)).toContain("Ran config/database/seed.sql on dev-a"); + }); + + test("a missing seed file exits 1", async () => { + tooling.runSeed.mockRejectedValue( + new Error("No seed file at config/database/seed.sql."), + ); + await dbCommand.parseAsync(["seed"], { from: "user" }); + expect(process.exitCode).toBe(1); + }); +}); diff --git a/packages/shared/src/cli/commands/db/init.ts b/packages/shared/src/cli/commands/db/init.ts new file mode 100644 index 000000000..a31838b8f --- /dev/null +++ b/packages/shared/src/cli/commands/db/init.ts @@ -0,0 +1,143 @@ +import { Command, Option } from "commander"; + +import { loadDatabaseTooling } from "./load-tooling"; +import { + fail, + info, + isInteractive, + ok, + printCleanReport, + printIntrospectReport, + printMigrateReport, + shortName, +} from "./output"; +import type { InitOptions, InitReport, SchemaStart } from "./tooling-contract"; + +interface InitFlags { + readonly session?: boolean; + readonly dev?: boolean; + readonly ttl?: string; + readonly schema?: SchemaStart; + readonly yes?: boolean; +} + +/** --schema clean drops data: a human confirms it, a script passes --yes. */ +async function confirmClean(flags: InitFlags): Promise { + if (flags.schema !== "clean" || flags.yes) return true; + if (!isInteractive()) { + throw new Error( + "Refusing to run `--schema clean` without confirmation in a non-interactive shell. Pass --yes.", + ); + } + const { confirm, isCancel } = await import("@clack/prompts"); + const answer = await confirm({ + message: + "Drop the app schema on your dev or session branch, with every table and row in it, then apply the committed migrations?", + initialValue: false, + }); + return !isCancel(answer) && answer; +} + +function printInitReport(report: InitReport): void { + const branch = shortName(report.branch); + const expires = report.expireTime ? `, expires ${report.expireTime}` : ""; + ok( + report.created + ? `Branch ${branch} (created from ${shortName(report.sourceBranch ?? "the default branch")}${expires})` + : `Branch ${branch} (reused${expires})`, + ); + ok( + `Endpoint ${shortName(report.endpoint)}${report.endpointCreated ? " (created)" : ""}: ${report.host}`, + ); + const role = { + created: "created", + updated: "added to DATABRICKS_SUPERUSER", + kept: "DATABRICKS_SUPERUSER", + }[report.role.action]; + ok(`Role ${report.user} (${role})`); + if (report.handedOff) { + ok(`The app schema is yours: it belonged to ${report.handedOff}`); + } + ok(`Pointed .env at ${branch} (database ${report.database})`); + if (report.installed) ok(`Installed drizzle-kit (${report.installed})`); + + const { start } = report; + if (start.kind === "introspect") printIntrospectReport(start.introspect); + else if (start.kind === "migrate") printMigrateReport(start.migrate); + else if (start.kind === "clean") { + printCleanReport(start.clean); + printMigrateReport(start.migrate); + } else info("Left the schema as it is"); +} + +async function runDbInit(flags: InitFlags): Promise { + try { + if (!(await confirmClean(flags))) { + info("Cancelled"); + process.exitCode = 1; + return; + } + const tooling = await loadDatabaseTooling(); + const mode = flags.session ? "session" : flags.dev ? "dev" : undefined; + const options: InitOptions = { + ...(mode ? { mode } : {}), + ...(flags.schema ? { schema: flags.schema } : {}), + ...(flags.ttl !== undefined ? { ttl: flags.ttl } : {}), + interactive: !flags.yes && isInteractive(), + onStep: info, + }; + printInitReport(await tooling.runInit(tooling.createDbContext(), options)); + } catch (error) { + fail(error); + } +} + +export const dbInitCommand = new Command("init") + .description( + "Set up your Lakebase dev or session branch for local development: create or reuse it, ensure an endpoint and your role, point .env at it, then set up the schema", + ) + .addOption( + new Option( + "--session", + "use this checkout's session branch, created with a TTL", + ).conflicts("dev"), + ) + .addOption( + new Option( + "--dev", + "use your dev branch, even when .env points at a session branch", + ), + ) + .option( + "--ttl ", + "TTL of a new session branch, e.g. 30m, 1h or 7d (default 7d)", + ) + .addOption( + new Option( + "--schema ", + "how the schema starts (default: migrate with schema.ts, introspect without)", + ).choices(["migrate", "introspect", "clean", "none"]), + ) + .option( + "-y, --yes", + "confirm --schema clean, and fail instead of prompting on an ambiguous rename", + ) + .addHelpText( + "after", + ` +init asks nothing. It keeps this checkout's session branch when .env points at +one, and otherwise uses your dev branch (dev--). New branches fork +the project's default branch, so they start with a copy of its rows, production +data included. init never changes the default branch, a protected branch or a +branch databricks.yml deploys to, and never edits databricks.yml. + +The project comes from LAKEBASE_PROJECT in .env (projects/ or ), else +from LAKEBASE_ENDPOINT, else from the bundle's postgres_branch variable. When +.env names neither DATABRICKS_CONFIG_PROFILE nor DATABRICKS_HOST, init adds the +one it authenticated with. + +Examples: + $ appkit db init + $ appkit db init --session --ttl 1h --schema clean --yes`, + ) + .action(runDbInit); diff --git a/packages/shared/src/cli/commands/db/output.ts b/packages/shared/src/cli/commands/db/output.ts index b73bbe163..72304f19d 100644 --- a/packages/shared/src/cli/commands/db/output.ts +++ b/packages/shared/src/cli/commands/db/output.ts @@ -3,7 +3,11 @@ import process from "node:process"; import pc from "picocolors"; -import type { IntrospectReport, MigrateReport } from "./tooling-contract"; +import type { + CleanReport, + IntrospectReport, + MigrateReport, +} from "./tooling-contract"; /** Prompts only when a human can answer them. */ export function isInteractive(): boolean { @@ -15,7 +19,7 @@ export function relative(file: string): string { } /** The last segment of a Lakebase resource name: projects/p/branches/dev-a → dev-a. */ -function shortName(name: string): string { +export function shortName(name: string): string { return name.slice(name.lastIndexOf("/") + 1); } @@ -79,6 +83,19 @@ export function printMigrateReport(report: MigrateReport): void { ok(`Generated ${relative(report.typesFile)}`); } +export function printCleanReport(report: CleanReport): void { + ok( + report.schemaName === "public" + ? `Removed your objects in schema public on ${shortName(report.branch)}` + : `Dropped schema ${report.schemaName} on ${shortName(report.branch)}`, + ); + if (report.dependents.length > 0) { + warn( + `Also removed, because they depended on it: ${report.dependents.join(", ")}`, + ); + } +} + export function printIntrospectReport(report: IntrospectReport): void { const count = report.tables.length; ok( diff --git a/packages/shared/src/cli/commands/db/seed.ts b/packages/shared/src/cli/commands/db/seed.ts new file mode 100644 index 000000000..72a7cf65e --- /dev/null +++ b/packages/shared/src/cli/commands/db/seed.ts @@ -0,0 +1,27 @@ +import { Command } from "commander"; + +import { loadDatabaseTooling } from "./load-tooling"; +import { fail, ok, relative, shortName } from "./output"; + +async function runDbSeed(): Promise { + try { + const tooling = await loadDatabaseTooling(); + const report = await tooling.runSeed(tooling.createDbContext()); + ok(`Ran ${relative(report.file)} on ${shortName(report.branch)}`); + } catch (error) { + fail(error); + } +} + +export const dbSeedCommand = new Command("seed") + .description( + "Run config/database/seed.sql in one transaction on your dev or session branch", + ) + .addHelpText( + "after", + ` +Unqualified names resolve to the app schema first (app, or the schemaName in +config/database/schema.ts), then public. The file already runs in one +transaction, so it must not contain BEGIN or COMMIT.`, + ) + .action(runDbSeed); diff --git a/packages/shared/src/cli/commands/db/tooling-contract.ts b/packages/shared/src/cli/commands/db/tooling-contract.ts index cb39f69b2..8298134b0 100644 --- a/packages/shared/src/cli/commands/db/tooling-contract.ts +++ b/packages/shared/src/cli/commands/db/tooling-contract.ts @@ -39,6 +39,55 @@ export interface IntrospectReport { readonly typesFile: string; } +export type BranchMode = "dev" | "session"; +export type SchemaStart = "migrate" | "introspect" | "clean" | "none"; + +export interface InitOptions { + readonly mode?: BranchMode; + readonly schema?: SchemaStart; + /** As `--ttl` takes it: 30m, 1h, 7d. */ + readonly ttl?: string; + readonly interactive?: boolean; + readonly onStep?: (message: string) => void; +} + +export interface CleanReport { + readonly branch: string; + readonly schemaName: string; + readonly dependents: readonly string[]; +} + +export type InitStart = + | { readonly kind: "none" } + | { readonly kind: "migrate"; readonly migrate: MigrateReport } + | { readonly kind: "introspect"; readonly introspect: IntrospectReport } + | { + readonly kind: "clean"; + readonly clean: CleanReport; + readonly migrate: MigrateReport; + }; + +export interface InitReport { + readonly mode: BranchMode; + readonly branch: string; + readonly created: boolean; + readonly sourceBranch?: string; + readonly expireTime: string | null; + readonly endpoint: string; + readonly endpointCreated: boolean; + readonly host: string; + readonly database: string; + readonly user: string; + readonly role: { + readonly name: string; + readonly action: "kept" | "created" | "updated"; + }; + readonly handedOff: string | null; + readonly envFile: string; + readonly installed: string | null; + readonly start: InitStart; +} + /** * The exports of `@databricks/appkit/database-tooling` the CLI calls. `shared` * must not depend on appkit, so the CLI imports the tooling dynamically and @@ -59,4 +108,13 @@ export interface DatabaseTooling { ctx: Ctx, options: { interactive: boolean; schemaName?: string }, ) => Promise; + readonly runInit: (ctx: Ctx, options?: InitOptions) => Promise; + readonly previewClean: (ctx: Ctx) => Promise; + readonly runClean: (ctx: Ctx) => Promise; + readonly deleteSessionBranch: ( + ctx: Ctx, + ) => Promise<{ readonly branch: string }>; + readonly runSeed: ( + ctx: Ctx, + ) => Promise<{ readonly branch: string; readonly file: string }>; }