@batadata/cli 0.1.1 → 0.1.3

package/dist/commands/db.js

@@ -2,8 +2,9 @@ import { execSync, spawn } from "node:child_process";
 import { api, apiError } from "../api.js";
 import { requireToken, loadConfig, isJsonMode } from "../config.js";
 import { colors, log, json, error, spinner, table, heading, info as logInfo } from "../utils/logger.js";
-import { prompt, confirm } from "../utils/prompts.js";
+import { prompt, confirmDestructive } from "../utils/prompts.js";
 import { openBrowser } from "../utils/open.js";
+import { emitError, isRetryable } from "../utils/errors.js";
 async function getConnectionInfo(projectId, token) {
     // reveal=true so the returned string is actually usable (the owner is asking).
     const res = await api.get(`/v1/connection-info/${projectId}`, token, { reveal: "true" });
@@ -29,20 +30,50 @@ function formatDate(iso) {
     const d = new Date(iso);
     return d.toLocaleDateString("en-US", { month: "short", day: "numeric", year: "numeric" });
 }
+/**
+ * Human-readable STATUS cell for a branch. Prefers the live compute lifecycle
+ * (computeStatus) over the static branch row status, and appends a green check
+ * when the compute is ready to serve a query without a cold start.
+ */
+function branchStatusLabel(b) {
+    const base = b.computeStatus || b.status || "-";
+    if (b.ready === true)
+        return `${base} ${colors.green("✓")}`;
+    return base;
+}
+/**
+ * Resolve a project's primary branch (the one to run a query against). Lists
+ * branches the same way `db branches` does (GET /v1/projects/:id) and returns
+ * the primary, falling back to the first branch. Returns null if the project
+ * has no branches or the lookup failed.
+ */
+async function getPrimaryBranch(projectId, token, teamId) {
+    const query = {};
+    if (teamId)
+        query.team_id = teamId;
+    const res = await api.get(`/v1/projects/${projectId}`, token, query);
+    if (!res.ok)
+        return null;
+    const list = res.data.branches ?? [];
+    return list.find((b) => b.is_primary) ?? list[0] ?? null;
+}
 export async function connect() {
+    // psql is an interactive session — there's no headless equivalent. Don't spawn
+    // it in --json or non-TTY contexts; point agents at the headless surfaces.
+    if (isJsonMode() || !process.stdin.isTTY) {
+        emitError("INTERACTIVE_ONLY", "bata db connect opens an interactive psql session and can't run headlessly.", "Use `bata db url` for a connection string or `bata db query <sql>` for headless execution.");
+    }
     const token = requireToken();
     const config = loadConfig();
     const projectId = config.defaultProject;
     if (!projectId) {
-        error("No default project. Run bata projects create or set one with bata projects info <id>.");
-        process.exit(1);
+        emitError("NO_PROJECT", "No default project.", "Run bata projects create or set one with bata projects info <id>.");
     }
     const s = spinner("Fetching connection info");
     const conn = await getConnectionInfo(projectId, token);
     s.stop();
     if (!conn || !conn.connection_uri) {
-        error("Could not fetch connection string for this project.");
-        process.exit(1);
+        emitError("NOT_FOUND", "Could not fetch connection string for this project.", "");
     }
     log();
     logInfo(`Connecting to project ${colors.cyan(projectId)}`);
@@ -70,17 +101,11 @@ export async function url() {
     const config = loadConfig();
     const projectId = config.defaultProject;
     if (!projectId) {
-        error("No default project set.");
-        process.exit(1);
+        emitError("NO_PROJECT", "No default project set.", "Set one with: bata projects info <id>");
     }
     const conn = await getConnectionInfo(projectId, token);
     if (!conn || !conn.connection_uri) {
-        const msg = "Could not fetch connection string.";
-        if (isJsonMode())
-            json({ error: msg });
-        else
-            error(msg);
-        process.exit(1);
+        emitError("NOT_FOUND", "Could not fetch connection string.", "");
     }
     if (isJsonMode()) {
         json({ direct: conn.connection_uri, pooled: conn.pooled_uri ?? null });
@@ -95,12 +120,7 @@ export async function branches() {
     const jsonMode = isJsonMode();
     const projectId = config.defaultProject;
     if (!projectId) {
-        const msg = "No default project set.";
-        if (jsonMode)
-            json({ error: msg });
-        else
-            error(msg);
-        process.exit(1);
+        emitError("NO_PROJECT", "No default project set.", "Set one with: bata projects info <id>");
     }
     const s = jsonMode ? null : spinner("Fetching branches");
     const query = {};
@@ -109,11 +129,10 @@ export async function branches() {
     const res = await api.get(`/v1/projects/${projectId}`, token, query);
     s?.stop();
     if (!res.ok) {
-        if (jsonMode)
-            json({ error: apiError(res, "Failed to fetch branches") });
-        else
-            error(apiError(res, "Failed to fetch project branches"));
-        process.exit(1);
+        emitError(res.status === 401 || res.status === 403 ? "INVALID_KEY"
+            : res.status === 404 ? "NO_PROJECT"
+                : res.status >= 500 || res.status === 0 ? "API_UNAVAILABLE"
+                    : "CLI_ERROR", apiError(res, "Failed to fetch branches"), "");
     }
     const branchList = res.data.branches || [];
     if (jsonMode) {
@@ -125,6 +144,10 @@ export async function branches() {
                 name: b.name,
                 is_primary: b.is_primary,
                 status: b.status,
+                // Compute readiness — what agents poll on between a create/cold start
+                // and a successful query.
+                computeStatus: b.computeStatus ?? null,
+                ready: b.ready ?? null,
                 created_at: b.created_at,
             })),
             count: branchList.length,
@@ -140,25 +163,33 @@ export async function branches() {
     table(["NAME", "PRIMARY", "STATUS", "CREATED"], branchList.map((b) => [
         b.name,
         b.is_primary ? colors.green("yes") : "-",
-        b.status || "-",
+        // Prefer the live compute lifecycle (computeStatus); fall back to the
+        // branch row status. Mark ready with a check so the column is scannable.
+        branchStatusLabel(b),
         formatDate(b.created_at),
     ]));
     log();
 }
 export async function branchCreate(name) {
+    const jsonMode = isJsonMode();
     const token = requireToken();
     const config = loadConfig();
     const projectId = config.defaultProject;
     if (!projectId) {
-        error("No default project set.");
-        process.exit(1);
+        emitError("NO_PROJECT", "No default project set.", "Set one with: bata projects info <id>");
     }
-    const branchName = name || (await prompt("Branch name"));
+    // Don't block on an interactive prompt headlessly.
+    let branchName = name;
     if (!branchName) {
-        error("Branch name is required.");
-        process.exit(1);
+        if (jsonMode || !process.stdin.isTTY) {
+            emitError("MISSING_ARG", "Branch name is required.", "Usage: bata db branch create <name>");
+        }
+        branchName = await prompt("Branch name");
+    }
+    if (!branchName) {
+        emitError("MISSING_ARG", "Branch name is required.", "Usage: bata db branch create <name>");
     }
-    const s = spinner(`Creating branch ${colors.cyan(branchName)}`);
+    const s = jsonMode ? null : spinner(`Creating branch ${colors.cyan(branchName)}`);
     const body = {
         name: branchName,
         project_id: projectId,
@@ -166,54 +197,66 @@ export async function branchCreate(name) {
     if (config.defaultTeam)
         body.team_id = config.defaultTeam;
     const res = await api.post("/v1/branches", body, token);
-    s.stop();
+    s?.stop();
     if (!res.ok) {
-        error(apiError(res, "Failed to create branch"));
-        process.exit(1);
+        emitError(res.status >= 500 || res.status === 0 ? "API_UNAVAILABLE" : "CLI_ERROR", apiError(res, "Failed to create branch"), "");
+    }
+    if (jsonMode) {
+        json({
+            branch: { id: res.data.id, name: res.data.name ?? branchName, project_id: projectId },
+            // The branch row exists immediately, but its compute may still be
+            // provisioning — poll `db branches` for status before connecting.
+            ready: false,
+            poll: "bata db branches --json",
+        });
+        return;
     }
     log();
     log(`  ${colors.green(">")} Branch ${colors.cyan(branchName)} created`);
+    log(`  ${colors.dim("Poll readiness with")} ${colors.cyan("bata db branches")}`);
     log();
 }
 export async function branchDelete(name) {
+    const jsonMode = isJsonMode();
     const token = requireToken();
     const config = loadConfig();
     const projectId = config.defaultProject;
     if (!projectId) {
-        error("No default project set.");
-        process.exit(1);
+        emitError("NO_PROJECT", "No default project set.", "Set one with: bata projects info <id>");
     }
     if (!name) {
-        error("Branch name is required. Usage: bata db branch delete <name>");
-        process.exit(1);
+        emitError("MISSING_ARG", "Branch name is required.", "Usage: bata db branch delete <name>");
     }
-    const ok = await confirm(`Delete branch ${colors.cyan(name)}?`, false);
+    // Skip the prompt headlessly (--yes / --json / no TTY) so agents and CI can
+    // actually delete — and never report success without deleting.
+    const ok = await confirmDestructive(`Delete branch ${colors.cyan(name)}?`);
     if (!ok) {
         log("  Aborted.");
         return;
     }
-    const s = spinner(`Deleting branch ${name}`);
+    const s = jsonMode ? null : spinner(`Deleting branch ${name}`);
     // We need to find the branch ID first
     const query = {};
     if (config.defaultTeam)
         query.team_id = config.defaultTeam;
     const projRes = await api.get(`/v1/projects/${projectId}`, token, query);
     if (!projRes.ok || !projRes.data.branches) {
-        s.stop();
-        error("Failed to fetch branches.");
-        process.exit(1);
+        s?.stop();
+        emitError(projRes.status >= 500 || projRes.status === 0 ? "API_UNAVAILABLE" : "CLI_ERROR", apiError(projRes, "Failed to fetch branches."), "");
     }
     const branch = projRes.data.branches.find((b) => b.name === name);
     if (!branch) {
-        s.stop();
-        error(`Branch "${name}" not found.`);
-        process.exit(1);
+        s?.stop();
+        emitError("BRANCH_NOT_FOUND", `Branch "${name}" not found.`, "List branches with: bata db branches --json");
     }
     const res = await api.del(`/v1/branches/${branch.id}`, token, query);
-    s.stop();
+    s?.stop();
     if (!res.ok) {
-        error("Failed to delete branch.");
-        process.exit(1);
+        emitError(res.status >= 500 || res.status === 0 ? "API_UNAVAILABLE" : "CLI_ERROR", apiError(res, "Failed to delete branch."), "");
+    }
+    if (jsonMode) {
+        json({ branch: { id: branch.id, name, project_id: projectId }, deleted: true });
+        return;
     }
     log();
     log(`  ${colors.green(">")} Branch ${colors.cyan(name)} deleted`);
@@ -231,60 +274,194 @@ export async function studio() {
     log();
     openBrowser(studioUrl);
 }
-export async function query(sql) {
+/**
+ * Pull a `--branch <id>` / `--branch=<id>` flag out of the query args and return
+ * the explicit branch id (if any) plus the remaining args, which join into the
+ * SQL string. Keeps `db query` order-independent (the flag can sit before or
+ * after the SQL) and consistent with how global flags are parsed.
+ */
+function parseBranchFlag(args) {
+    let branchId;
+    const rest = [];
+    for (let i = 0; i < args.length; i++) {
+        const arg = args[i];
+        if (arg === "--branch") {
+            branchId = args[++i];
+        }
+        else if (arg.startsWith("--branch=")) {
+            branchId = arg.slice("--branch=".length);
+        }
+        else {
+            rest.push(arg);
+        }
+    }
+    return { branchId, rest };
+}
+export async function query(args = []) {
     const jsonMode = isJsonMode();
+    const { branchId, rest } = parseBranchFlag(args);
+    const sql = rest.join(" ").trim();
     if (!sql) {
-        const msg = 'SQL query is required. Usage: bata db query "SELECT 1"';
-        if (jsonMode)
-            json({ error: msg });
-        else
-            error(msg);
-        process.exit(1);
+        emitError("MISSING_ARG", "SQL query is required.", 'Usage: bata db query "SELECT 1"');
     }
     const token = requireToken();
     const config = loadConfig();
     const projectId = config.defaultProject;
     if (!projectId) {
-        const msg = "No default project set.";
-        if (jsonMode)
-            json({ error: msg });
-        else
-            error(msg);
-        process.exit(1);
+        emitError("NO_PROJECT", "No default project set.", "Set one with: bata projects info <id>");
     }
     const s = jsonMode ? null : spinner("Running query");
-    const res = await api.post(`/v1/query/${projectId}`, { sql }, token);
+    // Target branch: an explicit --branch <id> wins; otherwise resolve the
+    // project's primary branch. /v1/sql/execute is keyed on branch_id, not
+    // project. (The old code POSTed to /v1/query/:projectId, which 404s.)
+    let branchId_ = branchId;
+    if (!branchId_) {
+        const branch = await getPrimaryBranch(projectId, token, config.defaultTeam);
+        if (!branch) {
+            s?.stop();
+            emitError("BRANCH_NOT_FOUND", "No branch found for this project to run the query against.", "Check the project with: bata db branches --json");
+        }
+        branchId_ = branch.id;
+    }
+    const res = await api.post("/v1/sql/execute", { branch_id: branchId_, query: sql }, token);
     s?.stop();
     if (!res.ok) {
-        if (jsonMode)
-            json({ error: apiError(res, "Query failed") });
-        else
-            error(apiError(res, "Query failed"));
-        process.exit(1);
+        // Cold-start / upstream blips are retryable (exit 6) so an agent backs off
+        // and retries instead of treating it as a hard CLI bug. The backend returns
+        // 503 + { code: "COMPUTE_STARTING" } while a just-created/idle branch's
+        // compute is still waking; connection-refused transients land here too.
+        const body = res.data;
+        if (isRetryable({ status: res.status, code: body?.code, message: body?.error })) {
+            emitError("COMPUTE_STARTING", apiError(res, "Compute is starting"), "compute is starting; retry in a few seconds");
+        }
+        const code = res.status === 401 || res.status === 403 ? "INVALID_KEY" : "CLI_ERROR";
+        emitError(code, apiError(res, "Query failed"), "");
+    }
+    // The endpoint returns 200 even when the SQL itself errored (it executed and
+    // failed). Treat that as a real error, not a silent success.
+    if (res.data.error) {
+        emitError("CLI_ERROR", res.data.error, "");
     }
-    const { columns, rows, rowCount } = res.data;
+    const columns = res.data.columns ?? [];
+    const rows = res.data.rows ?? [];
+    const rowCount = res.data.rowCount ?? rows.length;
     if (jsonMode) {
         json({
-            columns: columns ?? [],
-            rows: rows ?? [],
-            row_count: rowCount ?? rows?.length ?? 0,
+            columns,
+            rows, // column-keyed objects, exactly as the server returned them
+            row_count: rowCount,
+            command: res.data.command ?? null,
         });
         return;
     }
-    if (columns && rows) {
+    if (columns.length && rows.length) {
+        log();
+        table(columns.map((c) => c.toUpperCase()), rows.map((r) => columns.map((col) => fmtCell(r[col]))));
+        log();
+        log(`  ${colors.dim(`${rowCount} row(s)`)}`);
+    }
+    else if (rows.length) {
+        // Rows with no field metadata (rare) — render whatever keys came back.
+        const keys = Object.keys(rows[0]);
         log();
-        table(columns.map((c) => c.toUpperCase()), rows.map((r) => r.map(String)));
+        table(keys.map((k) => k.toUpperCase()), rows.map((r) => keys.map((k) => fmtCell(r[k]))));
         log();
-        log(`  ${colors.dim(`${rowCount ?? rows.length} row(s)`)}`);
+        log(`  ${colors.dim(`${rowCount} row(s)`)}`);
     }
     else {
         log();
-        log(`  ${colors.dim("Query executed successfully.")}`);
+        const cmd = res.data.command ? `${res.data.command} ` : "";
+        log(`  ${colors.dim(`${cmd}OK — ${rowCount} row(s) affected.`)}`);
+    }
+    log();
+}
+/** Render a column-keyed value for the table (null → empty, objects → JSON). */
+function fmtCell(v) {
+    if (v === null || v === undefined)
+        return "";
+    if (typeof v === "object")
+        return JSON.stringify(v);
+    return String(v);
+}
+/** Was a `--help` / `-h` flag passed anywhere in the args? */
+function hasHelpFlag(args) {
+    return args.some((a) => a === "--help" || a === "-h");
+}
+/**
+ * Print usage for a `db` subcommand (or general `db` help if none/unknown) and
+ * return. Pure stdout — NEVER touches the network. `db query --help` must print
+ * this, not run the SQL "--help".
+ */
+function dbHelp(sub) {
+    const usage = (line) => log(`  ${colors.cyan(line)}`);
+    const note = (line) => log(`  ${colors.dim(line)}`);
+    log();
+    switch (sub) {
+        case "query":
+            log(`  ${colors.bold("bata db query")} — run a SQL query against a branch`);
+            log();
+            usage('bata db query "SELECT 1"');
+            usage("bata db query <sql> [--branch <branch_id>] [--json]");
+            log();
+            note("--branch <branch_id>   Target a specific branch (default: the project's primary)");
+            note("--json                 Emit rows as JSON objects + row_count");
+            note("Cold-start note: a just-created/idle branch may answer with exit 6");
+            note("(retryable) while its compute wakes — retry in a few seconds.");
+            break;
+        case "branches":
+            log(`  ${colors.bold("bata db branches")} — list branches and their compute status`);
+            log();
+            usage("bata db branches [--json]");
+            log();
+            note("STATUS shows the live compute state; a ✓ means ready to query.");
+            note("--json includes computeStatus + ready — poll these after create/cold start.");
+            break;
+        case "branch":
+            log(`  ${colors.bold("bata db branch")} — create or delete a branch`);
+            log();
+            usage("bata db branch create <name>");
+            usage("bata db branch delete <name> [--yes]");
+            break;
+        case "url":
+            log(`  ${colors.bold("bata db url")} — print the connection string`);
+            log();
+            usage("bata db url [--json]");
+            break;
+        case "connect":
+            log(`  ${colors.bold("bata db connect")} — open an interactive psql session`);
+            log();
+            usage("bata db connect");
+            note("Interactive only — use `bata db query` / `bata db url` headlessly.");
+            break;
+        case "studio":
+            log(`  ${colors.bold("bata db studio")} — open the table browser in your browser`);
+            log();
+            usage("bata db studio");
+            break;
+        default:
+            log(`  ${colors.bold("bata db")} — database commands`);
+            log();
+            usage("bata db connect          Open psql to your database");
+            usage("bata db url              Print connection string");
+            usage("bata db branches         List branches + compute status");
+            usage("bata db branch create    Create a new branch");
+            usage("bata db branch delete    Delete a branch");
+            usage("bata db studio           Open the table browser");
+            usage("bata db query <sql>      Run a SQL query (--branch <id> to target a branch)");
+            log();
+            note("Add --help to any subcommand for details (e.g. bata db query --help).");
     }
     log();
 }
 export async function handleDb(args) {
     const sub = args[0];
+    // A --help / -h anywhere prints usage for the subcommand and exits 0 WITHOUT
+    // any network call. This must run before dispatch so `db query --help` never
+    // treats "--help" as SQL and hits /v1/sql/execute.
+    if (hasHelpFlag(args)) {
+        dbHelp(sub);
+        return;
+    }
     switch (sub) {
         case "connect":
             return connect();
@@ -298,18 +475,13 @@ export async function handleDb(args) {
                 return branchCreate(args[2]);
             if (action === "delete")
                 return branchDelete(args[2]);
-            error(`Unknown: db branch ${action || ""}`);
-            log(`  ${colors.dim("Available:")} create, delete`);
-            process.exit(1);
-            break;
+            emitError("INVALID_FLAG", `Unknown: db branch ${action || ""}`, "Available: create, delete");
         }
         case "studio":
             return studio();
         case "query":
-            return query(args.slice(1).join(" "));
+            return query(args.slice(1));
         default:
-            error(`Unknown subcommand: db ${sub || ""}`);
-            log(`  ${colors.dim("Available:")} connect, url, branches, branch, studio, query`);
-            process.exit(1);
+            emitError("INVALID_FLAG", `Unknown subcommand: db ${sub || ""}`, "Available: connect, url, branches, branch, studio, query");
     }
 }

package/dist/commands/db.test.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/commands/db.test.js

@@ -0,0 +1,173 @@
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+const { get, post } = vi.hoisted(() => ({
+    get: vi.fn(async () => ({
+        ok: true,
+        status: 200,
+        data: { id: "proj_1", name: "demo", branches: [{ id: "br_primary", name: "main", is_primary: true, status: "ready" }] },
+    })),
+    post: vi.fn(async () => ({ ok: true, status: 200, data: { columns: ["n"], rows: [{ n: "1" }], rowCount: 1 } })),
+}));
+vi.mock("../api.js", () => ({
+    api: { get, post, del: vi.fn() },
+    resolveTeamId: vi.fn(async () => "team_test"),
+    // Mirror the real apiError so retryable mapping (which reads the server code)
+    // gets a realistic message; we don't need its exact format here.
+    apiError: (res, fallback) => {
+        const msg = res?.data?.error ?? fallback;
+        const code = res?.data?.code ? ` (${res.data.code})` : "";
+        return `${msg}${code}`;
+    },
+    asList: (data) => (Array.isArray(data) ? data : []),
+}));
+import { query, handleDb } from "./db.js";
+import { setRuntime } from "../config.js";
+let tmpHome;
+let logSpy;
+// process.exit must NOT actually exit the test runner. Throw a tagged error so
+// the handler unwinds at the exit point; tests catch it and read the code.
+class ExitError extends Error {
+    code;
+    constructor(code) {
+        super(`exit ${code}`);
+        this.code = code;
+    }
+}
+beforeEach(() => {
+    tmpHome = fs.mkdtempSync(path.join(os.tmpdir(), "bata-cli-db-test-"));
+    vi.stubEnv("HOME", tmpHome);
+    vi.stubEnv("USERPROFILE", tmpHome);
+    // A real credential + a default project so query() reaches the network seam.
+    fs.writeFileSync(path.join(tmpHome, ".batarc"), JSON.stringify({ defaultProject: "proj_1" }));
+    setRuntime({ apiKey: "test-token", apiUrl: "https://api.test.local", json: true, yes: false });
+    get.mockClear();
+    post.mockClear();
+    vi.spyOn(process, "exit").mockImplementation(((code) => {
+        throw new ExitError(code ?? 0);
+    }));
+    // Silence the JSON/decorative output so test logs stay clean.
+    logSpy = vi.spyOn(console, "log").mockImplementation(() => { });
+    vi.spyOn(console, "error").mockImplementation(() => { });
+});
+afterEach(() => {
+    vi.unstubAllEnvs();
+    vi.restoreAllMocks();
+    setRuntime({ json: false, yes: false });
+    fs.rmSync(tmpHome, { recursive: true, force: true });
+});
+/** Run a handler that's expected to call process.exit; return the exit code. */
+async function exitCodeOf(run) {
+    try {
+        await run();
+    }
+    catch (e) {
+        if (e instanceof ExitError)
+            return e.code;
+        throw e;
+    }
+    return 0; // handler returned without exiting (success path)
+}
+describe("db query — --help short-circuits before any network call", () => {
+    it("`db query --help` prints help and makes NO api call (exit 0)", async () => {
+        const code = await exitCodeOf(() => handleDb(["query", "--help"]));
+        expect(post).not.toHaveBeenCalled();
+        expect(get).not.toHaveBeenCalled();
+        expect(code).toBe(0);
+        // It actually printed usage, not a row result.
+        const printed = logSpy.mock.calls.map((c) => String(c[0])).join("\n");
+        expect(printed).toContain("bata db query");
+        expect(printed).toContain("--branch");
+    });
+    it("`db query -h` (short flag) also short-circuits with no api call", async () => {
+        const code = await exitCodeOf(() => handleDb(["query", "-h"]));
+        expect(post).not.toHaveBeenCalled();
+        expect(get).not.toHaveBeenCalled();
+        expect(code).toBe(0);
+    });
+    it("`db --help` prints general db help with no api call", async () => {
+        const code = await exitCodeOf(() => handleDb(["--help"]));
+        expect(post).not.toHaveBeenCalled();
+        expect(get).not.toHaveBeenCalled();
+        expect(code).toBe(0);
+    });
+});
+describe("db query — branch targeting", () => {
+    it("--branch <id> sets branch_id in the request body and skips primary lookup", async () => {
+        await exitCodeOf(() => query(["SELECT 1", "--branch", "br_explicit"]));
+        expect(post).toHaveBeenCalledTimes(1);
+        expect(post).toHaveBeenCalledWith("/v1/sql/execute", expect.objectContaining({ branch_id: "br_explicit", query: "SELECT 1" }), "test-token");
+        // Explicit branch → no need to resolve the primary.
+        expect(get).not.toHaveBeenCalled();
+    });
+    it("--branch=<id> (equals form) is also honored", async () => {
+        await exitCodeOf(() => query(["--branch=br_eq", "SELECT 2"]));
+        expect(post).toHaveBeenCalledWith("/v1/sql/execute", expect.objectContaining({ branch_id: "br_eq", query: "SELECT 2" }), "test-token");
+    });
+    it("without --branch, resolves the project's primary branch (default behavior)", async () => {
+        await exitCodeOf(() => query(["SELECT 1"]));
+        expect(get).toHaveBeenCalledTimes(1); // resolved primary
+        expect(post).toHaveBeenCalledWith("/v1/sql/execute", expect.objectContaining({ branch_id: "br_primary", query: "SELECT 1" }), "test-token");
+    });
+});
+describe("db query — retryable cold-start mapping (exit 6)", () => {
+    it("503 + { code: COMPUTE_STARTING } maps to exit 6 (retryable), not 1", async () => {
+        post.mockResolvedValueOnce({
+            ok: false,
+            status: 503,
+            data: { error: "compute is starting", code: "COMPUTE_STARTING" },
+        });
+        const code = await exitCodeOf(() => query(["SELECT 1", "--branch", "br_cold"]));
+        expect(code).toBe(6);
+    });
+    it("REGRESSION GUARD: a non-5xx status carrying { code: COMPUTE_STARTING } still maps to exit 6", async () => {
+        // The pre-fix code only looked at the HTTP status (>=500 || 0), so a
+        // COMPUTE_STARTING surfaced on, say, a 425 would have been CLI_ERROR (exit
+        // 1). This is the exact guard for the body-code inspection.
+        post.mockResolvedValueOnce({
+            ok: false,
+            status: 425,
+            data: { error: "compute is starting", code: "COMPUTE_STARTING" },
+        });
+        const code = await exitCodeOf(() => query(["SELECT 1", "--branch", "br_cold"]));
+        expect(code).toBe(6);
+    });
+    it("a connection-refused error body maps to exit 6", async () => {
+        post.mockResolvedValueOnce({
+            ok: false,
+            status: 0,
+            data: { error: "connect ECONNREFUSED 127.0.0.1:5432" },
+        });
+        const code = await exitCodeOf(() => query(["SELECT 1", "--branch", "br_cold"]));
+        expect(code).toBe(6);
+    });
+    it("a generic 5xx maps to exit 6", async () => {
+        post.mockResolvedValueOnce({ ok: false, status: 502, data: { error: "bad gateway" } });
+        const code = await exitCodeOf(() => query(["SELECT 1", "--branch", "br_cold"]));
+        expect(code).toBe(6);
+    });
+    it("a genuine SQL syntax error (200 with error field) stays exit 1, NOT retryable", async () => {
+        post.mockResolvedValueOnce({
+            ok: true,
+            status: 200,
+            data: { error: 'syntax error at or near "SELCT"' },
+        });
+        const code = await exitCodeOf(() => query(["SELCT 1", "--branch", "br_ok"]));
+        expect(code).toBe(1);
+    });
+    it("a 400 bad-input SQL error stays exit 1, NOT retryable", async () => {
+        post.mockResolvedValueOnce({
+            ok: false,
+            status: 400,
+            data: { error: "relation does not exist" },
+        });
+        const code = await exitCodeOf(() => query(["SELECT * FROM nope", "--branch", "br_ok"]));
+        expect(code).toBe(1);
+    });
+    it("a 401 stays auth (exit 4), not retryable", async () => {
+        post.mockResolvedValueOnce({ ok: false, status: 401, data: { error: "invalid key" } });
+        const code = await exitCodeOf(() => query(["SELECT 1", "--branch", "br_ok"]));
+        expect(code).toBe(4);
+    });
+});