import { Database } from "bun:sqlite"; import { afterAll, beforeAll, describe, expect, it } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; import { removeWithRetries } from "@oh-my-pi/pi-utils"; import "@oh-my-pi/pi-coding-agent/tools/renderers"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { ReadTool } from "@oh-my-pi/pi-coding-agent/tools/read"; import { listTables, parseSqlitePathCandidates, parseSqliteSelector, renderTable, renderTableList, } from "@oh-my-pi/pi-coding-agent/tools/sqlite-reader"; import { WriteTool } from "@oh-my-pi/pi-coding-agent/tools/write"; type ToolTextResult = { content: Array<{ type: string; text?: string }>; }; type SessionLike = ConstructorParameters[0]; function getText(result: ToolTextResult): string { return result.content .filter(block => block.type === "text") .map(block => block.text ?? "") .join("\n"); } function createSession(cwd: string, overrides: Partial = {}): SessionLike { return { cwd, hasUI: false, enableLsp: false, getSessionFile: () => null, getSessionSpawns: () => "*", settings: Settings.isolated(), ...overrides, } as SessionLike; } /** * Builds the fixture database once in memory and serializes it to bytes. Tests * stamp these bytes onto disk with a single `writeFile` (one fsync) instead of * re-running the table creation and ~13 autocommit inserts (≈14 fsyncs) per * test — the original per-test on-disk build dominated the suite's wall time. */ function buildFixtureBytes(): Uint8Array { const db = new Database(":memory:"); try { db.run(` CREATE TABLE users ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, email TEXT NOT NULL, status TEXT NOT NULL, created INTEGER NOT NULL ); CREATE TABLE slugs ( slug TEXT PRIMARY KEY, title TEXT NOT NULL ); CREATE TABLE notes ( body TEXT NOT NULL ); CREATE TABLE composite ( team_id INTEGER NOT NULL, user_id INTEGER NOT NULL, value TEXT NOT NULL, PRIMARY KEY (team_id, user_id) ); CREATE TABLE wide_rows ( id INTEGER PRIMARY KEY, payload TEXT NOT NULL ); `); const insertUser = db.prepare("INSERT INTO users (name, email, status, created) VALUES (?, ?, ?, ?)"); const insertSlug = db.prepare("INSERT INTO slugs (slug, title) VALUES (?, ?)"); const insertNote = db.prepare("INSERT INTO notes (body) VALUES (?)"); const seed = db.transaction(() => { insertUser.run("Alice", "alice@example.com", "active", 1); insertUser.run("Bob", "bob@example.com", "inactive", 2); insertUser.run("Carol", "carol@example.com", "active", 3); insertUser.run("Dave", "dave@example.com", "inactive", 4); insertUser.run("Eve", "eve@example.com", "active", 5); insertUser.run("Frank", "frank@example.com", "active", 6); insertSlug.run("welcome", "Welcome"); insertSlug.run("about", "About"); insertNote.run("First note"); insertNote.run("Second note"); insertNote.run("Third; note"); db.prepare("INSERT INTO composite (team_id, user_id, value) VALUES (?, ?, ?)").run(1, 2, "pair"); db.prepare("INSERT INTO wide_rows (id, payload) VALUES (?, ?)").run(1, "x".repeat(320)); }); seed(); return db.serialize(); } finally { db.close(); } } function readUserEmail(dbPath: string, id: number): string | null { const db = new Database(dbPath, { readonly: true }); try { const row = db.prepare<{ email: string }, [number]>("SELECT email FROM users WHERE id = ?").get(id); return row?.email ?? null; } finally { db.close(); } } function readUserCount(dbPath: string): number { const db = new Database(dbPath, { readonly: true }); try { return db.prepare<{ count: number }, []>("SELECT COUNT(*) AS count FROM users").get()?.count ?? 0; } finally { db.close(); } } function readUserByEmail(dbPath: string, email: string): { name: string; email: string } | null { const db = new Database(dbPath, { readonly: true }); try { return db .prepare<{ name: string; email: string }, [string]>("SELECT name, email FROM users WHERE email = ?") .get(email); } finally { db.close(); } } describe("SQLite tool support", () => { let tmpDir: string; let sqlitePath: string; let sqliteDbPath: string; let invalidDbPath: string; let fixtureBytes: Uint8Array; let readTool: ReadTool; let writeTool: WriteTool; let originalEditVariant: string | undefined; // The shared fixture is only ever read by most tests; the few tests that // mutate a database stamp their own fresh copy via `stampFreshDb`, so the // shared file stays pristine and can be created exactly once. async function stampFreshDb(name: string): Promise { const dbPath = path.join(tmpDir, name); await fs.writeFile(dbPath, fixtureBytes); return dbPath; } beforeAll(async () => { tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), "sqlite-tool-test-")); sqlitePath = path.join(tmpDir, "app.sqlite"); sqliteDbPath = path.join(tmpDir, "app.db"); invalidDbPath = path.join(tmpDir, "thumbs.db"); originalEditVariant = Bun.env.PI_EDIT_VARIANT; Bun.env.PI_EDIT_VARIANT = "replace"; fixtureBytes = buildFixtureBytes(); await fs.writeFile(sqlitePath, fixtureBytes); await fs.writeFile(sqliteDbPath, fixtureBytes); await fs.writeFile(invalidDbPath, "not sqlite\nstill text\n"); const session = createSession(tmpDir); readTool = new ReadTool(session); writeTool = new WriteTool(session); }); afterAll(async () => { if (originalEditVariant === undefined) { delete Bun.env.PI_EDIT_VARIANT; } else { Bun.env.PI_EDIT_VARIANT = originalEditVariant; } await removeWithRetries(tmpDir); }); it("parses SQLite path candidates at the extension boundary", () => { expect(parseSqlitePathCandidates("data/app.db:users?limit=5")).toEqual([ { sqlitePath: "data/app.db", subPath: "users", queryString: "limit=5", }, ]); expect(parseSqlitePathCandidates("data/app.sqlite")).toEqual([ { sqlitePath: "data/app.sqlite", subPath: "", queryString: "", }, ]); }); it("parses SQLite selectors for row, query, and raw modes", () => { expect(parseSqliteSelector("users:42", "")).toEqual({ kind: "row", table: "users", key: "42" }); expect(parseSqliteSelector("users", "limit=2&offset=3&order=created:desc")).toEqual({ kind: "query", table: "users", limit: 2, offset: 3, order: "created:desc", where: undefined, }); expect(parseSqliteSelector("", "q=SELECT+1")).toEqual({ kind: "raw", sql: "SELECT 1" }); }); it("lists tables for a .sqlite database and excludes sqlite internal tables", async () => { const result = await readTool.execute("sqlite-list", { path: sqlitePath }); const text = getText(result); expect(text).toContain("users (6 rows)"); expect(text).toContain("slugs (2 rows)"); expect(text).toContain("notes (3 rows)"); expect(text).not.toContain("sqlite_sequence"); }); it("lists tables for a .db database when the magic bytes match SQLite", async () => { const result = await readTool.execute("sqlite-db-list", { path: sqliteDbPath }); expect(getText(result)).toContain("users (6 rows)"); }); it("falls through to plain file reading for non-SQLite .db files", async () => { const result = await readTool.execute("sqlite-invalid-db", { path: invalidDbPath }); expect(getText(result)).toContain("not sqlite"); }); it("shows table schema and sample rows", async () => { const result = await readTool.execute("sqlite-schema", { path: `${sqlitePath}:users` }); const text = getText(result); expect(text).toContain("CREATE TABLE users"); expect(text).toContain("Sample rows:"); expect(text).toContain("Alice"); }); it("returns a row by integer primary key", async () => { const result = await readTool.execute("sqlite-row-int", { path: `${sqlitePath}:users:2` }); const text = getText(result); expect(text).toContain("id: 2"); expect(text).toContain("name: Bob"); expect(text).toContain("email: bob@example.com"); }); it("returns a row by text primary key", async () => { const result = await readTool.execute("sqlite-row-text", { path: `${sqlitePath}:slugs:welcome` }); const text = getText(result); expect(text).toContain("slug: welcome"); expect(text).toContain("title: Welcome"); }); it("falls back to ROWID lookups for tables without a declared primary key", async () => { const result = await readTool.execute("sqlite-row-rowid", { path: `${sqlitePath}:notes:1` }); expect(getText(result)).toContain("body: First note"); }); it("errors on composite primary key row lookups", async () => { await expect(readTool.execute("sqlite-row-composite", { path: `${sqlitePath}:composite:1` })).rejects.toThrow( /composite primary key/i, ); }); it("supports pagination and includes a continuation hint", async () => { const result = await readTool.execute("sqlite-page", { path: `${sqlitePath}:users?limit=2&offset=1` }); const text = getText(result); expect(text).toContain("Bob"); expect(text).toContain("Carol"); expect(text).not.toContain("Alice"); expect(text).toContain("append :users?limit=2&offset=3 to the database path to continue"); }); it("supports where and order via the path selector", async () => { const result = await readTool.execute("sqlite-sel-query", { path: `${sqlitePath}:users?where=status='active'&order=created:desc&limit=2`, }); const text = getText(result); expect(text).toContain("Frank"); expect(text).toContain("Eve"); expect(text).not.toContain("Bob"); }); it("rejects where= clauses that try to bypass pagination", () => { expect(() => parseSqliteSelector("users", "where=1=1 LIMIT 1000000 --&limit=2&offset=0")).toThrow( /comments or statement terminators/i, ); expect(() => parseSqliteSelector("users", "where=status='active' LIMIT 1")).toThrow(/LIMIT\/OFFSET\/UNION/i); expect(() => parseSqliteSelector("users", "where=1=1; DROP TABLE users")).toThrow( /comments or statement terminators/i, ); }); it("allows semicolons inside quoted SQLite where string literals", async () => { const result = await readTool.execute("sqlite-sel-semicolon-literal", { path: `${sqlitePath}:notes?where=body LIKE '%;%'&limit=5`, }); const text = getText(result); expect(text).toContain("Third; note"); }); it("rejects SQLite where clauses that try to override pagination control syntax", async () => { await expect( readTool.execute("sqlite-where-pagination-bypass", { path: `${sqlitePath}:users?where=1=1 LIMIT 1000000 --&limit=2&offset=0`, }), ).rejects.toThrow(/comments or statement terminators/i); }); it("rejects mutating raw queries on the readonly connection", async () => { await expect( readTool.execute("sqlite-raw-write", { path: `${sqlitePath}?q=INSERT+INTO+users+(name,email,status,created)+VALUES+('X','x@example.com','active',7)`, }), ).rejects.toThrow(/readonly/i); }); it("caps raw ?q= queries at the row limit and surfaces a LIMIT hint", async () => { // Dedicated database: this is the only test that needs >1000 rows, and it // must not pollute the shared fixture. A single transaction = one commit. const capDbPath = path.join(tmpDir, "rawcap.sqlite"); const db = new Database(capDbPath); try { db.run("CREATE TABLE big (id INTEGER PRIMARY KEY, value TEXT NOT NULL)"); const insert = db.prepare("INSERT INTO big (value) VALUES (?)"); const fill = db.transaction(() => { for (let i = 1; i <= 1200; i++) { insert.run(`val_${i}_end`); } }); fill(); } finally { db.close(); } const result = await readTool.execute("sqlite-raw-row-cap", { path: `${capDbPath}?q=SELECT * FROM big` }); const text = getText(result); expect(text).toContain("val_1000_end"); expect(text).not.toContain("val_1001_end"); expect(text).toContain("Output capped at 1000 rows"); }); it("rejects table names that do not exist instead of interpolating them", async () => { await expect( readTool.execute("sqlite-injection-table", { path: `${sqlitePath}:users;DROP TABLE users;` }), ).rejects.toThrow(/not found/i); }); it("truncates wide rows to the configured table width", () => { const rendered = renderTable(["id", "payload"], [{ id: 1, payload: "x".repeat(320) }], { totalCount: 1, offset: 0, limit: 20, table: "wide_rows", dbPath: sqlitePath, }); for (const line of rendered.split("\n")) { expect(Bun.stringWidth(line)).toBeLessThanOrEqual(120); } }); it("falls back to vertical row blocks when the column count exceeds the horizontal budget (#3107)", () => { const columns = ["_id", ...Array.from({ length: 32 }, (_, i) => `col_${i + 1}`)]; const row: Record = { _id: 7 }; for (let i = 1; i <= 32; i++) row[`col_${i}`] = `value_${i}`; const rendered = renderTable(columns, [row], { totalCount: 1, offset: 0, limit: 20, table: "wide_columns", dbPath: sqlitePath, }); // Horizontal layout would shrink every column to width 1 and chop the // right edge — i.e., the line would look like `| … | … | … | …` (>=2 // ellipses chained by ` | `). Vertical mode renders one `col: value` // per line, so that signature must NOT be present. expect(rendered).not.toMatch(/…(?: \| …){2,}/); // Each declared column must appear with its real value on its own line. expect(rendered).toContain("── Row 1 ──"); expect(rendered).toContain("_id : 7"); expect(rendered).toContain("col_1 : value_1"); expect(rendered).toContain("col_32: value_32"); for (const line of rendered.split("\n")) { expect(Bun.stringWidth(line)).toBeLessThanOrEqual(120); } }); it("inserts rows through the write tool with JSON5 content", async () => { const dbPath = await stampFreshDb("write-insert.sqlite"); await writeTool.execute("sqlite-write-insert", { path: `${dbPath}:users`, content: "{ name: 'Grace', email: 'grace@example.com', status: 'active', created: 7 }", }); expect(readUserByEmail(dbPath, "grace@example.com")).toEqual({ name: "Grace", email: "grace@example.com", }); }); it("updates rows through the write tool by primary key", async () => { const dbPath = await stampFreshDb("write-update.sqlite"); await writeTool.execute("sqlite-write-update", { path: `${dbPath}:users:2`, content: "{ email: 'bob+new@example.com' }", }); expect(readUserEmail(dbPath, 2)).toBe("bob+new@example.com"); }); it("deletes rows through the write tool with empty content", async () => { const dbPath = await stampFreshDb("write-delete.sqlite"); await writeTool.execute("sqlite-write-delete", { path: `${dbPath}:users:2`, content: " ", }); expect(readUserCount(dbPath)).toBe(5); expect(readUserEmail(dbPath, 2)).toBeNull(); }); it("enforces plan mode for SQLite writes", async () => { const planSession = createSession(tmpDir, { getPlanModeState: () => ({ enabled: true, planFilePath: path.join(tmpDir, "plan.md"), }), }); const planWriteTool = new WriteTool(planSession); await expect( planWriteTool.execute("sqlite-plan-mode", { path: `${sqlitePath}:users:1`, content: "{ email: 'blocked@example.com' }", }), ).rejects.toThrow(/Plan mode/i); }); it("rejects writes to non-existent tables", async () => { await expect( writeTool.execute("sqlite-write-missing-table", { path: `${sqlitePath}:missing`, content: "{ value: 1 }", }), ).rejects.toThrow(/not found/i); }); it("rejects writes to non-existent databases", async () => { await expect( writeTool.execute("sqlite-write-missing-db", { path: path.join(tmpDir, "missing.sqlite:users"), content: "{ name: 'Nope' }", }), ).rejects.toThrow(/not found/i); }); it("rejects unknown columns in write content", async () => { await expect( writeTool.execute("sqlite-write-bad-column", { path: `${sqlitePath}:users`, content: "{ bogus: 1 }", }), ).rejects.toThrow(/no column named 'bogus'/i); }); }); describe("SQLite table listing row counts", () => { // These tests exercise `listTables`/`renderTableList` directly against a // `Database` handle, so an in-memory database preserves the row-count // contract with zero disk I/O. `base` is never analyzed (exact / lower-bound // behavior); `analyzed` carries planner estimates. let base: Database; let analyzed: Database; function buildCountsDb(analyze: boolean): Database { const db = new Database(":memory:"); db.run("CREATE TABLE big (id INTEGER PRIMARY KEY, v TEXT NOT NULL)"); db.run("CREATE TABLE small (id INTEGER PRIMARY KEY)"); const bigStmt = db.prepare("INSERT INTO big (v) VALUES (?)"); for (let i = 0; i < 10; i++) bigStmt.run("x"); const smallStmt = db.prepare("INSERT INTO small DEFAULT VALUES"); for (let i = 0; i < 2; i++) smallStmt.run(); if (analyze) db.run("ANALYZE"); return db; } beforeAll(() => { base = buildCountsDb(false); analyzed = buildCountsDb(true); }); afterAll(() => { base.close(); analyzed.close(); }); it("counts small tables exactly", () => { const rendered = renderTableList(listTables(base, { probeCap: 100 })); expect(rendered).toContain("big (10 rows)"); expect(rendered).toContain("small (2 rows)"); }); it("reports the planner estimate for tables larger than the probe cap", () => { // probeCap=5: big (estimate 10) exceeds it and is reported as an estimate // without scanning; small (estimate 2) is counted exactly. const rendered = renderTableList(listTables(analyzed, { probeCap: 5 })); expect(rendered).toContain("big (~10 rows)"); expect(rendered).toContain("small (2 rows)"); }); it("reports a lower bound when an unanalyzed table exceeds the probe cap", () => { // No ANALYZE, so no estimate exists; the bounded probe stops at the cap // and reports a lower bound instead of scanning the whole table. const rendered = renderTableList(listTables(base, { probeCap: 3 })); expect(rendered).toContain("big (3+ rows)"); expect(rendered).toContain("small (2 rows)"); }); });