postgresai 0.14.0-dev.7 → 0.14.0-dev.71

package/lib/init.ts

@@ -1,6 +1,37 @@
-import * as readline from "readline";
-import { URL } from "url";
+import { randomBytes } from "crypto";
+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";
+
+/**
+ * Database provider type. Affects which prepare-db steps are executed.
+ * Known providers have specific behavior adjustments; unknown providers use default behavior.
+ * TODO: Consider auto-detecting provider from connection string or server version string.
+ * TODO: Consider making this more flexible via a config that specifies which steps/checks to skip.
+ */
+export type DbProvider = string;
+
+/** Known providers with special handling. Unknown providers are treated as self-managed. */
+export const KNOWN_PROVIDERS = ["self-managed", "supabase"] as const;
+
+/** Providers where we skip role creation (users managed externally). */
+const SKIP_ROLE_CREATION_PROVIDERS = ["supabase"];
+
+/** Providers where we skip ALTER USER statements (restricted by provider). */
+const SKIP_ALTER_USER_PROVIDERS = ["supabase"];
+
+/** Providers where we skip search_path verification (not set via ALTER USER). */
+const SKIP_SEARCH_PATH_CHECK_PROVIDERS = ["supabase"];
+
+/** Check if a provider is known and return a warning message if not. */
+export function validateProvider(provider: string | undefined): string | null {
+  if (!provider || KNOWN_PROVIDERS.includes(provider as any)) return null;
+  return `Unknown provider "${provider}". Known providers: ${KNOWN_PROVIDERS.join(", ")}. Treating as self-managed.`;
+}
 
 export type PgClientConfig = {
   connectionString?: string;
@@ -9,14 +40,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;
@@ -30,11 +185,64 @@ export type InitPlan = {
   steps: InitStep[];
 };
 
+function sqlDir(): string {
+  // 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 {
+  const p = path.join(sqlDir(), filename);
+  return fs.readFileSync(p, "utf8");
+}
+
+function applyTemplate(sql: string, vars: Record<string, string>): string {
+  return sql.replace(/\{\{([A-Z0-9_]+)\}\}/g, (_, key) => {
+    const v = vars[key];
+    if (v === undefined) throw new Error(`Missing SQL template var: ${key}`);
+    return v;
+  });
+}
+
 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 {
@@ -98,6 +306,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("=");
@@ -126,12 +335,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;
 }
 
@@ -158,8 +375,12 @@ export function resolveAdminConnection(opts: {
   const conn = (opts.conn || "").trim();
   const dbUrlFlag = (opts.dbUrlFlag || "").trim();
 
-  const hasPsqlParts =
-    !!(opts.host || opts.port || opts.username || opts.dbname || opts.adminPassword || opts.envPassword);
+  // 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);
 
   if (conn && dbUrlFlag) {
     throw new Error("Provide either positional connection string or --db-url, not both");
@@ -168,84 +389,94 @@ 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 (!hasPsqlParts) {
-    throw new Error(
-      "Connection is required. Provide a connection string/conninfo as a positional arg, or use --db-url, or use -h/-p/-U/-d."
-    );
+  if (!hasConnDetails) {
+    // Keep this message short: the CLI prints full help (including examples) on this error.
+    throw new Error("Connection is required.");
   }
 
   const cfg: PgClientConfig = {};
   if (opts.host) cfg.host = opts.host;
-  if (opts.port !== undefined && opts.port !== "") cfg.port = Number(opts.port);
+  if (opts.port !== undefined && opts.port !== "") {
+    const p = Number(opts.port);
+    if (!Number.isFinite(p) || !Number.isInteger(p) || p <= 0 || p > 65535) {
+      throw new Error(`Invalid port value: ${String(opts.port)}`);
+    }
+    cfg.port = p;
+  }
   if (opts.username) cfg.user = opts.username;
   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) };
+  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 };
+  }
+  // Default: try SSL with fallback (sslmode=prefer behavior)
+  cfg.ssl = { rejectUnauthorized: false };
+  return { clientConfig: cfg, display: describePgConfig(cfg), sslFallbackEnabled: true };
 }
 
-export async function promptHidden(prompt: string): Promise<string> {
-  const rl = readline.createInterface({
-    input: process.stdin,
-    output: process.stdout,
-    terminal: true,
-  });
-
-  // Mask input by overriding internal write method.
-  const anyRl = rl as any;
-  const out = process.stdout as NodeJS.WriteStream;
-  anyRl._writeToOutput = (str: string) => {
-    // Keep newlines and carriage returns; mask everything else.
-    if (str === "\n" || str === "\r\n") {
-      out.write(str);
-    } else {
-      out.write("*");
-    }
-  };
-
-  try {
-    const answer = await new Promise<string>((resolve) => rl.question(prompt, resolve));
-    // Ensure we end the masked line cleanly.
-    process.stdout.write("\n");
-    return answer;
-  } finally {
-    rl.close();
+function generateMonitoringPassword(): string {
+  // 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<string> {
+}): Promise<{ password: string; generated: boolean }> {
   const fromFlag = (opts.passwordFlag || "").trim();
-  if (fromFlag) return fromFlag;
+  if (fromFlag) return { password: fromFlag, generated: false };
 
   const fromEnv = (opts.passwordEnv || "").trim();
-  if (fromEnv) return fromEnv;
-
-  if (!process.stdin.isTTY) {
-    throw new Error(
-      "Monitoring user password is required in non-interactive mode (use --password or PGAI_MON_PASSWORD)"
-    );
-  }
+  if (fromEnv) return { password: fromEnv, generated: false };
 
-  const prompter = opts.prompt || promptHidden;
-  while (true) {
-    const pw = (await prompter(`Enter password for monitoring user ${opts.monitoringUser}: `)).trim();
-    if (pw) return pw;
-    // eslint-disable-next-line no-console
-    console.error("Password cannot be empty");
-  }
+  // Default: auto-generate (safer than prompting; works in non-interactive mode).
+  return { password: generateMonitoringPassword(), generated: true };
 }
 
 export async function buildInitPlan(params: {
@@ -253,106 +484,87 @@ export async function buildInitPlan(params: {
   monitoringUser?: string;
   monitoringPassword: string;
   includeOptionalPermissions: boolean;
-  roleExists?: boolean;
+  /** Provider type. Affects which steps are included. Defaults to "self-managed". */
+  provider?: DbProvider;
 }): 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 provider = params.provider ?? "self-managed";
 
   const qRole = quoteIdent(monitoringUser);
   const qDb = quoteIdent(database);
+  const qPw = quoteLiteral(params.monitoringPassword);
+  const qRoleNameLit = quoteLiteral(monitoringUser);
 
   const steps: InitStep[] = [];
 
-  // Role creation/update is done in two alternative steps. Caller decides by checking role existence.
-  if (params.roleExists === false) {
-    steps.push({
-      name: "create monitoring user",
-      sql: `create user ${qRole} with password $1;`,
-      params: [params.monitoringPassword],
-    });
-  } else if (params.roleExists === true) {
-    steps.push({
-      name: "update monitoring user password",
-      sql: `alter user ${qRole} with password $1;`,
-      params: [params.monitoringPassword],
-    });
-  } else {
-    // Unknown: caller will rebuild after probing role existence.
+  const vars: Record<string, string> = {
+    ROLE_IDENT: qRole,
+    DB_IDENT: qDb,
+  };
+
+  // Some providers (e.g., Supabase) manage users externally - skip role creation.
+  // TODO: Make this more flexible by allowing users to specify which steps to skip via config.
+  if (!SKIP_ROLE_CREATION_PROVIDERS.includes(provider)) {
+    // Role creation/update is done in one template file.
+    // 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
+    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 });
   }
 
-  steps.push(
-    {
-      name: "grant connect on database",
-      sql: `grant connect on database ${qDb} to ${qRole};`,
-    },
-    {
-      name: "grant pg_monitor",
-      sql: `grant pg_monitor to ${qRole};`,
-    },
-    {
-      name: "grant select on pg_index",
-      sql: `grant select on pg_catalog.pg_index to ${qRole};`,
-    },
-    {
-      name: "create or replace public.pg_statistic view",
-      sql: `create or replace view public.pg_statistic as
-select
-    n.nspname as schemaname,
-    c.relname as tablename,
-    a.attname,
-    s.stanullfrac as null_frac,
-    s.stawidth as avg_width,
-    false as inherited
-from pg_catalog.pg_statistic s
-join pg_catalog.pg_class c on c.oid = s.starelid
-join pg_catalog.pg_namespace n on n.oid = c.relnamespace
-join pg_catalog.pg_attribute a on a.attrelid = s.starelid and a.attnum = s.staattnum
-where a.attnum > 0 and not a.attisdropped;`,
-    },
-    {
-      name: "grant select on public.pg_statistic",
-      sql: `grant select on public.pg_statistic to ${qRole};`,
-    },
-    {
-      name: "ensure access to public schema (for hardened clusters)",
-      sql: `grant usage on schema public to ${qRole};`,
-    },
-    {
-      name: "set monitoring user search_path",
-      sql: `alter user ${qRole} set search_path = "$user", public, pg_catalog;`,
-    }
-  );
+  let permissionsSql = applyTemplate(loadSqlTemplate("02.permissions.sql"), vars);
+
+  // Some providers restrict ALTER USER - remove those statements.
+  // TODO: Make this more flexible by allowing users to specify which statements to skip via config.
+  if (SKIP_ALTER_USER_PROVIDERS.includes(provider)) {
+    permissionsSql = permissionsSql
+      .split("\n")
+      .filter((line) => {
+        const trimmed = line.trim();
+        // Keep comments and empty lines
+        if (trimmed.startsWith("--") || trimmed === "") return true;
+        // Filter out ALTER USER statements (case-insensitive, flexible whitespace)
+        return !/^\s*alter\s+user\s+/i.test(line);
+      })
+      .join("\n");
+  }
+
+  steps.push({
+    name: "02.permissions",
+    sql: permissionsSql,
+  });
+
+  // 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(
       {
-        name: "create rds_tools extension (optional)",
-        sql: "create extension if not exists rds_tools;",
-        optional: true,
-      },
-      {
-        name: "grant rds_tools.pg_ls_multixactdir() (optional)",
-        sql: `grant execute on function rds_tools.pg_ls_multixactdir() to ${qRole};`,
-        optional: true,
-      },
-      {
-        name: "grant pg_stat_file(text) (optional)",
-        sql: `grant execute on function pg_catalog.pg_stat_file(text) to ${qRole};`,
-        optional: true,
-      },
-      {
-        name: "grant pg_stat_file(text, boolean) (optional)",
-        sql: `grant execute on function pg_catalog.pg_stat_file(text, boolean) to ${qRole};`,
-        optional: true,
-      },
-      {
-        name: "grant pg_ls_dir(text) (optional)",
-        sql: `grant execute on function pg_catalog.pg_ls_dir(text) to ${qRole};`,
+        name: "03.optional_rds",
+        sql: applyTemplate(loadSqlTemplate("03.optional_rds.sql"), vars),
         optional: true,
       },
       {
-        name: "grant pg_ls_dir(text, boolean, boolean) (optional)",
-        sql: `grant execute on function pg_catalog.pg_ls_dir(text, boolean, boolean) to ${qRole};`,
+        name: "04.optional_self_managed",
+        sql: applyTemplate(loadSqlTemplate("04.optional_self_managed.sql"), vars),
         optional: true,
       }
     );
@@ -369,28 +581,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);
-        throw new Error(`Failed at step "${step.name}": ${msg}`);
+        await params.client.query("rollback;");
+      } catch {
+        // ignore
       }
+      throw e;
+    }
+  };
+
+  // Apply non-optional steps, each in its own transaction
+  for (const step of params.plan.steps.filter((s) => !s.optional)) {
+    try {
+      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;
     }
-    await params.client.query("commit;");
-  } catch (e) {
-    await params.client.query("rollback;");
-    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);
@@ -401,4 +651,167 @@ export async function applyInitPlan(params: {
   return { applied, skippedOptional };
 }
 
+export type VerifyInitResult = {
+  ok: boolean;
+  missingRequired: string[];
+  missingOptional: string[];
+};
+
+export async function verifyInitSetup(params: {
+  client: PgClient;
+  database: string;
+  monitoringUser: string;
+  includeOptionalPermissions: boolean;
+  /** Provider type. Affects which checks are performed. */
+  provider?: DbProvider;
+}): Promise<VerifyInitResult> {
+  // 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 provider = params.provider ?? "self-managed";
+
+    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 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");
+    }
+
+    // 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 (!schemaExistsRes.rows?.[0]?.ok) {
+      missingRequired.push("USAGE on schema postgres_ai");
+    }
+
+    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 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");
+    }
+
+    // Some providers don't allow setting search_path via ALTER USER - skip this check.
+    // TODO: Make this more flexible by allowing users to specify which checks to skip via config.
+    if (!SKIP_SEARCH_PATH_CHECK_PROVIDERS.includes(provider)) {
+      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 ${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
+    }
+  }
+}
+