@nifrajs/mcp-db 5.5.0 → 5.6.0

package/src/index.ts

@@ -31,14 +31,9 @@ import {
   fitToBytes,
   gateSqliteStatement,
   type SqliteDatabaseLike,
-  unexposedPlanRelation,
   unexposedProgramRelation,
 } from "./internal/shared.ts"
-import {
-  openQuerySession,
-  type QuerySession,
-  QueryTimeoutError,
-} from "./internal/sqlite-session.ts"
+import { openQuerySession, QueryTimeoutError } from "./internal/sqlite-session.ts"
 
 export type { SqliteDatabaseLike } from "./internal/shared.ts"
 
@@ -65,7 +60,8 @@ export interface RunQueryOptions {
    */
   readonly maxResultBytes?: number
   /**
-   * Maximum wall-clock time for planning, execution, and optional counting (default 5 seconds).
+   * Maximum wall-clock time for planning, execution, and optional counting (default 5 seconds). It
+   * also bounds `list_tables`' row counts, at the same 5-second default when `runQuery` is omitted.
    * The deadline always bounds the RESPONSE. Whether it also stops the WORK depends on where the
    * query runs: a file-backed database runs it in a reusable read-only worker that is terminated
    * on expiry, while an in-memory database (or a shim with no `filename`) runs it on this thread,
@@ -94,7 +90,7 @@ export interface ServeDatabaseAsMcpOptions {
   /**
    * Set `PRAGMA query_only = ON` on the connection at construction (default true). Disable ONLY
    * when the same connection must also serve writes elsewhere in the app - with it disabled,
-   * `run_query` still gates on SELECT/WITH + plan verification, but the engine-level guarantee is
+   * `run_query` still gates on SELECT/WITH + compiled-statement verification, but the engine-level guarantee is
    * yours to provide (e.g. open a second, read-only connection for the MCP mount - preferred).
    */
   readonly enforceQueryOnly?: boolean
@@ -113,6 +109,10 @@ export class McpDbConfigError extends Error {
   }
 }
 
+// A count reads the whole table and the schema tools answer anyone, so one count per table serves
+// every listing for this long instead of each call scanning again.
+const ROW_COUNT_TTL_MS = 10_000
+
 /** A safe SQL identifier - quoting is not enough for PRAGMA args, so reject instead. */
 const SAFE_IDENTIFIER = /^[A-Za-z_][A-Za-z0-9_]*$/
 
@@ -176,21 +176,44 @@ export function serveDatabaseAsMcp(
   assertResultLimit(queryTimeoutMs, "queryTimeoutMs")
   assertResultLimit(maxConcurrentQueries, "maxConcurrentQueries")
   let activeQueries = 0
-  let session: QuerySession | undefined
+  // One lane for the life of the server, picked once here and shared by the row counts and run_query.
+  const session = openQuerySession(db)
+
+  const countTables = async (): Promise<{ table: string; rows: unknown }[]> => {
+    const deadline = Date.now() + queryTimeoutMs
+    const tables: { table: string; rows: unknown }[] = []
+    for (const table of options.tables) {
+      let rows: unknown = null
+      try {
+        const [count] = await session.run(
+          // nifra-expect sql-dynamic: table name from the operator-configured allowlist, escaped by sqlIdentifier; SQL cannot bind an identifier as a parameter
+          sqlIdentifiers`SELECT count(*) AS n FROM "${table}"`,
+          deadline,
+        )
+        if (typeof count === "object" && count !== null && "n" in count) rows = count.n
+      } catch (error) {
+        if (!(error instanceof QueryTimeoutError)) throw error
+      }
+      tables.push({ table, rows })
+    }
+    return tables
+  }
+  let listing: { readonly at: number; readonly tables: ReturnType<typeof countTables> } | undefined
 
   const listTables = defineMcpTool({
     name: "list_tables",
-    description: "List the tables this server exposes, with row counts.",
-    handler: () => {
-      const rows = options.tables.map((table) => {
-        // nifra-expect sql-dynamic: table name from the operator-configured allowlist, escaped by sqlIdentifier; SQL cannot bind an identifier as a parameter
-        const [count] = db
-          .prepare(sqlIdentifiers`SELECT count(*) AS n FROM "${table}"`)
-          .all() as Array<{
-          n: number
-        }>
-        return { table, rows: count?.n ?? 0 }
-      })
+    description:
+      "List the tables this server exposes, with row counts (refreshed at most every 10 seconds; null when a count runs past the query time limit).",
+    handler: async () => {
+      if (listing === undefined || Date.now() - listing.at >= ROW_COUNT_TTL_MS) {
+        const tables = countTables()
+        const current = { at: Date.now(), tables }
+        listing = current
+        tables.catch(() => {
+          if (listing === current) listing = undefined
+        })
+      }
+      const rows = await listing.tables
       return { text: JSON.stringify(rows), structuredContent: { tables: rows } }
     },
   })
@@ -218,9 +241,6 @@ export function serveDatabaseAsMcp(
   const tools = [listTables, describeTable]
 
   if (options.runQuery !== undefined) {
-    // One lane for the life of the server: picked once here, so the handler below always has it.
-    const querySession = openQuerySession(db)
-    session = querySession
     const runQuery = defineMcpTool({
       name: "run_query",
       description:
@@ -245,35 +265,24 @@ export function serveDatabaseAsMcp(
           if (!gate.ok) return { isError: true, text: gate.text }
 
           const deadline = Date.now() + queryTimeoutMs
-
-          // Verify via the query plan that only allowlisted tables are touched. SQLite names every
-          // scanned/searched relation in EXPLAIN QUERY PLAN detail rows.
           const query = gate.query
-          let planRows: Array<{ detail?: unknown }>
-          try {
-            planRows = (await querySession.run(`EXPLAIN QUERY PLAN ${query}`, deadline)) as Array<{
-              detail?: unknown
-            }>
-          } catch (error) {
-            if (error instanceof QueryTimeoutError) return { isError: true, text: error.message }
-            return { isError: true, text: `query failed to plan: ${String(error)}` }
-          }
           let catalog: unknown[]
           let program: unknown[]
           try {
-            catalog = await querySession.run(
+            catalog = await session.run(
               "SELECT tbl_name, rootpage FROM sqlite_master WHERE rootpage > 0",
               deadline,
             )
-            program = await querySession.run(`EXPLAIN ${query}`, deadline)
+            program = await session.run(`EXPLAIN ${query}`, deadline)
           } catch (error) {
             if (error instanceof QueryTimeoutError) return { isError: true, text: error.message }
             return { isError: true, text: `query failed to plan: ${String(error)}` }
           }
-          // The plan names a scan by its alias; the program names the table by its root page.
-          const unexposed =
-            unexposedPlanRelation(planRows, (relation) => allowlist.has(relation)) ??
-            unexposedProgramRelation(program, catalog, (relation) => allowlist.has(relation))
+          // The program names each table it opens by root page. EXPLAIN QUERY PLAN cannot stand in:
+          // it names a scan by its alias (`FROM habits h` plans as `SCAN h`).
+          const unexposed = unexposedProgramRelation(program, catalog, (relation) =>
+            allowlist.has(relation),
+          )
           if (unexposed !== undefined) {
             return {
               isError: true,
@@ -288,11 +297,11 @@ export function serveDatabaseAsMcp(
             // cannot make the driver materialize every row. Where the driver names the columns,
             // oversized cells and rows past the byte budget are never copied out either.
             const bounded = boundedSqliteQuery(query, maxRows)
-            const columns = await querySession.columns(bounded, deadline)
+            const columns = await session.columns(bounded, deadline)
             if (columns.length === 0) {
-              limitedRows = await querySession.run(bounded, deadline)
+              limitedRows = await session.run(bounded, deadline)
             } else {
-              const read = await querySession.read(
+              const read = await session.read(
                 cellBoundedSqliteQuery(query, columns, maxRows, maxResultBytes),
                 deadline,
                 maxResultBytes,
@@ -309,10 +318,9 @@ export function serveDatabaseAsMcp(
           let total: number | null = limitedRows.length
           if (wasLimited && options.runQuery?.exactTotal === true) {
             try {
-              const countRows = (await querySession.run(
-                countSqliteQuery(query),
-                deadline,
-              )) as Array<{ __nifra_total?: unknown }>
+              const countRows = (await session.run(countSqliteQuery(query), deadline)) as Array<{
+                __nifra_total?: unknown
+              }>
               const count = countRows[0]?.__nifra_total
               if (
                 typeof count !== "number" ||
@@ -357,7 +365,7 @@ export function serveDatabaseAsMcp(
   })
 
   const runQueryOptions = options.runQuery
-  if (session === undefined || runQueryOptions === undefined) return server
+  if (runQueryOptions === undefined) return { ...server, close: () => session.close() }
 
   // `run_query` is authorized at the TRANSPORT boundary, where the inbound Request's credentials
   // are visible - the handler itself performs no auth, so every path to it must pass this gate:

package/src/internal/pg-exec.ts

@@ -42,8 +42,10 @@ export interface PostgresClient {
 
 /** Run caller text as ONE extended-protocol statement with no parameters. */
 export function extended(connection: PgConnection, text: string): PgQuery {
-  const strings: TemplateStringsArray = Object.assign([text], { raw: [text] })
-  return connection(strings)
+  const strings = Object.assign([text], { raw: Object.freeze([text]) })
+  // Shaped like a real template object: Bun.SQL (1.4.4+) refuses one whose `raw` is enumerable.
+  Object.defineProperty(strings, "raw", { enumerable: false })
+  return connection(Object.freeze(strings))
 }
 
 const WRITE_KEYWORD =
@@ -98,6 +100,10 @@ function millis(timeoutMs: number): number {
  * Run `work` inside `BEGIN READ ONLY` on one reserved connection, with the statement, lock and
  * idle-in-transaction timeouts set for this transaction only. Always rolls back, so nothing a call
  * does (a `set_config`, a session setting) outlives it.
+ *
+ * JIT is off too: the catalog reads' recursive subplans cross `jit_above_cost`, and compiling them
+ * took past the timeout on a stock Ubuntu server where running them takes under a millisecond.
+ * PostgreSQL 10 has no `jit` setting, so the version is checked first.
  */
 export async function inReadOnlyTransaction<T>(
   client: PostgresClient,
@@ -114,7 +120,7 @@ export async function inReadOnlyTransaction<T>(
   try {
     // nifra-expect sql-dynamic: fixed statements; the one interpolation is a validated integer of milliseconds
     await connection.unsafe(
-      `BEGIN READ ONLY; SET LOCAL statement_timeout = ${ms}; SET LOCAL lock_timeout = ${ms}; SET LOCAL idle_in_transaction_session_timeout = ${ms}`,
+      `BEGIN READ ONLY; SET LOCAL statement_timeout = ${ms}; SET LOCAL lock_timeout = ${ms}; SET LOCAL idle_in_transaction_session_timeout = ${ms}; SELECT pg_catalog.set_config('jit', 'off', true) WHERE pg_catalog.current_setting('server_version_num')::int >= 110000`,
     )
     return await work(connection)
   } catch (error) {

package/src/internal/pg-lexer.ts

@@ -184,6 +184,10 @@ export const DENIED_FUNCTION_PREFIXES: readonly string[] = [
   "pg_try_advisory",
   "pg_notify",
   "query_to_xml",
+  "ts_stat",
+  "ts_rewrite",
+  "crosstab",
+  "connectby",
   "table_to_xml",
   "cursor_to_xml",
   "schema_to_xml",
@@ -272,6 +276,106 @@ export function lintPostgres(sql: string): PgLintRefusal | undefined {
   return undefined
 }
 
+// Postgres' reserved and type/function-name keywords: none of them can be a bare table name or alias.
+const NOT_A_NAME = new Set(
+  `
+  all analyse analyze and any array as asc asymmetric both case cast check collate column
+  constraint create current_catalog current_date current_role current_time current_timestamp
+  current_user default deferrable desc distinct do else end except false fetch for foreign from
+  grant group having in initially intersect into lateral leading limit localtime localtimestamp
+  not null offset on only or order placing primary references returning select session_user some
+  symmetric system_user table then to trailing true union unique user using variadic when where
+  window with authorization binary collation concurrently cross current_schema freeze full ilike
+  inner is isnull join left like natural notnull outer overlaps right similar tablesample verbose
+  `
+    .trim()
+    .split(/\s+/),
+)
+
+// Words that close a FROM list at their own parenthesis depth.
+const FROM_LIST_ENDS = new Set(
+  `
+  where group having window order limit offset fetch for union intersect except returning select
+  `
+    .trim()
+    .split(/\s+/),
+)
+
+const isPunct = (token: PgToken | undefined, value: string): boolean =>
+  token?.kind === "punct" && token.value === value
+const isKeyword = (token: PgToken | undefined, word: string): boolean =>
+  token?.kind === "word" && token.bare && token.value.toLowerCase() === word
+const isName = (token: PgToken | undefined): token is PgToken =>
+  token?.kind === "word" && !(token.bare && NOT_A_NAME.has(token.value.toLowerCase()))
+
+/** A table reference starting at `start` that carries a column alias list: `[ONLY] name [*] [AS] alias (`. */
+function aliasListAt(
+  tokens: readonly PgToken[],
+  start: number,
+): { readonly table: string; readonly alias: string } | undefined {
+  let at = start
+  const only = isKeyword(tokens[at], "only")
+  if (only) at++
+  const parenthesized = only && isPunct(tokens[at], "(")
+  if (parenthesized) at++
+  let name = tokens[at]
+  if (!isName(name)) return undefined
+  at++
+  for (let part = tokens[at + 1]; isPunct(tokens[at], ".") && isName(part); part = tokens[at + 1]) {
+    name = part
+    at += 2
+  }
+  if (parenthesized) {
+    if (!isPunct(tokens[at], ")")) return undefined
+    at++
+  } else if (isPunct(tokens[at], "(")) return undefined // a function in FROM, whose alias list is safe
+  if (isPunct(tokens[at], "*")) at++
+  if (isKeyword(tokens[at], "as")) at++
+  const alias = tokens[at]
+  if (!isName(alias) || !isPunct(tokens[at + 1], "(")) return undefined
+  return { table: name.value, alias: alias.value }
+}
+
+/**
+ * Every table the statement gives a column alias list, as in `FROM users AS t(a, b)`. A verbose plan
+ * prints such a scan's columns by the alias names (its partitions' scans too), and the list can hand a
+ * real column's name to another one, so no check on the plan's column names can see what is read.
+ * A function's or a subquery's alias list is not reported: the plan prints the base columns under them.
+ */
+export function tableAliasLists(
+  sql: string,
+): readonly { readonly table: string; readonly alias: string }[] {
+  const { tokens } = lexPostgres(sql)
+  const found: { readonly table: string; readonly alias: string }[] = []
+  // Per parenthesis depth: is a FROM list open there.
+  const inFromList: boolean[] = [false]
+  let itemStart = false
+  for (const [index, token] of tokens.entries()) {
+    const starts: boolean = itemStart
+    itemStart = false
+    if (starts) {
+      const listed = aliasListAt(tokens, index)
+      if (listed !== undefined) found.push(listed)
+    }
+    if (isPunct(token, "(")) {
+      // A parenthesized join holds a FROM list of its own; a subquery's SELECT closes it.
+      inFromList.push(starts)
+      itemStart = starts
+    } else if (isPunct(token, ")")) {
+      if (inFromList.length > 1) inFromList.pop()
+    } else if (isPunct(token, ",")) {
+      itemStart = inFromList.at(-1) === true
+    } else if (token.kind === "word" && token.bare) {
+      const word = token.value.toLowerCase()
+      if (word === "from" || word === "join") {
+        inFromList[inFromList.length - 1] = true
+        itemStart = true
+      } else if (FROM_LIST_ENDS.has(word)) inFromList[inFromList.length - 1] = false
+    }
+  }
+  return found
+}
+
 /** The statement without a final `;` (comments after it included), ready to wrap in a subquery. */
 export function withoutTerminator(sql: string): string {
   const last = lexPostgres(sql).tokens.at(-1)

package/src/internal/shared.ts

@@ -110,13 +110,83 @@ export function isDbRefusal(value: unknown): value is DbRefusal {
 // Statement gates
 // ---------------------------------------------------------------------------------------------------
 
-const stripSqlNoise = (sql: string): string =>
-  sql
-    // Strings first so comment markers inside literals don't count.
-    .replace(/'(?:[^']|'')*'/g, "''")
-    .replace(/"(?:[^"]|"")*"/g, '""')
-    .replace(/--[^\n]*/g, " ")
-    .replace(/\/\*[\s\S]*?\*\//g, " ")
+// SQLite's four quotings, by opening delimiter: a 'literal' and the "x", `x` and [x] identifiers. Only
+// the bracket form has no doubled-delimiter escape.
+const QUOTE_CLOSE = new Map([
+  ["'", "'"],
+  ['"', '"'],
+  ["`", "`"],
+  ["[", "]"],
+])
+
+/** The index just past the quoted run opening at `start`, or the text's length when it never closes. */
+function quotedEnd(sql: string, start: number, open: string, close: string): number {
+  for (let i = start + 1; i < sql.length; i++) {
+    if (sql[i] !== close) continue
+    if (open !== "[" && sql[i + 1] === close) i++
+    else return i + 1
+  }
+  return sql.length
+}
+
+/** The index just past a comment opening at `start`, or `undefined` when none opens there. An open
+ * block comment runs to the end of the text, as SQLite reads it. */
+function commentEnd(sql: string, start: number): number | undefined {
+  if (sql[start] === "-" && sql[start + 1] === "-") {
+    const end = sql.indexOf("\n", start + 2)
+    return end === -1 ? sql.length : end + 1
+  }
+  if (sql[start] === "/" && sql[start + 1] === "*") {
+    const end = sql.indexOf("*/", start + 2)
+    return end === -1 ? sql.length : end + 2
+  }
+  return undefined
+}
+
+/** `sql` with each literal and quoted identifier emptied and each comment turned into a space. */
+function stripSqlNoise(sql: string): string {
+  let out = ""
+  let from = 0
+  let i = 0
+  while (i < sql.length) {
+    const char = sql.charAt(i)
+    const comment = commentEnd(sql, i)
+    const close = QUOTE_CLOSE.get(char)
+    if (comment !== undefined) {
+      out += `${sql.slice(from, i)} `
+      i = from = comment
+    } else if (close !== undefined) {
+      out += `${sql.slice(from, i)}${char}${close}`
+      i = from = quotedEnd(sql, i, char, close)
+    } else {
+      i++
+    }
+  }
+  return out + sql.slice(from)
+}
+
+/**
+ * Where the statement in `sql` ends: `end` is just past its last character of code (a trailing comment
+ * is not code, so the bounded wrapper's closing parenthesis can never land in one), and `terminator` is
+ * the index of its first `;` outside a literal, quoted identifier or comment, or -1.
+ */
+function statementEnd(sql: string): { readonly end: number; readonly terminator: number } {
+  let end = 0
+  let i = 0
+  while (i < sql.length) {
+    const char = sql.charAt(i)
+    if (char === ";") return { end, terminator: i }
+    const comment = commentEnd(sql, i)
+    if (comment !== undefined) {
+      i = comment
+      continue
+    }
+    const close = QUOTE_CLOSE.get(char)
+    i = close === undefined ? i + 1 : quotedEnd(sql, i, char, close)
+    if (char !== " " && char !== "\t" && char !== "\n" && char !== "\r" && char !== "\f") end = i
+  }
+  return { end, terminator: -1 }
+}
 
 // SQLite's bare-identifier alphabet. Non-ASCII is included because SQLite treats any byte >= 0x80 as
 // an identifier character; leaving it out would end a word early and split one table name into two.
@@ -167,7 +237,7 @@ function tokenizeSql(sql: string): SqlToken[] {
     tokens.push({ kind: "word", value, bare: false })
   }
   while (i < sql.length) {
-    const char = sql[i] as string
+    const char = sql.charAt(i)
     if (char === " " || char === "\t" || char === "\n" || char === "\r" || char === "\f") {
       i++
     } else if (char === "-" && sql[i + 1] === "-") {
@@ -303,7 +373,6 @@ function relationNames(sql: string): string[] {
   return [...new Set(names)]
 }
 
-/** Reject multi-statement input: allow one terminator only when it is the final character. */
 /**
  * Names a statement defines as CTEs, column lists included (`WITH RECURSIVE c(x) AS (...)`), which
  * {@link gateSqliteStatement}'s relation scan does not see past. Only for callers whose own check of
@@ -334,50 +403,10 @@ export function sqliteCteNames(sql: string): string[] {
   return [...names]
 }
 
+/** Reject multi-statement input: allow one terminator only when it is the final character. */
 const isSingleStatement = (sql: string): boolean => {
-  let quote: "'" | '"' | null = null
-  let lineComment = false
-  let blockComment = false
-  for (let i = 0; i < sql.length; i++) {
-    const char = sql[i]
-    const next = sql[i + 1]
-    if (lineComment) {
-      if (char === "\n") lineComment = false
-      continue
-    }
-    if (blockComment) {
-      if (char === "*" && next === "/") {
-        blockComment = false
-        i++
-      }
-      continue
-    }
-    if (quote !== null) {
-      if (char === quote) {
-        if (next === quote) i++
-        else quote = null
-      }
-      continue
-    }
-    if (char === "'" || char === '"') {
-      quote = char
-      continue
-    }
-    if (char === "-" && next === "-") {
-      lineComment = true
-      i++
-      continue
-    }
-    if (char === "/" && next === "*") {
-      blockComment = true
-      i++
-      continue
-    }
-    if (char === ";") {
-      return sql.slice(i + 1).trim() === ""
-    }
-  }
-  return true
+  const { terminator } = statementEnd(sql)
+  return terminator === -1 || sql.slice(terminator + 1).trim() === ""
 }
 
 const isReadStatement = (sql: string): boolean => {
@@ -385,58 +414,8 @@ const isReadStatement = (sql: string): boolean => {
   return head.startsWith("select") || head.startsWith("with")
 }
 
-/** Remove one trailing statement terminator (and anything after it) while preserving literals and
- * identifiers. Wrapping the query in a bounded subquery must not let a trailing comment consume the
- * wrapper's closing parenthesis. */
-function executableSql(sql: string): string {
-  let quote: "'" | '"' | null = null
-  let lineComment = false
-  let lineCommentStart = -1
-  let blockComment = false
-  for (let i = 0; i < sql.length; i++) {
-    const char = sql[i]
-    const next = sql[i + 1]
-    if (lineComment) {
-      if (char === "\n") {
-        lineComment = false
-        lineCommentStart = -1
-      }
-      continue
-    }
-    if (blockComment) {
-      if (char === "*" && next === "/") {
-        blockComment = false
-        i++
-      }
-      continue
-    }
-    if (quote !== null) {
-      if (char === quote) {
-        if (next === quote) i++
-        else quote = null
-      }
-      continue
-    }
-    if (char === "'" || char === '"') {
-      quote = char
-      continue
-    }
-    if (char === "-" && next === "-") {
-      lineComment = true
-      lineCommentStart = i
-      i++
-      continue
-    }
-    if (char === "/" && next === "*") {
-      blockComment = true
-      i++
-      continue
-    }
-    if (char === ";") return sql.slice(0, i).trim()
-  }
-  if (lineComment && lineCommentStart >= 0) return sql.slice(0, lineCommentStart).trim()
-  return sql.trim()
-}
+/** The statement without its terminator or a trailing comment, literals and identifiers intact. */
+const executableSql = (sql: string): string => sql.slice(0, statementEnd(sql).end)
 
 /** The outcome of {@link gateSqliteStatement}. A refusal's `text` is the message `run_query` returns. */
 export type SqliteGate =
@@ -451,8 +430,8 @@ 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`/`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`).
+ * to wrap. The compiled program ({@link unexposedProgramRelation}) remains the authority on the
+ * tables a statement reads; this check refuses a named relation before SQLite compiles anything.
  */
 export function gateSqliteStatement(
   input: string,
@@ -486,7 +465,9 @@ export function gateSqliteStatement(
 
 /**
  * The first relation an `EXPLAIN QUERY PLAN` result scans or searches that `exposed` rejects, as the
- * plan spells it, or `undefined` when every scanned relation is exposed.
+ * plan spells it, or `undefined` when every scanned relation is exposed. The plan names a scan by its
+ * alias (`FROM habits h` plans as `SCAN h`, `users AS habits` as `SCAN habits`), so it cannot decide
+ * on its own which tables a statement reads: the compiled program names them by root page.
  */
 export function unexposedPlanRelation(
   planRows: readonly { readonly detail?: unknown }[],

package/src/postgres.ts

@@ -39,7 +39,12 @@ import {
   type PostgresClient,
   records,
 } from "./internal/pg-exec.ts"
-import { DENIED_FUNCTION_PREFIXES, lintPostgres, withoutTerminator } from "./internal/pg-lexer.ts"
+import {
+  DENIED_FUNCTION_PREFIXES,
+  lintPostgres,
+  tableAliasLists,
+  withoutTerminator,
+} from "./internal/pg-lexer.ts"
 import {
   type DbPlan,
   type DbRedaction,
@@ -342,9 +347,14 @@ function scopeMatcher(scope: PostgresScope): (relation: PlanRelation) => string
  * The plan-level table check: each relation, and every table it inherits from (partitions too). A
  * `pg_*` function in FROM reads server state the way a catalog does - a system view such as
  * `pg_stat_activity` plans as one once the planner drops its catalog joins - so it is out of scope too.
- * With `masked`, a plan expression that uses a masked column, or a whole row holding one, is refused.
+ * With `masked`, a plan expression that uses a masked column, or a whole row holding one, is refused,
+ * and so is any masked column of a relation `statement` reads when it gives a table a column alias list.
  */
-function checkScope(scope: PostgresScope, masked?: (column: string) => boolean) {
+function checkScope(
+  scope: PostgresScope,
+  masked: ((column: string) => boolean) | undefined,
+  statement: string,
+) {
   const outOfScope = scopeMatcher(scope)
   return async (connection: PgConnection, reads: PlanReads): Promise<DbRefusal | undefined> => {
     for (const fn of reads.functions) {
@@ -393,6 +403,9 @@ function checkScope(scope: PostgresScope, masked?: (column: string) => boolean)
       if (column === "*") wholeRows.push(relation)
       else if (masked(column)) hits.add(`${relation.name}.${column}`)
     }
+    // The plan names an alias-listed scan's columns by the list, which can reuse a real column's name.
+    const renamed = tableAliasLists(statement)
+    if (renamed.length > 0) wholeRows.push(...reads.relations)
     const listed = [...wholeRows, ...reads.outputs.map((output) => output.relation)]
     if (listed.length > 0) {
       const live = new Map<string, string[]>()
@@ -423,6 +436,13 @@ function checkScope(scope: PostgresScope, masked?: (column: string) => boolean)
       }
     }
     if (hits.size === 0) return undefined
+    const [aliased] = renamed
+    if (aliased !== undefined) {
+      return dbRefusal(
+        "NIFRA_DB_COLUMN_REFUSED",
+        `the query renames the columns of ${aliased.table} with the alias list ${aliased.alias}(...), so the plan cannot show which columns it reads, and it reads a relation holding ${[...hits].join(", ")}, masked as ${hits.size === 1 ? "a credential" : "credentials"}; alias columns in the select list instead`,
+      )
+    }
     return dbRefusal(
       "NIFRA_DB_COLUMN_REFUSED",
       `the query reads ${[...hits].join(", ")}, masked as ${hits.size === 1 ? "a credential" : "credentials"}`,
@@ -465,7 +485,7 @@ export async function queryPostgres(
   const fetched = await gated(client, sql, options, (connection, statement) =>
     fetchThroughCursor(connection, statement, {
       maxRows: options.maxRows,
-      scope: checkScope(options, options.redaction?.column),
+      scope: checkScope(options, options.redaction?.column, statement),
     }),
   )
   if (isDbRefusal(fetched)) return fetched
@@ -494,7 +514,7 @@ export async function explainPostgres(
   const explained = await gated(client, sql, options, (connection, statement) =>
     explainThroughServer(connection, statement, {
       analyze: options.analyze === true,
-      scope: checkScope(options, options.redaction?.column),
+      scope: checkScope(options, options.redaction?.column, statement),
     }),
   )
   if (isDbRefusal(explained)) return explained