postgresai 0.16.0-dev.2 → 0.16.0-dev.3

package/test/dblab.test.ts

@@ -0,0 +1,408 @@
+import { describe, test, expect, mock, afterEach } from "bun:test";
+import {
+  resolveDblabInstanceId,
+  isNumericProjectRef,
+  createClone,
+  listClones,
+  getClone,
+  resetClone,
+  destroyClone,
+  listBranches,
+  createBranch,
+  deleteBranch,
+  branchLog,
+  listSnapshots,
+  createSnapshot,
+  destroySnapshot,
+} from "../lib/dblab";
+import { handleToolCall } from "../lib/mcp-server";
+
+const originalFetch = globalThis.fetch;
+const API = "https://api.example.com";
+
+interface Captured {
+  url: string;
+  options: RequestInit;
+}
+
+/** Stub fetch, returning `body` for every call, capturing the last request. */
+function stubFetch(body: unknown, status = 200): { get: () => Captured | null } {
+  let captured: Captured | null = null;
+  globalThis.fetch = mock((url: string, options: RequestInit) => {
+    captured = { url, options };
+    return Promise.resolve(
+      new Response(typeof body === "string" ? body : JSON.stringify(body), {
+        status,
+        headers: { "Content-Type": "application/json" },
+      })
+    );
+  }) as unknown as typeof fetch;
+  return { get: () => captured };
+}
+
+/** Route fetch by URL: `/dblab_instances` → instances, `/rpc/dblab_api_call` → reply. */
+function routeFetch(instances: unknown[], reply: unknown): { calls: Captured[] } {
+  const calls: Captured[] = [];
+  globalThis.fetch = mock((url: string, options: RequestInit) => {
+    calls.push({ url, options });
+    if (String(url).includes("/dblab_instances")) {
+      return Promise.resolve(new Response(JSON.stringify(instances), { status: 200 }));
+    }
+    return Promise.resolve(new Response(JSON.stringify(reply), { status: 200 }));
+  }) as unknown as typeof fetch;
+  return { calls };
+}
+
+function bodyOf(c: Captured | null): any {
+  return JSON.parse((c!.options.body as string) ?? "{}");
+}
+
+afterEach(() => {
+  globalThis.fetch = originalFetch;
+});
+
+// ---------------------------------------------------------------------------
+// resolveDblabInstanceId
+// ---------------------------------------------------------------------------
+
+describe("resolveDblabInstanceId", () => {
+  const INSTANCES = [
+    { id: 7, project_id: 12, project_name: "Main DB", project_alias: "main-db", project_label_or_name: "Main DB" },
+    { id: 9, project_id: 34, project_name: "Analytics", project_alias: "analytics", project_label_or_name: "Analytics" },
+  ];
+
+  test("throws when apiKey is missing", async () => {
+    await expect(
+      resolveDblabInstanceId({ apiKey: "", apiBaseUrl: API, project: "12" })
+    ).rejects.toThrow("API key is required");
+  });
+
+  test("throws when project is missing", async () => {
+    await expect(
+      resolveDblabInstanceId({ apiKey: "k", apiBaseUrl: API, project: "  " })
+    ).rejects.toThrow("project is required");
+  });
+
+  test("isNumericProjectRef distinguishes id from alias", () => {
+    expect(isNumericProjectRef("12")).toBe(true);
+    expect(isNumericProjectRef(" 34 ")).toBe(true);
+    expect(isNumericProjectRef("main-db")).toBe(false);
+  });
+
+  test("resolves a numeric project id to its instance id (string)", async () => {
+    const cap = stubFetch(INSTANCES);
+    const id = await resolveDblabInstanceId({ apiKey: "k", apiBaseUrl: API, project: "34" });
+    expect(id).toBe("9");
+    const c = cap.get()!;
+    expect(c.url).toContain(`${API}/dblab_instances`);
+    expect(c.url).toContain("select=");
+    expect(c.options.method).toBe("GET");
+    expect((c.options.headers as Record<string, string>)["access-token"]).toBe("k");
+  });
+
+  test("resolves an alias (case-insensitive) to its instance id", async () => {
+    stubFetch(INSTANCES);
+    const id = await resolveDblabInstanceId({ apiKey: "k", apiBaseUrl: API, project: "Main-DB" });
+    expect(id).toBe("7");
+  });
+
+  test("resolves a project name to its instance id", async () => {
+    stubFetch(INSTANCES);
+    const id = await resolveDblabInstanceId({ apiKey: "k", apiBaseUrl: API, project: "analytics" });
+    expect(id).toBe("9");
+  });
+
+  test("adds an org_id filter when orgId is provided", async () => {
+    const cap = stubFetch(INSTANCES);
+    await resolveDblabInstanceId({ apiKey: "k", apiBaseUrl: API, project: "12", orgId: 5 });
+    expect(decodeURIComponent(cap.get()!.url)).toContain("org_id=eq.5");
+  });
+
+  test("throws a helpful error when no instance matches", async () => {
+    stubFetch(INSTANCES);
+    await expect(
+      resolveDblabInstanceId({ apiKey: "k", apiBaseUrl: API, project: "nope" })
+    ).rejects.toThrow(/No DBLab instance found for project 'nope'/);
+  });
+
+  test("surfaces an HTTP error from the listing", async () => {
+    stubFetch('{"message":"boom"}', 500);
+    await expect(
+      resolveDblabInstanceId({ apiKey: "k", apiBaseUrl: API, project: "12" })
+    ).rejects.toThrow(/Failed to resolve project's DBLab instance/);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// dblab_api_call proxy — action/method/data per verb
+// ---------------------------------------------------------------------------
+
+describe("clone verbs → dblab_api_call", () => {
+  const common = { apiKey: "k", apiBaseUrl: API, instanceId: "7" };
+
+  test("createClone posts /clone with a minimal data body", async () => {
+    const cap = stubFetch({ id: "c-1" });
+    await createClone({ ...common });
+    const c = cap.get()!;
+    expect(c.url).toBe(`${API}/rpc/dblab_api_call`);
+    expect(c.options.method).toBe("POST");
+    const b = bodyOf(c);
+    expect(b.instance_id).toBe("7");
+    expect(b.action).toBe("/clone");
+    expect(b.method).toBe("post");
+    expect(b.data).toEqual({ protected: false });
+  });
+
+  test("createClone maps branch / snapshot / db / protected into data", async () => {
+    const cap = stubFetch({ id: "c-1" });
+    await createClone({
+      ...common,
+      cloneId: "c-1",
+      branch: "feature-idx",
+      snapshotId: "s-9",
+      dbUser: "u",
+      dbPassword: "p",
+      isProtected: true,
+    });
+    expect(bodyOf(cap.get()).data).toEqual({
+      protected: true,
+      id: "c-1",
+      branch: "feature-idx",
+      snapshot: { id: "s-9" },
+      db: { username: "u", password: "p" },
+    });
+  });
+
+  test("listClones gets /clones with no data", async () => {
+    const cap = stubFetch([]);
+    await listClones({ ...common });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/clones");
+    expect(b.method).toBe("get");
+    expect(b.data).toBeUndefined();
+  });
+
+  test("getClone gets /clone/<id> (url-encoded)", async () => {
+    const cap = stubFetch({ id: "c/1" });
+    await getClone({ ...common, cloneId: "c/1" });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/clone/c%2F1");
+    expect(b.method).toBe("get");
+  });
+
+  test("resetClone posts /clone/<id>/reset with latest:true when no snapshot", async () => {
+    const cap = stubFetch(true);
+    await resetClone({ ...common, cloneId: "c-1" });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/clone/c-1/reset");
+    expect(b.method).toBe("post");
+    expect(b.data).toEqual({ latest: true });
+  });
+
+  test("resetClone with a snapshot sends snapshotID + latest:false", async () => {
+    const cap = stubFetch(true);
+    await resetClone({ ...common, cloneId: "c-1", snapshotId: "s-3" });
+    expect(bodyOf(cap.get()).data).toEqual({ latest: false, snapshotID: "s-3" });
+  });
+
+  test("destroyClone deletes /clone/<id>", async () => {
+    const cap = stubFetch("");
+    await destroyClone({ ...common, cloneId: "c-1" });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/clone/c-1");
+    expect(b.method).toBe("delete");
+  });
+
+  test("destroyClone requires cloneId", async () => {
+    await expect(destroyClone({ ...common, cloneId: "" })).rejects.toThrow("cloneId is required");
+  });
+});
+
+describe("branch verbs → dblab_api_call", () => {
+  const common = { apiKey: "k", apiBaseUrl: API, instanceId: "7" };
+
+  test("listBranches gets /branches", async () => {
+    const cap = stubFetch([]);
+    await listBranches({ ...common });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/branches");
+    expect(b.method).toBe("get");
+  });
+
+  test("createBranch posts /branch with branchName and optional base/snapshot", async () => {
+    const cap = stubFetch({ name: "feature-idx" });
+    await createBranch({ ...common, branchName: "feature-idx", baseBranch: "main", snapshotId: "s-1" });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/branch");
+    expect(b.method).toBe("post");
+    expect(b.data).toEqual({ branchName: "feature-idx", baseBranch: "main", snapshotID: "s-1" });
+  });
+
+  test("createBranch omits base/snapshot when not given", async () => {
+    const cap = stubFetch({ name: "b" });
+    await createBranch({ ...common, branchName: "b" });
+    expect(bodyOf(cap.get()).data).toEqual({ branchName: "b" });
+  });
+
+  test("deleteBranch deletes /branch/<name> (encoded)", async () => {
+    const cap = stubFetch(true);
+    await deleteBranch({ ...common, branchName: "feature/x" });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/branch/feature%2Fx");
+    expect(b.method).toBe("delete");
+  });
+
+  test("branchLog gets /branch/<name>/log", async () => {
+    const cap = stubFetch([]);
+    await branchLog({ ...common, branchName: "feature-idx" });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/branch/feature-idx/log");
+    expect(b.method).toBe("get");
+  });
+});
+
+describe("snapshot verbs → dblab_api_call", () => {
+  const common = { apiKey: "k", apiBaseUrl: API, instanceId: "7" };
+
+  test("listSnapshots gets /snapshots (no query when unfiltered)", async () => {
+    const cap = stubFetch([]);
+    await listSnapshots({ ...common });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/snapshots");
+    expect(b.method).toBe("get");
+  });
+
+  test("listSnapshots appends branch and dataset query params", async () => {
+    const cap = stubFetch([]);
+    await listSnapshots({ ...common, branchName: "main", dataset: "ds1" });
+    expect(bodyOf(cap.get()).action).toBe("/snapshots?branch=main&dataset=ds1");
+  });
+
+  test("createSnapshot posts /branch/snapshot with cloneID + message", async () => {
+    const cap = stubFetch({ snapshotID: "s-1" });
+    await createSnapshot({ ...common, cloneId: "c-1", message: "before idx" });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/branch/snapshot");
+    expect(b.method).toBe("post");
+    expect(b.data).toEqual({ cloneID: "c-1", message: "before idx" });
+  });
+
+  test("destroySnapshot deletes /snapshot/<id>?force=<bool>", async () => {
+    const cap = stubFetch(true);
+    await destroySnapshot({ ...common, snapshotId: "s-1", force: true });
+    const b = bodyOf(cap.get());
+    expect(b.action).toBe("/snapshot/s-1?force=true");
+    expect(b.method).toBe("delete");
+  });
+
+  test("destroySnapshot defaults force=false", async () => {
+    const cap = stubFetch(true);
+    await destroySnapshot({ ...common, snapshotId: "s-1" });
+    expect(bodyOf(cap.get()).action).toBe("/snapshot/s-1?force=false");
+  });
+});
+
+// ---------------------------------------------------------------------------
+// error / scope-gating paths
+// ---------------------------------------------------------------------------
+
+describe("proxy error paths", () => {
+  const common = { apiKey: "k", apiBaseUrl: API, instanceId: "7" };
+
+  test("a PT403 (scope) surfaces as a formatted 403 error", async () => {
+    stubFetch('{"message":"insufficient scope: joe:admin required"}', 403);
+    await expect(destroyClone({ ...common, cloneId: "c-1" })).rejects.toThrow(/Failed to destroy clone/);
+  });
+
+  test("a 500 surfaces the operation label", async () => {
+    stubFetch('{"message":"boom"}', 500);
+    await expect(listClones({ ...common })).rejects.toThrow(/Failed to list clones/);
+  });
+
+  test("proxy requires an instanceId", async () => {
+    await expect(listClones({ ...common, instanceId: "" })).rejects.toThrow("instanceId is required");
+  });
+});
+
+// ---------------------------------------------------------------------------
+// resolve → proxy end-to-end (project → instance → call)
+// ---------------------------------------------------------------------------
+
+describe("project → instance → dblab_api_call", () => {
+  test("a verb driven off a resolved instance sends both requests", async () => {
+    const instances = [{ id: 7, project_id: 12, project_alias: "main-db" }];
+    const { calls } = routeFetch(instances, { id: "c-1", status: "ready" });
+    const instanceId = await resolveDblabInstanceId({ apiKey: "k", apiBaseUrl: API, project: "main-db" });
+    await createClone({ apiKey: "k", apiBaseUrl: API, instanceId });
+    expect(calls.length).toBe(2);
+    expect(calls[0].url).toContain("/dblab_instances");
+    expect(calls[1].url).toBe(`${API}/rpc/dblab_api_call`);
+    expect(bodyOf(calls[1]).instance_id).toBe("7");
+  });
+});
+
+// ---------------------------------------------------------------------------
+// MCP tool parity — each tool resolves the instance then proxies dblab_api_call
+// ---------------------------------------------------------------------------
+
+describe("MCP DBLab tools (handleToolCall)", () => {
+  const rootOpts = { apiKey: "k", apiBaseUrl: API };
+
+  test("clone_create resolves project then proxies /clone POST", async () => {
+    const instances = [{ id: 7, project_id: 12, project_alias: "main-db" }];
+    const { calls } = routeFetch(instances, { ok: true });
+    const res = await handleToolCall(
+      { params: { name: "clone_create", arguments: { project: "main-db" } } },
+      rootOpts
+    );
+    expect(res.isError).toBeFalsy();
+    expect(calls[0].url).toContain("/dblab_instances");
+    const b = bodyOf(calls[1]);
+    expect(b.instance_id).toBe("7");
+    expect(b.action).toBe("/clone");
+    expect(b.method).toBe("post");
+  });
+
+  test("clone_destroy proxies /clone/<id> DELETE", async () => {
+    const instances = [{ id: 7, project_id: 12, project_alias: "main-db" }];
+    const { calls } = routeFetch(instances, { ok: true });
+    const res = await handleToolCall(
+      { params: { name: "clone_destroy", arguments: { project: "main-db", clone_id: "c-1" } } },
+      rootOpts
+    );
+    expect(res.isError).toBeFalsy();
+    const b = bodyOf(calls[1]);
+    expect(b.action).toBe("/clone/c-1");
+    expect(b.method).toBe("delete");
+  });
+
+  test("snapshot_create proxies /branch/snapshot POST with cloneID", async () => {
+    const instances = [{ id: 7, project_id: 12, project_alias: "main-db" }];
+    const { calls } = routeFetch(instances, { ok: true });
+    await handleToolCall(
+      { params: { name: "snapshot_create", arguments: { project: "main-db", clone_id: "c-1", message: "m" } } },
+      rootOpts
+    );
+    const b = bodyOf(calls[1]);
+    expect(b.action).toBe("/branch/snapshot");
+    expect(b.data).toEqual({ cloneID: "c-1", message: "m" });
+  });
+
+  test("a DBLab tool without project returns an isError response", async () => {
+    const res = await handleToolCall(
+      { params: { name: "clone_list", arguments: {} } },
+      rootOpts
+    );
+    expect(res.isError).toBe(true);
+    expect(res.content[0].text).toContain("project is required");
+  });
+
+  test("clone_status requires clone_id", async () => {
+    const res = await handleToolCall(
+      { params: { name: "clone_status", arguments: { project: "main-db" } } },
+      rootOpts
+    );
+    expect(res.isError).toBe(true);
+    expect(res.content[0].text).toContain("clone_id is required");
+  });
+});

package/test/init.integration.test.ts

@@ -407,12 +407,14 @@ describe.skipIf(skipTests)("integration: prepare-db", () => {
     }
   }, { timeout: TEST_TIMEOUT });
 
-  // 60s timeout for PostgreSQL startup + multiple SQL queries in slow CI
-  test("explain_generic validates input and prevents SQL injection", async () => {
+  // Security regression for #387: explain_generic was a SECURITY DEFINER helper that
+  // ran caller-supplied SQL through EXPLAIN. The planner folds IMMUTABLE/STABLE functions
+  // at plan time in the definer's (superuser) context, making it an RCE/data-exfil vector.
+  // It has been removed; prepare-db must NOT create it.
+  test("explain_generic is not created by prepare-db (RCE helper removed, #387)", async () => {
     pg = await createTempPostgres();
 
     try {
-      // Run init first
       {
         const r = runCliInit([pg.adminUri, "--password", "pw1", "--skip-optional-permissions"]);
         expect(r.status).toBe(0);
@@ -422,82 +424,10 @@ describe.skipIf(skipTests)("integration: prepare-db", () => {
       await c.connect();
 
       try {
-        // Check PostgreSQL version - generic_plan requires 16+
-        const versionRes = await c.query("show server_version_num");
-        const version = parseInt(versionRes.rows[0].server_version_num, 10);
-
-        if (version < 160000) {
-          // Skip this test on older PostgreSQL versions
-          console.log("Skipping explain_generic tests: requires PostgreSQL 16+");
-          return;
-        }
-
-        // Test 1: Empty query should be rejected
-        await expect(
-          c.query("select postgres_ai.explain_generic('')")
-        ).rejects.toThrow(/query cannot be empty/);
-
-        // Test 2: Null query should be rejected
-        await expect(
-          c.query("select postgres_ai.explain_generic(null)")
-        ).rejects.toThrow(/query cannot be empty/);
-
-        // Test 3: Multiple statements (semicolon in middle) should be rejected
-        await expect(
-          c.query("select postgres_ai.explain_generic('select 1; select 2')")
-        ).rejects.toThrow(/semicolon|multiple statements/i);
-
-        // Test 4: Trailing semicolon should be stripped and work
-        {
-          const res = await c.query("select postgres_ai.explain_generic('select 1;') as result");
-          expect(res.rows[0].result).toBeTruthy();
-          expect(res.rows[0].result).toMatch(/Result/i);
-        }
-
-        // Test 5: Valid query should work
-        {
-          const res = await c.query("select postgres_ai.explain_generic('select $1::int', 'text') as result");
-          expect(res.rows[0].result).toBeTruthy();
-        }
-
-        // Test 6: JSON format should work
-        {
-          const res = await c.query("select postgres_ai.explain_generic('select 1', 'json') as result");
-          const plan = JSON.parse(res.rows[0].result);
-          expect(Array.isArray(plan)).toBe(true);
-          expect(plan[0].Plan).toBeTruthy();
-        }
-
-        // Test 7: Whitespace-only query should be rejected
-        await expect(
-          c.query("select postgres_ai.explain_generic('   ')")
-        ).rejects.toThrow(/query cannot be empty/);
-
-        // Test 8: Semicolon in string literal is rejected (documented limitation)
-        // Note: This is a known limitation - the simple heuristic cannot parse SQL strings
-        await expect(
-          c.query("select postgres_ai.explain_generic('select ''hello;world''')")
-        ).rejects.toThrow(/semicolon/i);
-
-        // Test 9: SQL comments should work (no semicolons)
-        {
-          const res = await c.query("select postgres_ai.explain_generic('select 1 -- comment') as result");
-          expect(res.rows[0].result).toBeTruthy();
-        }
-
-        // Test 10: Escaped quotes should work (no semicolons)
-        {
-          const res = await c.query("select postgres_ai.explain_generic('select ''test''''s value''') as result");
-          expect(res.rows[0].result).toBeTruthy();
-        }
-
-        // Test 11: Case-insensitive format parameter
-        {
-          const res = await c.query("select postgres_ai.explain_generic('select 1', 'JSON') as result");
-          const plan = JSON.parse(res.rows[0].result);
-          expect(Array.isArray(plan)).toBe(true);
-        }
-
+        const res = await c.query(
+          "select count(*)::int as n from pg_proc where proname = 'explain_generic' and pronamespace = (select oid from pg_namespace where nspname = 'postgres_ai')"
+        );
+        expect(res.rows[0].n).toBe(0);
       } finally {
         await c.end();
       }