@nifrajs/mcp-db 4.0.2 → 5.1.0

package/src/internal/shared.ts

@@ -6,8 +6,13 @@
 
 /** The structural slice of `bun:sqlite`'s `Database` this package needs. */
 export interface SqliteDatabaseLike {
-  /** Prepare a statement; `all` runs it and returns row objects. */
-  prepare(sql: string): { all(...params: unknown[]): unknown[] }
+  /** Prepare a statement; `all` runs it and returns row objects. `columnNames` and `iterate`, which
+   * bun:sqlite has, let `run_query` bound cell bytes and stop reading rows once the budget is spent. */
+  prepare(sql: string): {
+    all(...params: unknown[]): unknown[]
+    readonly columnNames?: readonly string[]
+    iterate?(): Iterable<unknown>
+  }
   /** Execute a statement for its side effect (used only for `PRAGMA query_only`). */
   run(sql: string): unknown
   /** Optional path this database was opened from; a file-backed one can be reopened read-only in a
@@ -210,7 +215,7 @@ const isKeyword = (token: SqlToken | undefined, word: string): boolean =>
   token?.bare === true && token.value.toLowerCase() === word
 
 /**
- * Every relation the statement names after `FROM`/`JOIN`, lowercased, with CTE names excluded.
+ * Every relation the statement names after `FROM`/`JOIN`/`IN`, lowercased, with CTE names excluded.
  * A schema qualifier resolves to its table (`main.habits` -> `habits`); a subquery contributes
  * nothing of its own because its inner `FROM` is tokenized alongside everything else.
  */
@@ -228,18 +233,21 @@ function relationNames(sql: string): string[] {
   const names: string[] = []
   const relationAt = (
     index: number,
-  ): { readonly name: string; readonly next: number } | undefined => {
+  ): { readonly name: string; readonly qualified: boolean; readonly next: number } | undefined => {
     const first = tokens[index]
     if (first === undefined || first.kind !== "word") return undefined // `FROM (subquery)` - nothing to name here
     // A qualified reference names the table second; an unqualified one names it first.
     const qualified = tokens[index + 1]?.value === "." && tokens[index + 2]?.kind === "word"
     const table = qualified ? (tokens[index + 2] as SqlToken).value : first.value
-    return { name: table.toLowerCase(), next: qualified ? index + 3 : index + 1 }
+    return { name: table.toLowerCase(), qualified, next: qualified ? index + 3 : index + 1 }
   }
   const addRelation = (index: number): number | undefined => {
     const relation = relationAt(index)
     if (relation === undefined) return undefined
-    if (relation.name !== "" && !ctes.has(relation.name)) names.push(relation.name)
+    // A CTE is never schema-qualified: `main.users` is the table even beside a CTE named `users`.
+    if (relation.name !== "" && (relation.qualified || !ctes.has(relation.name))) {
+      names.push(relation.name)
+    }
     return relation.next
   }
   // Comma joins are part of the FROM table-source grammar but do not have a JOIN keyword. Walk the
@@ -260,6 +268,11 @@ function relationNames(sql: string): string[] {
     "returning",
   ])
   for (let i = 0; i < tokens.length; i++) {
+    // `x IN users` and `x IN json_each(...)` read a relation without a FROM.
+    if (isKeyword(tokens[i], "in")) {
+      addRelation(i + 1)
+      continue
+    }
     const from = isKeyword(tokens[i], "from")
     const join = isKeyword(tokens[i], "join")
     if (!from && !join) continue
@@ -437,7 +450,7 @@ export type SqliteGate =
 
 /**
  * Gate one statement before SQLite sees it: non-empty, a single statement, SELECT/WITH only, and every
- * relation it names after `FROM`/`JOIN` exposed. Returns the statement without its terminator, ready
+ * relation it names after `FROM`/`JOIN`/`IN` exposed. Returns the statement without its terminator, ready
  * to wrap. Relation names are checked here as well as in the plan because SQLite's plan names an
  * alias as the scan target (`users AS habits` plans as `SCAN habits`).
  */
@@ -504,12 +517,79 @@ export function unexposedPlanRelation(
   return undefined
 }
 
+// The opcodes that open a b-tree cursor.
+const BTREE_OPEN_OPCODES = new Set(["OpenRead", "OpenWrite", "ReopenIdx"])
+
+const field = (row: unknown, key: string): unknown =>
+  typeof row === "object" && row !== null ? Reflect.get(row, key) : undefined
+
+/**
+ * The first table a compiled statement opens that `exposed` rejects, or `undefined` when every b-tree
+ * it opens belongs to an exposed table (an index counts as its table). `program` is the statement's
+ * `EXPLAIN` rows and `catalog` is `SELECT tbl_name, rootpage FROM sqlite_master WHERE rootpage > 0`.
+ * The program names what SQLite will read by root page, so no alias, CTE name or quoting changes it.
+ * A virtual table or table-valued function (`VOpen`) names no table at all and is never exposed.
+ */
+export function unexposedProgramRelation(
+  program: readonly unknown[],
+  catalog: readonly unknown[],
+  exposed: (relation: string) => boolean,
+): string | undefined {
+  const tables = new Map<number, string>([[1, "sqlite_master"]])
+  for (const entry of catalog) {
+    const table = field(entry, "tbl_name")
+    if (typeof table === "string") tables.set(Number(field(entry, "rootpage")), table)
+  }
+  for (const step of program) {
+    const opcode = field(step, "opcode")
+    if (opcode === "VOpen") return "<virtual table>"
+    if (typeof opcode !== "string" || !BTREE_OPEN_OPCODES.has(opcode)) continue
+    // P3 is the schema: 0 is main; temp (1) and attached databases (2+) are not in this catalog.
+    const schema = Number(field(step, "p3"))
+    if (schema !== 0) return `<schema ${schema}>`
+    const page = Number(field(step, "p2"))
+    const table = tables.get(page)
+    if (table === undefined) return `<root page ${page}>`
+    if (!exposed(table.toLowerCase())) return table
+  }
+  return undefined
+}
+
 /** Wrap a gated query so SQLite materializes at most one row beyond `maxRows`. */
 export function boundedSqliteQuery(query: string, maxRows: number): string {
   // nifra-expect sql-dynamic: wraps a query already through gateSqliteStatement in a bounding subselect; it runs on a query-only connection
   return `SELECT * FROM (${query}) AS "__nifra_result" LIMIT ${maxRows + 1}`
 }
 
+/** The column {@link cellBoundedSqliteQuery} adds to carry each row's cell bytes. */
+export const ROW_BYTES_COLUMN = "__nifra_row_bytes"
+
+/**
+ * {@link boundedSqliteQuery}, plus two bounds the driver cannot apply itself: a text or blob cell
+ * longer than `maxCellBytes` arrives as `"<n bytes>"` instead of being copied out, and each row
+ * carries its cells' byte total in {@link ROW_BYTES_COLUMN} so a reader can stop once the result
+ * budget is spent. `columns` are the result column names of the bounded query.
+ */
+export function cellBoundedSqliteQuery(
+  query: string,
+  columns: readonly string[],
+  maxRows: number,
+  maxCellBytes: number,
+): string {
+  const selected: string[] = []
+  const sizes: string[] = []
+  for (const name of columns) {
+    const column = `"${name.replaceAll('"', '""')}"`
+    const bytes = `length(CAST(${column} AS BLOB))`
+    selected.push(
+      `CASE WHEN typeof(${column}) IN ('text', 'blob') AND ${bytes} > ${maxCellBytes} THEN '<' || ${bytes} || ' bytes>' ELSE ${column} END AS ${column}`,
+    )
+    sizes.push(`min(coalesce(${bytes}, 0), ${maxCellBytes})`)
+  }
+  // nifra-expect sql-dynamic: wraps a query already through gateSqliteStatement; column names are quoted identifiers and the bounds are validated integers
+  return `SELECT ${selected.join(", ")}, ${sizes.join(" + ")} AS "${ROW_BYTES_COLUMN}" FROM (${query}) AS "__nifra_result" LIMIT ${maxRows + 1}`
+}
+
 /** Wrap a gated query to count its rows (re-executes it). */
 export function countSqliteQuery(query: string): string {
   // nifra-expect sql-dynamic: wraps a query already through gateSqliteStatement in a counting subselect; it runs on a query-only connection

package/src/internal/sqlite-session.ts

@@ -4,7 +4,7 @@
  * tool runs each query in its own subprocess instead, which bounds the work by killing the process.
  */
 
-import type { SqliteDatabaseLike } from "./shared.ts"
+import { ROW_BYTES_COLUMN, type SqliteDatabaseLike } from "./shared.ts"
 
 export class QueryTimeoutError extends Error {
   constructor() {
@@ -13,9 +13,19 @@ export class QueryTimeoutError extends Error {
   }
 }
 
-/** One execution lane for `run_query`: `run` resolves rows, or rejects once the deadline passes. */
+/** Rows read until their bytes passed the budget; `cut` says reading stopped early. */
+export interface BoundedRows {
+  readonly rows: unknown[]
+  readonly cut: boolean
+}
+
+/** One execution lane for `run_query`: each call resolves, or rejects once the deadline passes. */
 export interface QuerySession {
   run(sql: string, deadline: number): Promise<unknown[]>
+  /** The result column names of `sql`, without running it. Empty when the driver cannot say. */
+  columns(sql: string, deadline: number): Promise<readonly string[]>
+  /** Rows of a `cellBoundedSqliteQuery`, read one at a time until their bytes pass `maxBytes`. */
+  read(sql: string, deadline: number, maxBytes: number): Promise<BoundedRows>
   close(): Promise<void>
 }
 
@@ -29,7 +39,7 @@ const SQLITE_WORKER_SOURCE = `
 const { Database } = require("bun:sqlite")
 let db
 self.onmessage = (event) => {
-  const { id, filename, sql, type } = event.data
+  const { id, filename, sql, type, maxBytes } = event.data
   if (type === "close") {
     let response
     try {
@@ -54,7 +64,30 @@ self.onmessage = (event) => {
       db = new Database(filename, { readonly: true })
       db.run("PRAGMA query_only = ON")
     }
-    self.postMessage({ id, ok: true, rows: db.prepare(sql).all() })
+    const statement = db.prepare(sql)
+    if (type === "columns") {
+      const rows = statement.columnNames
+      statement.finalize()
+      self.postMessage({ id, ok: true, rows })
+      return
+    }
+    if (type === "read") {
+      const rows = []
+      let bytes = 0
+      let cut = false
+      for (const { ${ROW_BYTES_COLUMN}: size, ...row } of statement.iterate()) {
+        rows.push(row)
+        bytes += Number(size)
+        if (bytes > maxBytes) {
+          cut = true
+          break
+        }
+      }
+      statement.finalize()
+      self.postMessage({ id, ok: true, rows, cut })
+      return
+    }
+    self.postMessage({ id, ok: true, rows: statement.all() })
   } catch (error) {
     self.postMessage({ id, ok: false, error: String(error) })
   }
@@ -75,10 +108,16 @@ function reopenableFilename(db: SqliteDatabaseLike): string | undefined {
 
 interface PendingQuery {
   readonly deadline: number
-  resolve(rows: unknown[]): void
+  resolve(result: BoundedRows): void
   reject(error: Error): void
 }
 
+interface WorkerQuery {
+  readonly sql: string
+  readonly type?: "columns" | "read"
+  readonly maxBytes?: number
+}
+
 interface PendingClose {
   readonly id: number
   resolve(graceful?: boolean): void
@@ -135,11 +174,12 @@ function workerSession(filename: string): QuerySession {
         id?: number
         ok?: boolean
         rows?: unknown[]
+        cut?: boolean
         error?: string
         closed?: boolean
       }>,
     ) => {
-      const { id, ok, rows, error } = event.data
+      const { id, ok, rows, cut, error } = event.data
       if (typeof id !== "number") return
       if (event.data.closed === true) {
         const closing = pendingClose
@@ -155,7 +195,7 @@ function workerSession(filename: string): QuerySession {
         return
       }
       pending.delete(id)
-      if (ok === true) waiter.resolve(rows ?? [])
+      if (ok === true) waiter.resolve({ rows: rows ?? [], cut: cut === true })
       else waiter.reject(new Error(error ?? "query failed in worker"))
     }
     spawned.onerror = () => discard(new Error("query failed in worker"))
@@ -166,7 +206,7 @@ function workerSession(filename: string): QuerySession {
     return spawned
   }
 
-  const dispatch = (sql: string, deadline: number): Promise<unknown[]> => {
+  const dispatch = (query: WorkerQuery, deadline: number): Promise<BoundedRows> => {
     if (closed) return Promise.reject(new Error("query session is closed"))
     const active = ensureWorker()
     const remaining = deadline - Date.now()
@@ -175,13 +215,13 @@ function workerSession(filename: string): QuerySession {
       return Promise.reject(new QueryTimeoutError())
     }
     const id = ++nextId
-    return new Promise<unknown[]>((resolve, reject) => {
+    return new Promise<BoundedRows>((resolve, reject) => {
       const timer = setTimeout(() => discard(new QueryTimeoutError()), remaining)
       pending.set(id, {
         deadline,
-        resolve: (rows) => {
+        resolve: (result) => {
           clearTimeout(timer)
-          resolve(rows)
+          resolve(result)
         },
         reject: (error) => {
           clearTimeout(timer)
@@ -189,7 +229,7 @@ function workerSession(filename: string): QuerySession {
         },
       })
       try {
-        active.postMessage({ id, filename, sql })
+        active.postMessage({ id, filename, ...query })
       } catch (error) {
         pending.delete(id)
         clearTimeout(timer)
@@ -198,14 +238,21 @@ function workerSession(filename: string): QuerySession {
     })
   }
 
+  // One worker owns the connection, so overlapping calls queue behind each other rather than
+  // interleaving statements on it.
+  const queue = (query: WorkerQuery, deadline: number): Promise<BoundedRows> => {
+    const result = lane.then(() => dispatch(query, deadline))
+    lane = result.catch(() => undefined)
+    return result
+  }
+
   return {
-    run(sql, deadline) {
-      // One worker owns the connection, so overlapping calls queue behind each other rather than
-      // interleaving statements on it.
-      const result = lane.then(() => dispatch(sql, deadline))
-      lane = result.catch(() => undefined)
-      return result
-    },
+    run: async (sql, deadline) => (await queue({ sql }, deadline)).rows,
+    columns: async (sql, deadline) =>
+      (await queue({ sql, type: "columns" }, deadline)).rows.filter(
+        (name): name is string => typeof name === "string",
+      ),
+    read: (sql, deadline, maxBytes) => queue({ sql, type: "read", maxBytes }, deadline),
     close() {
       if (closePromise !== undefined) return closePromise
       closed = true
@@ -257,6 +304,29 @@ function inProcessSession(db: SqliteDatabaseLike): QuerySession {
       if (Date.now() > deadline) throw new QueryTimeoutError()
       return rows
     },
+    async columns(sql, deadline) {
+      if (Date.now() >= deadline) throw new QueryTimeoutError()
+      return db.prepare(sql).columnNames ?? []
+    },
+    async read(sql, deadline, maxBytes) {
+      if (Date.now() >= deadline) throw new QueryTimeoutError()
+      const statement = db.prepare(sql)
+      const rows: unknown[] = []
+      let bytes = 0
+      let cut = false
+      for (const row of statement.iterate?.() ?? statement.all()) {
+        if (typeof row !== "object" || row === null) continue
+        bytes += Number(Reflect.get(row, ROW_BYTES_COLUMN))
+        Reflect.deleteProperty(row, ROW_BYTES_COLUMN)
+        rows.push(row)
+        if (bytes > maxBytes) {
+          cut = true
+          break
+        }
+      }
+      if (Date.now() > deadline) throw new QueryTimeoutError()
+      return { rows, cut }
+    },
     async close() {},
   }
 }

package/src/postgres.ts

@@ -494,7 +494,7 @@ export async function explainPostgres(
   const explained = await gated(client, sql, options, (connection, statement) =>
     explainThroughServer(connection, statement, {
       analyze: options.analyze === true,
-      scope: checkScope(options),
+      scope: checkScope(options, options.redaction?.column),
     }),
   )
   if (isDbRefusal(explained)) return explained
@@ -709,10 +709,10 @@ export interface PostgresRoleSql {
 }
 
 /**
- * Write the SQL for a login role that can read the allowed schemas and nothing else: `pg_read_all_data`
- * on Postgres 14+ when nothing is excluded; otherwise `GRANT SELECT` per schema plus default
- * privileges, and a `REVOKE` for each excluded table and its partitions (`pg_read_all_data` would
- * override a REVOKE, so it is not used then). Nothing is executed.
+ * Write the SQL for a login role that can read the allowed schemas and nothing else: `GRANT SELECT`
+ * per schema plus default privileges, and a `REVOKE` for each excluded table and its partitions.
+ * `pg_read_all_data` is never used: it reaches every schema, present and future, and overrides a
+ * REVOKE. Nothing is executed.
  */
 export async function postgresRoleSql(
   client: PostgresClient,
@@ -756,21 +756,17 @@ export async function postgresRoleSql(
       `ALTER ROLE ${who} SET default_transaction_read_only = on;`,
       `GRANT CONNECT ON DATABASE ${quoteIdent(database)} TO ${who};`,
     ]
-    if (version >= 140000 && excluded.length === 0) {
-      statements.push(`GRANT pg_read_all_data TO ${who};`)
-    } else {
-      for (const schema of schemas) {
-        statements.push(
-          `GRANT USAGE ON SCHEMA ${quoteIdent(schema)} TO ${who};`,
-          `GRANT SELECT ON ALL TABLES IN SCHEMA ${quoteIdent(schema)} TO ${who};`,
-          `ALTER DEFAULT PRIVILEGES IN SCHEMA ${quoteIdent(schema)} GRANT SELECT ON TABLES TO ${who};`,
-        )
-      }
-      for (const row of excluded) {
-        statements.push(
-          `REVOKE SELECT ON TABLE ${quoteIdent(String(row.schema))}.${quoteIdent(String(row.name))} FROM ${who};`,
-        )
-      }
+    for (const schema of schemas) {
+      statements.push(
+        `GRANT USAGE ON SCHEMA ${quoteIdent(schema)} TO ${who};`,
+        `GRANT SELECT ON ALL TABLES IN SCHEMA ${quoteIdent(schema)} TO ${who};`,
+        `ALTER DEFAULT PRIVILEGES IN SCHEMA ${quoteIdent(schema)} GRANT SELECT ON TABLES TO ${who};`,
+      )
+    }
+    for (const row of excluded) {
+      statements.push(
+        `REVOKE SELECT ON TABLE ${quoteIdent(String(row.schema))}.${quoteIdent(String(row.name))} FROM ${who};`,
+      )
     }
     return { role, database, serverVersionNum: version, statements }
   })