postgresai 0.14.0-dev.14 → 0.14.0-dev.15

package/lib/init.ts

@@ -1,10 +1,13 @@
 import * as readline from "readline";
 import { randomBytes } from "crypto";
 import { URL } 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,7 +15,7 @@ export type PgClientConfig = {
   user?: string;
   password?: string;
   database?: string;
-  ssl?: any;
+  ssl?: boolean | TlsConnectionOptions;
 };
 
 export type AdminConnection = {
@@ -57,15 +60,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 {
@@ -307,8 +321,13 @@ export async function promptHidden(prompt: string): Promise<string> {
 }
 
 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: {
@@ -332,9 +351,8 @@ export async function buildInitPlan(params: {
   monitoringUser?: string;
   monitoringPassword: string;
   includeOptionalPermissions: boolean;
-  roleExists?: boolean;
 }): Promise<InitPlan> {
-  const monitoringUser = params.monitoringUser || "postgres_ai_mon";
+  const monitoringUser = params.monitoringUser || DEFAULT_MONITORING_USER;
   const database = params.database;
 
   const qRole = quoteIdent(monitoringUser);
@@ -350,21 +368,20 @@ export async function buildInitPlan(params: {
   };
 
   // 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 });
@@ -411,9 +428,31 @@ export async function applyInitPlan(params: {
         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;
+        // 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;
       }
@@ -432,11 +471,24 @@ export async function applyInitPlan(params: {
   // Apply optional steps outside of the transaction so a failure doesn't abort everything.
   for (const step of params.plan.steps.filter((s) => s.optional)) {
     try {
-      await params.client.query(step.sql, step.params as any);
-      applied.push(step.name);
+      // Run each optional step in its own mini-transaction to avoid partial application.
+      await params.client.query("begin;");
+      try {
+        await params.client.query(step.sql, step.params as any);
+        await params.client.query("commit;");
+        applied.push(step.name);
+      } catch {
+        try {
+          await params.client.query("rollback;");
+        } catch {
+          // ignore rollback errors
+        }
+        skippedOptional.push(step.name);
+        // best-effort: ignore
+      }
     } catch {
+      // If we can't even begin/commit, treat as skipped.
       skippedOptional.push(step.name);
-      // best-effort: ignore
     }
   }
 
@@ -455,111 +507,122 @@ 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 };
-  }
-
-  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}"`);
-  }
+  // 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 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 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 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 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 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",
+    const pgIndexRes = await params.client.query(
+      "select has_table_privilege($1, 'pg_catalog.pg_index', 'SELECT') as ok",
       [role]
     );
-    if (!viewPrivRes.rows?.[0]?.ok) {
-      missingRequired.push("SELECT on view public.pg_statistic");
+    if (!pgIndexRes.rows?.[0]?.ok) {
+      missingRequired.push("SELECT on pg_catalog.pg_index");
     }
-  }
 
-  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('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",
+        [role]
+      );
+      if (!viewPrivRes.rows?.[0]?.ok) {
+        missingRequired.push("SELECT on view public.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 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");
+      }
+    }
+
+    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 };
 }
 
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "postgresai",
-  "version": "0.14.0-dev.14",
+  "version": "0.14.0-dev.15",
   "description": "postgres_ai CLI (Node.js)",
   "license": "Apache-2.0",
   "private": false,

package/test/init.integration.test.cjs

@@ -92,25 +92,41 @@ async function withTempPostgres(t) {
 
   const port = await getFreePort();
 
-  const postgresProc = spawn(postgresBin, ["-D", dataDir, "-k", socketDir, "-h", "127.0.0.1", "-p", String(port)], {
-    stdio: ["ignore", "pipe", "pipe"],
-  });
-
-  // Register cleanup immediately so failures below don't leave a running postgres and hang CI.
-  t.after(async () => {
-    postgresProc.kill("SIGTERM");
+  let postgresProc;
+  try {
+    postgresProc = spawn(
+      postgresBin,
+      ["-D", dataDir, "-k", socketDir, "-h", "127.0.0.1", "-p", String(port)],
+      {
+        stdio: ["ignore", "pipe", "pipe"],
+      }
+    );
+
+    // Register cleanup immediately so failures below don't leave a running postgres and hang CI.
+    t.after(async () => {
+      postgresProc.kill("SIGTERM");
+      try {
+        await waitFor(
+          async () => {
+            if (postgresProc.exitCode === null) throw new Error("still running");
+          },
+          { timeoutMs: 5000, intervalMs: 100 }
+        );
+      } catch {
+        postgresProc.kill("SIGKILL");
+      }
+      fs.rmSync(tmpRoot, { recursive: true, force: true });
+    });
+  } catch (e) {
+    // If anything goes wrong before cleanup is registered, ensure we don't leak a running postgres.
     try {
-      await waitFor(
-        async () => {
-          if (postgresProc.exitCode === null) throw new Error("still running");
-        },
-        { timeoutMs: 5000, intervalMs: 100 }
-      );
+      if (postgresProc) postgresProc.kill("SIGKILL");
     } catch {
-      postgresProc.kill("SIGKILL");
+      // ignore
     }
     fs.rmSync(tmpRoot, { recursive: true, force: true });
-  });
+    throw e;
+  }
 
   const { Client } = require("pg");
 
@@ -208,8 +224,8 @@ test(
     {
       const r = await runCliInit([pg.adminUri, "--print-password", "--skip-optional-permissions"]);
       assert.equal(r.status, 0, r.stderr || r.stdout);
-      assert.match(r.stdout, /Generated monitoring password for postgres_ai_mon/i);
-      assert.match(r.stdout, /PGAI_MON_PASSWORD=/);
+      assert.match(r.stderr, /Generated monitoring password for postgres_ai_mon/i);
+      assert.match(r.stderr, /PGAI_MON_PASSWORD='/);
     }
   }
 );
@@ -294,12 +310,10 @@ test("integration: init reports nicely when lacking permissions", { skip: !haveP
   const limitedUri = `postgresql://limited:${limitedPw}@127.0.0.1:${pg.port}/testdb`;
   const r = await runCliInit([limitedUri, "--password", "monpw", "--skip-optional-permissions"]);
   assert.notEqual(r.status, 0);
-  assert.match(r.stderr, /init failed:/);
+  assert.match(r.stderr, /Error: init failed:/);
   // Should include step context and hint.
   assert.match(r.stderr, /Failed at step "/);
-  assert.match(r.stderr, /Permission error:/i);
-  assert.match(r.stderr, /How to fix:/i);
-  assert.match(r.stderr, /Hint: connect as a superuser/i);
+  assert.match(r.stderr, /Fix: connect as a superuser/i);
 });
 
 test("integration: init --verify returns 0 when ok and non-zero when missing", { skip: !havePostgresBinaries() }, async (t) => {