postgres-engine: scope search statement_timeout to the transaction

garagon · garagon · commit 6146c3b470dc · 2026-04-16T15:04:36.000-03:00
searchKeyword and searchVector run on a pooled postgres.js client
(max: 10 by default). The original code bounded each search with

  await sql`SET statement_timeout = '8s'`
  try { await sql`&lt;query&gt;` }
  finally { await sql`SET statement_timeout = '0'` }

but every tagged template is an independent round-trip that picks an
arbitrary connection from the pool. The SET, the query, and the reset
could all land on DIFFERENT connections. In practice the GUC sticks
to whichever connection ran the SET and then gets returned to the
pool — the next unrelated caller on that connection inherits the 8s
timeout (clipping legitimate long queries) or the reset-to-0 (disabling
the guard for whoever expected it). A crash in the middle leaves the
state set permanently.

Wrap each search in sql.begin(async sql =&gt; …). postgres.js reserves
a single connection for the transaction body, so the SET LOCAL, the
query, and the implicit COMMIT all run on the same connection. SET
LOCAL scopes the GUC to the transaction — COMMIT or ROLLBACK restores
the previous value automatically, regardless of the code path out.
Error paths can no longer leak the GUC.

No API change. Timeout value and semantics are identical (8s cap on
search queries, no effect on embed --all / bulk import which runs
outside these methods). Only one transaction per search — BEGIN +
COMMIT round-trips are negligible next to a ranked FTS or pgvector
query.

Also closes the earlier audit finding R4-F002 which reported the same
pattern on searchKeyword. This PR covers both searchKeyword and
searchVector so the pool-leak class is fully closed.

Tests (test/postgres-engine.test.ts, new file):
- No bare SET statement_timeout remains after stripping comments.
- searchKeyword and searchVector each wrap their query in sql.begin.
- Both use SET LOCAL.
- Neither explicitly clears the timeout with SET statement_timeout=0.

Source-level guardrails keep the fast unit suite DB-free. Live
Postgres coverage of the search path is in test/e2e/search-quality.test.ts,
which continues to exercise these methods end-to-end against
pgvector when DATABASE_URL is set.
diff --git a/src/core/postgres-engine.ts b/src/core/postgres-engine.ts
@@ -191,11 +191,17 @@ export class PostgresEngine implements BrainEngine {
     const detailLow = opts?.detail === 'low';
 
     // Search-only timeout: prevents DoS via expensive queries without
-    // affecting long-running operations like embed --all or bulk import
-    await sql`SET statement_timeout = '8s'`;
-    try {
+    // affecting long-running operations like embed --all or bulk import.
+    // SET LOCAL inside sql.begin() scopes the GUC to the transaction so
+    // it can never leak onto a pooled connection returned to other
+    // callers. A bare `SET statement_timeout` goes to an arbitrary
+    // connection from the pool, lives past this method, and either
+    // clips an unrelated caller's long-running query (DoS) or — via
+    // `SET statement_timeout = 0` — disables the guard for them.
+    const rows = await sql.begin(async sql => {
+      await sql`SET LOCAL statement_timeout = '8s'`;
       // CTE: rank pages by FTS score, then pick the best chunk per page in SQL
-      const rows = await sql`
+      return await sql`
         WITH ranked_pages AS (
           SELECT p.id, p.slug, p.title, p.type,
             ts_rank(p.search_vector, websearch_to_tsquery('english', ${query})) AS score
@@ -221,10 +227,8 @@ export class PostgresEngine implements BrainEngine {
         FROM best_chunks
         ORDER BY score DESC
       `;
-      return rows.map(rowToSearchResult);
-    } finally {
-      await sql`SET statement_timeout = '0'`;
-    }
+    });
+    return rows.map(rowToSearchResult);
   }
 
   async searchVector(embedding: Float32Array, opts?: SearchOpts): Promise<SearchResult[]> {
@@ -241,10 +245,12 @@ export class PostgresEngine implements BrainEngine {
 
     const vecStr = '[' + Array.from(embedding).join(',') + ']';
 
-    // Search-only timeout (see searchKeyword for rationale)
-    await sql`SET statement_timeout = '8s'`;
-    try {
-      const rows = await sql`
+    // Search-only timeout (see searchKeyword for rationale). SET LOCAL +
+    // sql.begin ensures the GUC stays transaction-scoped on the pooled
+    // connection.
+    const rows = await sql.begin(async sql => {
+      await sql`SET LOCAL statement_timeout = '8s'`;
+      return await sql`
         SELECT
           p.slug, p.id as page_id, p.title, p.type,
           cc.id as chunk_id, cc.chunk_index, cc.chunk_text, cc.chunk_source,
@@ -260,10 +266,8 @@ export class PostgresEngine implements BrainEngine {
         LIMIT ${limit}
         OFFSET ${offset}
       `;
-      return rows.map(rowToSearchResult);
-    } finally {
-      await sql`SET statement_timeout = '0'`;
-    }
+    });
+    return rows.map(rowToSearchResult);
   }
 
   async getEmbeddingsByChunkIds(ids: number[]): Promise<Map<number, Float32Array>> {
diff --git a/test/postgres-engine.test.ts b/test/postgres-engine.test.ts
@@ -0,0 +1,112 @@
+/**
+ * postgres-engine.ts source-level guardrails.
+ *
+ * Live Postgres coverage for search paths lives in test/e2e/search-quality.test.ts.
+ * This file stays fast and DB-free: it inspects the source of
+ * src/core/postgres-engine.ts to lock in decisions that protect the
+ * shared connection pool from per-request GUC leaks.
+ *
+ * Regression: R6-F006 / R4-F002.
+ * searchKeyword and searchVector used to call bare
+ *   await sql`SET statement_timeout = '8s'`
+ *   ...query...
+ *   finally { await sql`SET statement_timeout = '0'` }
+ * against the shared pool. Each tagged template picks an arbitrary
+ * connection, so the SET, the query, and the reset could all land on
+ * DIFFERENT connections. Worst case: the 8s GUC sticks on some pooled
+ * connection and clips the next caller's long-running query; or the
+ * reset to 0 lands on a connection that other code expected to be
+ * protected. The fix wraps each query in sql.begin() and uses
+ * SET LOCAL so the GUC is transaction-scoped and auto-resets on
+ * COMMIT/ROLLBACK, regardless of error path.
+ */
+
+import { describe, test, expect } from 'bun:test';
+import { readFileSync } from 'fs';
+import { join } from 'path';
+
+const SRC = readFileSync(
+  join(import.meta.dir, '..', 'src', 'core', 'postgres-engine.ts'),
+  'utf-8',
+);
+
+describe('postgres-engine / search path timeout isolation', () => {
+  test('no bare `SET statement_timeout` statement survives', () => {
+    // Strip comments so the commentary mentioning the anti-pattern does
+    // not trigger a false positive. Block-comment + line-comment strip.
+    const stripped = SRC
+      .replace(/\/\*[\s\S]*?\*\//g, '')
+      .replace(/(^|\s)\/\/[^\n]*/g, '$1');
+
+    // Match a tagged-template statement of the form
+    //   sql`SET statement_timeout = ...`
+    // that is NOT preceded by LOCAL. This is the exact shape that bleeds
+    // onto pooled connections; SET LOCAL is safe inside a transaction.
+    const bare = stripped.match(
+      /sql`\s*SET\s+(?!LOCAL\s)statement_timeout\b[^`]*`/gi,
+    );
+    expect(bare).toBeNull();
+  });
+
+  test('searchKeyword wraps its query in sql.begin()', () => {
+    const fn = extractMethod(SRC, 'searchKeyword');
+    expect(fn).toMatch(/sql\.begin\s*\(\s*async\s+sql\s*=>/);
+  });
+
+  test('searchVector wraps its query in sql.begin()', () => {
+    const fn = extractMethod(SRC, 'searchVector');
+    expect(fn).toMatch(/sql\.begin\s*\(\s*async\s+sql\s*=>/);
+  });
+
+  test('both search methods use SET LOCAL for the timeout', () => {
+    const keyword = extractMethod(SRC, 'searchKeyword');
+    const vector = extractMethod(SRC, 'searchVector');
+    expect(keyword).toMatch(/SET\s+LOCAL\s+statement_timeout/);
+    expect(vector).toMatch(/SET\s+LOCAL\s+statement_timeout/);
+  });
+
+  test('neither search method clears the timeout with `SET statement_timeout = 0`', () => {
+    // The reset-to-zero pattern was the other half of the leak: if SET
+    // LOCAL is in play, COMMIT handles the reset and an explicit
+    // `SET statement_timeout = '0'` would itself leak the GUC change
+    // onto the returned connection. Strip comments first so the
+    // commentary in the method itself (which quotes the anti-pattern
+    // to explain it) does not trigger a false positive.
+    const keyword = stripComments(extractMethod(SRC, 'searchKeyword'));
+    const vector = stripComments(extractMethod(SRC, 'searchVector'));
+    expect(keyword).not.toMatch(/SET\s+statement_timeout\s*=\s*['"]?0/);
+    expect(vector).not.toMatch(/SET\s+statement_timeout\s*=\s*['"]?0/);
+  });
+});
+
+function stripComments(s: string): string {
+  return s
+    .replace(/\/\*[\s\S]*?\*\//g, '')
+    .replace(/(^|\s)\/\/[^\n]*/g, '$1');
+}
+
+// extractMethod grabs the body of a class method by brace-matching from
+// its opening line. Returns the method body up to the matching closing
+// brace. Good enough for the small number of methods in this file.
+function extractMethod(source: string, name: string): string {
+  // Find "async <name>(" at method-definition indentation (2 spaces).
+  const openRe = new RegExp(`^\\s+async\\s+${name}\\s*\\(`, 'm');
+  const match = openRe.exec(source);
+  if (!match) {
+    throw new Error(`method ${name} not found in postgres-engine.ts`);
+  }
+  // Scan forward balancing braces.
+  let i = source.indexOf('{', match.index);
+  if (i < 0) throw new Error(`no opening brace for ${name}`);
+  const start = i;
+  let depth = 0;
+  for (; i < source.length; i++) {
+    const c = source[i];
+    if (c === '{') depth++;
+    else if (c === '}') {
+      depth--;
+      if (depth === 0) return source.slice(start, i + 1);
+    }
+  }
+  throw new Error(`unbalanced braces in ${name}`);
+}
