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 }>; }