postgresai 0.14.0-beta.1 → 0.14.0-beta.10

package/lib/init.ts

@@ -1,10 +1,12 @@
-import * as readline from "readline";
 import { randomBytes } from "crypto";
-import { URL } from "url";
+import { URL, fileURLToPath } from "url";
+import type { ConnectionOptions as TlsConnectionOptions } from "tls";
 import type { Client as PgClient } from "pg";
 import * as fs from "fs";
 import * as path from "path";
 
+export const DEFAULT_MONITORING_USER = "postgres_ai_mon";
+
 export type PgClientConfig = {
   connectionString?: string;
   host?: string;
@@ -12,14 +14,138 @@ export type PgClientConfig = {
   user?: string;
   password?: string;
   database?: string;
-  ssl?: any;
+  ssl?: boolean | TlsConnectionOptions;
 };
 
+/**
+ * Convert PostgreSQL sslmode to node-postgres ssl config.
+ */
+function sslModeToConfig(mode: string): boolean | TlsConnectionOptions {
+  if (mode.toLowerCase() === "disable") return false;
+  if (mode.toLowerCase() === "verify-full" || mode.toLowerCase() === "verify-ca") return true;
+  // For require/prefer/allow: encrypt without certificate verification
+  return { rejectUnauthorized: false };
+}
+
+/** Extract sslmode from a PostgreSQL connection URI. */
+function extractSslModeFromUri(uri: string): string | undefined {
+  try {
+    return new URL(uri).searchParams.get("sslmode") ?? undefined;
+  } catch {
+    return uri.match(/[?&]sslmode=([^&]+)/i)?.[1];
+  }
+}
+
+/** Remove sslmode parameter from a PostgreSQL connection URI. */
+function stripSslModeFromUri(uri: string): string {
+  try {
+    const u = new URL(uri);
+    u.searchParams.delete("sslmode");
+    return u.toString();
+  } catch {
+    // Fallback regex for malformed URIs
+    return uri
+      .replace(/[?&]sslmode=[^&]*/gi, "")
+      .replace(/\?&/, "?")
+      .replace(/\?$/, "");
+  }
+}
+
 export type AdminConnection = {
   clientConfig: PgClientConfig;
   display: string;
+  /** True if SSL fallback is enabled (try SSL first, fall back to non-SSL on failure). */
+  sslFallbackEnabled?: boolean;
 };
 
+/**
+ * Check if an error indicates SSL negotiation failed and fallback to non-SSL should be attempted.
+ * This mimics libpq's sslmode=prefer behavior.
+ * 
+ * IMPORTANT: This should NOT match certificate errors (expired, invalid, self-signed)
+ * as those are real errors the user needs to fix, not negotiation failures.
+ */
+function isSslNegotiationError(err: unknown): boolean {
+  if (!err || typeof err !== "object") return false;
+  const e = err as any;
+  const msg = typeof e.message === "string" ? e.message.toLowerCase() : "";
+  const code = typeof e.code === "string" ? e.code : "";
+  
+  // Specific patterns that indicate server doesn't support SSL (should fallback)
+  const fallbackPatterns = [
+    "the server does not support ssl",
+    "ssl off",
+    "server does not support ssl connections",
+  ];
+  
+  for (const pattern of fallbackPatterns) {
+    if (msg.includes(pattern)) return true;
+  }
+  
+  // PostgreSQL error code 08P01 (protocol violation) during initial connection
+  // often indicates SSL negotiation mismatch, but only if the message suggests it
+  if (code === "08P01" && (msg.includes("ssl") || msg.includes("unsupported"))) {
+    return true;
+  }
+  
+  return false;
+}
+
+/**
+ * Connect to PostgreSQL with sslmode=prefer-like behavior.
+ * If sslFallbackEnabled is true, tries SSL first, then falls back to non-SSL on failure.
+ */
+export async function connectWithSslFallback(
+  ClientClass: new (config: PgClientConfig) => PgClient,
+  adminConn: AdminConnection,
+  verbose?: boolean
+): Promise<{ client: PgClient; usedSsl: boolean }> {
+  const tryConnect = async (config: PgClientConfig): Promise<PgClient> => {
+    const client = new ClientClass(config);
+    await client.connect();
+    return client;
+  };
+
+  // If SSL was explicitly set or no SSL configured, just try once
+  if (!adminConn.sslFallbackEnabled) {
+    const client = await tryConnect(adminConn.clientConfig);
+    return { client, usedSsl: !!adminConn.clientConfig.ssl };
+  }
+
+  // sslmode=prefer behavior: try SSL first, fallback to non-SSL
+  try {
+    const client = await tryConnect(adminConn.clientConfig);
+    return { client, usedSsl: true };
+  } catch (sslErr) {
+    if (!isSslNegotiationError(sslErr)) {
+      // Not an SSL error, don't retry
+      throw sslErr;
+    }
+
+    if (verbose) {
+      console.log("SSL connection failed, retrying without SSL...");
+    }
+
+    // Retry without SSL
+    const noSslConfig: PgClientConfig = { ...adminConn.clientConfig, ssl: false };
+    try {
+      const client = await tryConnect(noSslConfig);
+      return { client, usedSsl: false };
+    } catch (noSslErr) {
+      // If non-SSL also fails, check if it's "SSL required" - throw that instead
+      if (isSslNegotiationError(noSslErr)) {
+        const msg = (noSslErr as any)?.message || "";
+        if (msg.toLowerCase().includes("ssl") && msg.toLowerCase().includes("required")) {
+          // Server requires SSL but SSL attempt failed - throw original SSL error
+          throw sslErr;
+        }
+      }
+      // Throw the non-SSL error (it's more relevant since SSL attempt also failed)
+      throw noSslErr;
+    }
+  }
+}
+
 export type InitStep = {
   name: string;
   sql: string;
@@ -33,13 +159,27 @@ export type InitPlan = {
   steps: InitStep[];
 };
 
-function packageRootDirFromCompiled(): string {
-  // dist/lib/init.js -> <pkg>/dist/lib ; package root is ../..
-  return path.resolve(__dirname, "..", "..");
-}
-
 function sqlDir(): string {
-  return path.join(packageRootDirFromCompiled(), "sql");
+  // Handle both development and production paths
+  // Development: lib/init.ts -> ../sql
+  // Production (bundled): dist/bin/postgres-ai.js -> ../sql (copied during build)
+  //
+  // IMPORTANT: Use import.meta.url instead of __dirname because bundlers (bun/esbuild)
+  // bake in __dirname at build time, while import.meta.url resolves at runtime.
+  const currentFile = fileURLToPath(import.meta.url);
+  const currentDir = path.dirname(currentFile);
+
+  const candidates = [
+    path.resolve(currentDir, "..", "sql"),       // bundled: dist/bin -> dist/sql
+    path.resolve(currentDir, "..", "..", "sql"), // dev from lib: lib -> ../sql
+  ];
+
+  for (const candidate of candidates) {
+    if (fs.existsSync(candidate)) {
+      return candidate;
+    }
+  }
+  throw new Error(`SQL directory not found. Searched: ${candidates.join(", ")}`);
 }
 
 function loadSqlTemplate(filename: string): string {
@@ -57,15 +197,26 @@ function applyTemplate(sql: string, vars: Record<string, string>): string {
 
 function quoteIdent(ident: string): string {
   // Always quote. Escape embedded quotes by doubling.
+  if (ident.includes("\0")) {
+    throw new Error("Identifier cannot contain null bytes");
+  }
   return `"${ident.replace(/"/g, "\"\"")}"`;
 }
 
 function quoteLiteral(value: string): string {
   // Single-quote and escape embedded quotes by doubling.
   // This is used where Postgres grammar requires a literal (e.g., CREATE/ALTER ROLE PASSWORD).
+  if (value.includes("\0")) {
+    throw new Error("Literal cannot contain null bytes");
+  }
   return `'${value.replace(/'/g, "''")}'`;
 }
 
+export function redactPasswordsInSql(sql: string): string {
+  // Replace PASSWORD '<literal>' (handles doubled quotes inside).
+  return sql.replace(/password\s+'(?:''|[^'])*'/gi, "password '<redacted>'");
+}
+
 export function maskConnectionString(dbUrl: string): string {
   // Hide password if present (postgresql://user:pass@host/db).
   try {
@@ -129,6 +280,7 @@ function tokenizeConninfo(input: string): string[] {
 export function parseLibpqConninfo(input: string): PgClientConfig {
   const tokens = tokenizeConninfo(input);
   const cfg: PgClientConfig = {};
+  let sslmode: string | undefined;
 
   for (const t of tokens) {
     const eq = t.indexOf("=");
@@ -157,12 +309,20 @@ export function parseLibpqConninfo(input: string): PgClientConfig {
       case "database":
         cfg.database = val;
         break;
-      // ignore everything else (sslmode, options, application_name, etc.)
+      case "sslmode":
+        sslmode = val;
+        break;
+      // ignore everything else (options, application_name, etc.)
       default:
         break;
     }
   }
 
+  // Apply SSL configuration based on sslmode
+  if (sslmode) {
+    cfg.ssl = sslModeToConfig(sslmode);
+  }
+
   return cfg;
 }
 
@@ -189,6 +349,9 @@ export function resolveAdminConnection(opts: {
   const conn = (opts.conn || "").trim();
   const dbUrlFlag = (opts.dbUrlFlag || "").trim();
 
+  // Resolve explicit SSL setting from environment (undefined = auto-detect)
+  const explicitSsl = process.env.PGSSLMODE;
+
   // NOTE: passwords alone (PGPASSWORD / --admin-password) do NOT constitute a connection.
   // We require at least some connection addressing (host/port/user/db) if no positional arg / --db-url is provided.
   const hasConnDetails = !!(opts.host || opts.port || opts.username || opts.dbname);
@@ -200,28 +363,45 @@ export function resolveAdminConnection(opts: {
   if (conn || dbUrlFlag) {
     const v = conn || dbUrlFlag;
     if (isLikelyUri(v)) {
-      return { clientConfig: { connectionString: v }, display: maskConnectionString(v) };
+      const urlSslMode = extractSslModeFromUri(v);
+      const effectiveSslMode = explicitSsl || urlSslMode;
+      // SSL priority: PGSSLMODE env > URL param > auto (sslmode=prefer behavior)
+      const sslConfig = effectiveSslMode
+        ? sslModeToConfig(effectiveSslMode)
+        : { rejectUnauthorized: false }; // Default: try SSL (with fallback)
+      // Enable fallback for: no explicit mode OR explicit "prefer"/"allow"
+      const shouldFallback = !effectiveSslMode || 
+        effectiveSslMode.toLowerCase() === "prefer" || 
+        effectiveSslMode.toLowerCase() === "allow";
+      // Strip sslmode from URI so pg uses our ssl config object instead
+      const cleanUri = stripSslModeFromUri(v);
+      return {
+        clientConfig: { connectionString: cleanUri, ssl: sslConfig },
+        display: maskConnectionString(v),
+        sslFallbackEnabled: shouldFallback,
+      };
     }
     // libpq conninfo (dbname=... host=...)
     const cfg = parseLibpqConninfo(v);
     if (opts.envPassword && !cfg.password) cfg.password = opts.envPassword;
-    return { clientConfig: cfg, display: describePgConfig(cfg) };
+    const cfgHadSsl = cfg.ssl !== undefined;
+    if (cfg.ssl === undefined) {
+      if (explicitSsl) cfg.ssl = sslModeToConfig(explicitSsl);
+      else cfg.ssl = { rejectUnauthorized: false }; // Default: try SSL (with fallback)
+    }
+    // Enable fallback for: no explicit mode OR explicit "prefer"/"allow"
+    const shouldFallback = (!explicitSsl && !cfgHadSsl) || 
+      (!!explicitSsl && (explicitSsl.toLowerCase() === "prefer" || explicitSsl.toLowerCase() === "allow"));
+    return {
+      clientConfig: cfg,
+      display: describePgConfig(cfg),
+      sslFallbackEnabled: shouldFallback,
+    };
   }
 
   if (!hasConnDetails) {
-    throw new Error(
-      [
-        "Connection is required.",
-        "",
-        "Examples:",
-        "  postgresai init postgresql://admin@host:5432/dbname",
-        "  postgresai init \"dbname=dbname host=host user=admin\"",
-        "  postgresai init -h host -p 5432 -U admin -d dbname",
-        "",
-        "Admin password:",
-        "  --admin-password <password>  (or set PGPASSWORD)",
-      ].join("\n")
-    );
+    // Keep this message short: the CLI prints full help (including examples) on this error.
+    throw new Error("Connection is required.");
   }
 
   const cfg: PgClientConfig = {};
@@ -237,84 +417,30 @@ export function resolveAdminConnection(opts: {
   if (opts.dbname) cfg.database = opts.dbname;
   if (opts.adminPassword) cfg.password = opts.adminPassword;
   if (opts.envPassword && !cfg.password) cfg.password = opts.envPassword;
-  return { clientConfig: cfg, display: describePgConfig(cfg) };
-}
-
-export async function promptHidden(prompt: string): Promise<string> {
-  // Implement our own hidden input reader so:
-  // - prompt text is visible
-  // - only user input is masked
-  // - we don't rely on non-public readline internals
-  if (!process.stdin.isTTY) {
-    throw new Error("Cannot prompt for password in non-interactive mode");
+  if (explicitSsl) {
+    cfg.ssl = sslModeToConfig(explicitSsl);
+    // Enable fallback for explicit "prefer"/"allow"
+    const shouldFallback = explicitSsl.toLowerCase() === "prefer" || explicitSsl.toLowerCase() === "allow";
+    return { clientConfig: cfg, display: describePgConfig(cfg), sslFallbackEnabled: shouldFallback };
   }
-
-  const stdin = process.stdin;
-  const stdout = process.stdout as NodeJS.WriteStream;
-
-  stdout.write(prompt);
-
-  return await new Promise<string>((resolve, reject) => {
-    let value = "";
-
-    const cleanup = () => {
-      try {
-        stdin.setRawMode(false);
-      } catch {
-        // ignore
-      }
-      stdin.removeListener("keypress", onKeypress);
-    };
-
-    const onKeypress = (str: string, key: any) => {
-      if (key?.ctrl && key?.name === "c") {
-        stdout.write("\n");
-        cleanup();
-        reject(new Error("Cancelled"));
-        return;
-      }
-
-      if (key?.name === "return" || key?.name === "enter") {
-        stdout.write("\n");
-        cleanup();
-        resolve(value);
-        return;
-      }
-
-      if (key?.name === "backspace") {
-        if (value.length > 0) {
-          value = value.slice(0, -1);
-          // Erase one mask char.
-          stdout.write("\b \b");
-        }
-        return;
-      }
-
-      // Ignore other control keys.
-      if (key?.ctrl || key?.meta) return;
-
-      if (typeof str === "string" && str.length > 0) {
-        value += str;
-        stdout.write("*");
-      }
-    };
-
-    readline.emitKeypressEvents(stdin);
-    stdin.setRawMode(true);
-    stdin.on("keypress", onKeypress);
-    stdin.resume();
-  });
+  // Default: try SSL with fallback (sslmode=prefer behavior)
+  cfg.ssl = { rejectUnauthorized: false };
+  return { clientConfig: cfg, display: describePgConfig(cfg), sslFallbackEnabled: true };
 }
 
 function generateMonitoringPassword(): string {
-  // URL-safe and easy to copy/paste; length ~32 chars.
-  return randomBytes(24).toString("base64url");
+  // URL-safe and easy to copy/paste; 24 bytes => 32 base64url chars (no padding).
+  // Note: randomBytes() throws on failure; we add a tiny sanity check for unexpected output.
+  const password = randomBytes(24).toString("base64url");
+  if (password.length < 30) {
+    throw new Error("Password generation failed: unexpected output length");
+  }
+  return password;
 }
 
 export async function resolveMonitoringPassword(opts: {
   passwordFlag?: string;
   passwordEnv?: string;
-  prompt?: (prompt: string) => Promise<string>;
   monitoringUser: string;
 }): Promise<{ password: string; generated: boolean }> {
   const fromFlag = (opts.passwordFlag || "").trim();
@@ -332,9 +458,9 @@ export async function buildInitPlan(params: {
   monitoringUser?: string;
   monitoringPassword: string;
   includeOptionalPermissions: boolean;
-  roleExists?: boolean;
 }): Promise<InitPlan> {
-  const monitoringUser = params.monitoringUser || "postgres_ai_mon";
+  // NOTE: kept async for API stability / potential future async template loading.
+  const monitoringUser = params.monitoringUser || DEFAULT_MONITORING_USER;
   const database = params.database;
 
   const qRole = quoteIdent(monitoringUser);
@@ -344,27 +470,26 @@ export async function buildInitPlan(params: {
 
   const steps: InitStep[] = [];
 
-  const vars = {
+  const vars: Record<string, string> = {
     ROLE_IDENT: qRole,
     DB_IDENT: qDb,
   };
 
   // Role creation/update is done in one template file.
-  // If roleExists is unknown, use a single DO block to create-or-alter safely.
-  let roleStmt: string | null = null;
-  if (params.roleExists === false) {
-    roleStmt = `create user ${qRole} with password ${qPw};`;
-  } else if (params.roleExists === true) {
-    roleStmt = `alter user ${qRole} with password ${qPw};`;
-  } else {
-    roleStmt = `do $$ begin
+  // Always use a single DO block to avoid race conditions between "role exists?" checks and CREATE USER.
+  // We:
+  // - create role if missing (and handle duplicate_object in case another session created it concurrently),
+  // - then ALTER ROLE to ensure the password is set to the desired value.
+  const roleStmt = `do $$ begin
   if not exists (select 1 from pg_catalog.pg_roles where rolname = ${qRoleNameLit}) then
-    create user ${qRole} with password ${qPw};
-  else
-    alter user ${qRole} with password ${qPw};
+    begin
+      create user ${qRole} with password ${qPw};
+    exception when duplicate_object then
+      null;
+    end;
   end if;
+  alter user ${qRole} with password ${qPw};
 end $$;`;
-  }
 
   const roleSql = applyTemplate(loadSqlTemplate("01.role.sql"), { ...vars, ROLE_STMT: roleStmt });
   steps.push({ name: "01.role", sql: roleSql });
@@ -374,6 +499,12 @@ end $$;`;
     sql: applyTemplate(loadSqlTemplate("02.permissions.sql"), vars),
   });
 
+  // Helper functions (SECURITY DEFINER) for plan analysis and table info
+  steps.push({
+    name: "05.helpers",
+    sql: applyTemplate(loadSqlTemplate("05.helpers.sql"), vars),
+  });
+
   if (params.includeOptionalPermissions) {
     steps.push(
       {
@@ -400,39 +531,66 @@ export async function applyInitPlan(params: {
   const applied: string[] = [];
   const skippedOptional: string[] = [];
 
-  // Apply non-optional steps in a single transaction.
-  await params.client.query("begin;");
-  try {
-    for (const step of params.plan.steps.filter((s) => !s.optional)) {
+  // Helper to wrap a step execution in begin/commit
+  const executeStep = async (step: InitStep): Promise<void> => {
+    await params.client.query("begin;");
+    try {
+      await params.client.query(step.sql, step.params as any);
+      await params.client.query("commit;");
+    } catch (e) {
+      // Rollback errors should never mask the original failure.
       try {
-        await params.client.query(step.sql, step.params as any);
-        applied.push(step.name);
-      } catch (e) {
-        const msg = e instanceof Error ? e.message : String(e);
-        const errAny = e as any;
-        const wrapped: any = new Error(`Failed at step "${step.name}": ${msg}`);
-        // Preserve Postgres error code so callers can provide better hints (e.g., 42501 insufficient_privilege).
-        if (errAny && typeof errAny === "object" && typeof errAny.code === "string") {
-          wrapped.code = errAny.code;
-        }
-        throw wrapped;
+        await params.client.query("rollback;");
+      } catch {
+        // ignore
       }
+      throw e;
     }
-    await params.client.query("commit;");
-  } catch (e) {
-    // Rollback errors should never mask the original failure.
+  };
+
+  // Apply non-optional steps, each in its own transaction
+  for (const step of params.plan.steps.filter((s) => !s.optional)) {
     try {
-      await params.client.query("rollback;");
-    } catch {
-      // ignore
+      await executeStep(step);
+      applied.push(step.name);
+    } catch (e) {
+      const msg = e instanceof Error ? e.message : String(e);
+      const errAny = e as any;
+      const wrapped: any = new Error(`Failed at step "${step.name}": ${msg}`);
+      // Preserve useful Postgres error fields so callers can provide better hints / diagnostics.
+      const pgErrorFields = [
+        "code",
+        "detail",
+        "hint",
+        "position",
+        "internalPosition",
+        "internalQuery",
+        "where",
+        "schema",
+        "table",
+        "column",
+        "dataType",
+        "constraint",
+        "file",
+        "line",
+        "routine",
+      ] as const;
+      if (errAny && typeof errAny === "object") {
+        for (const field of pgErrorFields) {
+          if (errAny[field] !== undefined) wrapped[field] = errAny[field];
+        }
+      }
+      if (e instanceof Error && e.stack) {
+        wrapped.stack = e.stack;
+      }
+      throw wrapped;
     }
-    throw e;
   }
 
-  // Apply optional steps outside of the transaction so a failure doesn't abort everything.
+  // Apply optional steps, each in its own transaction (failure doesn't abort)
   for (const step of params.plan.steps.filter((s) => s.optional)) {
     try {
-      await params.client.query(step.sql, step.params as any);
+      await executeStep(step);
       applied.push(step.name);
     } catch {
       skippedOptional.push(step.name);
@@ -455,111 +613,148 @@ export async function verifyInitSetup(params: {
   monitoringUser: string;
   includeOptionalPermissions: boolean;
 }): Promise<VerifyInitResult> {
-  const missingRequired: string[] = [];
-  const missingOptional: string[] = [];
-
-  const role = params.monitoringUser;
-  const db = params.database;
-
-  const roleRes = await params.client.query("select 1 from pg_catalog.pg_roles where rolname = $1", [role]);
-  const roleExists = (roleRes.rowCount ?? 0) > 0;
-  if (!roleExists) {
-    missingRequired.push(`role "${role}" does not exist`);
-    // If role is missing, other checks will error or be meaningless.
-    return { ok: false, missingRequired, missingOptional };
-  }
+  // Use a repeatable-read snapshot so all checks see a consistent view.
+  await params.client.query("begin isolation level repeatable read;");
+  try {
+    const missingRequired: string[] = [];
+    const missingOptional: string[] = [];
+
+    const role = params.monitoringUser;
+    const db = params.database;
+
+    const roleRes = await params.client.query("select 1 from pg_catalog.pg_roles where rolname = $1", [role]);
+    const roleExists = (roleRes.rowCount ?? 0) > 0;
+    if (!roleExists) {
+      missingRequired.push(`role "${role}" does not exist`);
+      // If role is missing, other checks will error or be meaningless.
+      return { ok: false, missingRequired, missingOptional };
+    }
 
-  const connectRes = await params.client.query(
-    "select has_database_privilege($1, $2, 'CONNECT') as ok",
-    [role, db]
-  );
-  if (!connectRes.rows?.[0]?.ok) {
-    missingRequired.push(`CONNECT on database "${db}"`);
-  }
+    const connectRes = await params.client.query(
+      "select has_database_privilege($1, $2, 'CONNECT') as ok",
+      [role, db]
+    );
+    if (!connectRes.rows?.[0]?.ok) {
+      missingRequired.push(`CONNECT on database "${db}"`);
+    }
 
-  const pgMonitorRes = await params.client.query(
-    "select pg_has_role($1, 'pg_monitor', 'member') as ok",
-    [role]
-  );
-  if (!pgMonitorRes.rows?.[0]?.ok) {
-    missingRequired.push("membership in role pg_monitor");
-  }
+    const pgMonitorRes = await params.client.query(
+      "select pg_has_role($1, 'pg_monitor', 'member') as ok",
+      [role]
+    );
+    if (!pgMonitorRes.rows?.[0]?.ok) {
+      missingRequired.push("membership in role pg_monitor");
+    }
 
-  const pgIndexRes = await params.client.query(
-    "select has_table_privilege($1, 'pg_catalog.pg_index', 'SELECT') as ok",
-    [role]
-  );
-  if (!pgIndexRes.rows?.[0]?.ok) {
-    missingRequired.push("SELECT on pg_catalog.pg_index");
-  }
+    const pgIndexRes = await params.client.query(
+      "select has_table_privilege($1, 'pg_catalog.pg_index', 'SELECT') as ok",
+      [role]
+    );
+    if (!pgIndexRes.rows?.[0]?.ok) {
+      missingRequired.push("SELECT on pg_catalog.pg_index");
+    }
 
-  const viewExistsRes = await params.client.query("select to_regclass('public.pg_statistic') is not null as ok");
-  if (!viewExistsRes.rows?.[0]?.ok) {
-    missingRequired.push("view public.pg_statistic exists");
-  } else {
-    const viewPrivRes = await params.client.query(
-      "select has_table_privilege($1, 'public.pg_statistic', 'SELECT') as ok",
+    // Check postgres_ai schema exists and is usable
+    const schemaExistsRes = await params.client.query(
+      "select has_schema_privilege($1, 'postgres_ai', 'USAGE') as ok",
       [role]
     );
-    if (!viewPrivRes.rows?.[0]?.ok) {
-      missingRequired.push("SELECT on view public.pg_statistic");
+    if (!schemaExistsRes.rows?.[0]?.ok) {
+      missingRequired.push("USAGE on schema postgres_ai");
     }
-  }
 
-  const schemaUsageRes = await params.client.query(
-    "select has_schema_privilege($1, 'public', 'USAGE') as ok",
-    [role]
-  );
-  if (!schemaUsageRes.rows?.[0]?.ok) {
-    missingRequired.push("USAGE on schema public");
-  }
+    const viewExistsRes = await params.client.query("select to_regclass('postgres_ai.pg_statistic') is not null as ok");
+    if (!viewExistsRes.rows?.[0]?.ok) {
+      missingRequired.push("view postgres_ai.pg_statistic exists");
+    } else {
+      const viewPrivRes = await params.client.query(
+        "select has_table_privilege($1, 'postgres_ai.pg_statistic', 'SELECT') as ok",
+        [role]
+      );
+      if (!viewPrivRes.rows?.[0]?.ok) {
+        missingRequired.push("SELECT on view postgres_ai.pg_statistic");
+      }
+    }
 
-  const rolcfgRes = await params.client.query("select rolconfig from pg_catalog.pg_roles where rolname = $1", [role]);
-  const rolconfig = rolcfgRes.rows?.[0]?.rolconfig;
-  const spLine = Array.isArray(rolconfig) ? rolconfig.find((v: any) => String(v).startsWith("search_path=")) : undefined;
-  if (typeof spLine !== "string" || !spLine) {
-    missingRequired.push("role search_path is set");
-  } else {
-    // We accept any ordering as long as public and pg_catalog are included.
-    const sp = spLine.toLowerCase();
-    if (!sp.includes("public") || !sp.includes("pg_catalog")) {
-      missingRequired.push("role search_path includes public and pg_catalog");
+    const schemaUsageRes = await params.client.query(
+      "select has_schema_privilege($1, 'public', 'USAGE') as ok",
+      [role]
+    );
+    if (!schemaUsageRes.rows?.[0]?.ok) {
+      missingRequired.push("USAGE on schema public");
     }
-  }
 
-  if (params.includeOptionalPermissions) {
-    // Optional RDS/Aurora extras
-    {
-      const extRes = await params.client.query("select 1 from pg_extension where extname = 'rds_tools'");
-      if ((extRes.rowCount ?? 0) === 0) {
-        missingOptional.push("extension rds_tools");
-      } else {
-        const fnRes = await params.client.query(
-          "select has_function_privilege($1, 'rds_tools.pg_ls_multixactdir()', 'EXECUTE') as ok",
-          [role]
-        );
+    const rolcfgRes = await params.client.query("select rolconfig from pg_catalog.pg_roles where rolname = $1", [role]);
+    const rolconfig = rolcfgRes.rows?.[0]?.rolconfig;
+    const spLine = Array.isArray(rolconfig) ? rolconfig.find((v: any) => String(v).startsWith("search_path=")) : undefined;
+    if (typeof spLine !== "string" || !spLine) {
+      missingRequired.push("role search_path is set");
+    } else {
+      // We accept any ordering as long as postgres_ai, public, and pg_catalog are included.
+      const sp = spLine.toLowerCase();
+      if (!sp.includes("postgres_ai") || !sp.includes("public") || !sp.includes("pg_catalog")) {
+        missingRequired.push("role search_path includes postgres_ai, public and pg_catalog");
+      }
+    }
+
+    // Check for helper functions
+    const explainFnRes = await params.client.query(
+      "select has_function_privilege($1, 'postgres_ai.explain_generic(text, text, text)', 'EXECUTE') as ok",
+      [role]
+    );
+    if (!explainFnRes.rows?.[0]?.ok) {
+      missingRequired.push("EXECUTE on postgres_ai.explain_generic(text, text, text)");
+    }
+
+    const tableDescribeFnRes = await params.client.query(
+      "select has_function_privilege($1, 'postgres_ai.table_describe(text)', 'EXECUTE') as ok",
+      [role]
+    );
+    if (!tableDescribeFnRes.rows?.[0]?.ok) {
+      missingRequired.push("EXECUTE on postgres_ai.table_describe(text)");
+    }
+
+    if (params.includeOptionalPermissions) {
+      // Optional RDS/Aurora extras
+      {
+        const extRes = await params.client.query("select 1 from pg_extension where extname = 'rds_tools'");
+        if ((extRes.rowCount ?? 0) === 0) {
+          missingOptional.push("extension rds_tools");
+        } else {
+          const fnRes = await params.client.query(
+            "select has_function_privilege($1, 'rds_tools.pg_ls_multixactdir()', 'EXECUTE') as ok",
+            [role]
+          );
+          if (!fnRes.rows?.[0]?.ok) {
+            missingOptional.push("EXECUTE on rds_tools.pg_ls_multixactdir()");
+          }
+        }
+      }
+
+      // Optional self-managed extras
+      const optionalFns = [
+        "pg_catalog.pg_stat_file(text)",
+        "pg_catalog.pg_stat_file(text, boolean)",
+        "pg_catalog.pg_ls_dir(text)",
+        "pg_catalog.pg_ls_dir(text, boolean, boolean)",
+      ];
+      for (const fn of optionalFns) {
+        const fnRes = await params.client.query("select has_function_privilege($1, $2, 'EXECUTE') as ok", [role, fn]);
         if (!fnRes.rows?.[0]?.ok) {
-          missingOptional.push("EXECUTE on rds_tools.pg_ls_multixactdir()");
+          missingOptional.push(`EXECUTE on ${fn}`);
         }
       }
     }
 
-    // Optional self-managed extras
-    const optionalFns = [
-      "pg_catalog.pg_stat_file(text)",
-      "pg_catalog.pg_stat_file(text, boolean)",
-      "pg_catalog.pg_ls_dir(text)",
-      "pg_catalog.pg_ls_dir(text, boolean, boolean)",
-    ];
-    for (const fn of optionalFns) {
-      const fnRes = await params.client.query("select has_function_privilege($1, $2, 'EXECUTE') as ok", [role, fn]);
-      if (!fnRes.rows?.[0]?.ok) {
-        missingOptional.push(`EXECUTE on ${fn}`);
-      }
+    return { ok: missingRequired.length === 0, missingRequired, missingOptional };
+  } finally {
+    // Read-only: rollback to release snapshot; do not mask original errors.
+    try {
+      await params.client.query("rollback;");
+    } catch {
+      // ignore
     }
   }
-
-  return { ok: missingRequired.length === 0, missingRequired, missingOptional };
 }