postgresai 0.12.0-beta.7 → 0.14.0-dev.7

package/lib/init.ts

@@ -0,0 +1,404 @@
+import * as readline from "readline";
+import { URL } from "url";
+import type { Client as PgClient } from "pg";
+
+export type PgClientConfig = {
+  connectionString?: string;
+  host?: string;
+  port?: number;
+  user?: string;
+  password?: string;
+  database?: string;
+  ssl?: any;
+};
+
+export type AdminConnection = {
+  clientConfig: PgClientConfig;
+  display: string;
+};
+
+export type InitStep = {
+  name: string;
+  sql: string;
+  params?: unknown[];
+  optional?: boolean;
+};
+
+export type InitPlan = {
+  monitoringUser: string;
+  database: string;
+  steps: InitStep[];
+};
+
+function quoteIdent(ident: string): string {
+  // Always quote. Escape embedded quotes by doubling.
+  return `"${ident.replace(/"/g, "\"\"")}"`;
+}
+
+export function maskConnectionString(dbUrl: string): string {
+  // Hide password if present (postgresql://user:pass@host/db).
+  try {
+    const u = new URL(dbUrl);
+    if (u.password) u.password = "*****";
+    return u.toString();
+  } catch {
+    return dbUrl.replace(/\/\/([^:/?#]+):([^@/?#]+)@/g, "//$1:*****@");
+  }
+}
+
+function isLikelyUri(value: string): boolean {
+  return /^postgres(ql)?:\/\//i.test(value.trim());
+}
+
+function tokenizeConninfo(input: string): string[] {
+  const s = input.trim();
+  const tokens: string[] = [];
+  let i = 0;
+
+  const isSpace = (ch: string) => ch === " " || ch === "\t" || ch === "\n" || ch === "\r";
+
+  while (i < s.length) {
+    while (i < s.length && isSpace(s[i]!)) i++;
+    if (i >= s.length) break;
+
+    let tok = "";
+    let inSingle = false;
+    while (i < s.length) {
+      const ch = s[i]!;
+      if (!inSingle && isSpace(ch)) break;
+
+      if (ch === "'" && !inSingle) {
+        inSingle = true;
+        i++;
+        continue;
+      }
+      if (ch === "'" && inSingle) {
+        inSingle = false;
+        i++;
+        continue;
+      }
+
+      if (ch === "\\" && i + 1 < s.length) {
+        tok += s[i + 1]!;
+        i += 2;
+        continue;
+      }
+
+      tok += ch;
+      i++;
+    }
+
+    tokens.push(tok);
+    while (i < s.length && isSpace(s[i]!)) i++;
+  }
+
+  return tokens;
+}
+
+export function parseLibpqConninfo(input: string): PgClientConfig {
+  const tokens = tokenizeConninfo(input);
+  const cfg: PgClientConfig = {};
+
+  for (const t of tokens) {
+    const eq = t.indexOf("=");
+    if (eq <= 0) continue;
+    const key = t.slice(0, eq).trim();
+    const rawVal = t.slice(eq + 1);
+    const val = rawVal.trim();
+    if (!key) continue;
+
+    switch (key) {
+      case "host":
+        cfg.host = val;
+        break;
+      case "port": {
+        const p = Number(val);
+        if (Number.isFinite(p)) cfg.port = p;
+        break;
+      }
+      case "user":
+        cfg.user = val;
+        break;
+      case "password":
+        cfg.password = val;
+        break;
+      case "dbname":
+      case "database":
+        cfg.database = val;
+        break;
+      // ignore everything else (sslmode, options, application_name, etc.)
+      default:
+        break;
+    }
+  }
+
+  return cfg;
+}
+
+export function describePgConfig(cfg: PgClientConfig): string {
+  if (cfg.connectionString) return maskConnectionString(cfg.connectionString);
+  const user = cfg.user ? cfg.user : "<user>";
+  const host = cfg.host ? cfg.host : "<host>";
+  const port = cfg.port ? String(cfg.port) : "<port>";
+  const db = cfg.database ? cfg.database : "<db>";
+  // Don't include password
+  return `postgresql://${user}:*****@${host}:${port}/${db}`;
+}
+
+export function resolveAdminConnection(opts: {
+  conn?: string;
+  dbUrlFlag?: string;
+  host?: string;
+  port?: string | number;
+  username?: string;
+  dbname?: string;
+  adminPassword?: string;
+  envPassword?: string;
+}): AdminConnection {
+  const conn = (opts.conn || "").trim();
+  const dbUrlFlag = (opts.dbUrlFlag || "").trim();
+
+  const hasPsqlParts =
+    !!(opts.host || opts.port || opts.username || opts.dbname || opts.adminPassword || opts.envPassword);
+
+  if (conn && dbUrlFlag) {
+    throw new Error("Provide either positional connection string or --db-url, not both");
+  }
+
+  if (conn || dbUrlFlag) {
+    const v = conn || dbUrlFlag;
+    if (isLikelyUri(v)) {
+      return { clientConfig: { connectionString: v }, display: maskConnectionString(v) };
+    }
+    // libpq conninfo (dbname=... host=...)
+    const cfg = parseLibpqConninfo(v);
+    if (opts.envPassword && !cfg.password) cfg.password = opts.envPassword;
+    return { clientConfig: cfg, display: describePgConfig(cfg) };
+  }
+
+  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."
+    );
+  }
+
+  const cfg: PgClientConfig = {};
+  if (opts.host) cfg.host = opts.host;
+  if (opts.port !== undefined && opts.port !== "") cfg.port = Number(opts.port);
+  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) };
+}
+
+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();
+  }
+}
+
+export async function resolveMonitoringPassword(opts: {
+  passwordFlag?: string;
+  passwordEnv?: string;
+  prompt?: (prompt: string) => Promise<string>;
+  monitoringUser: string;
+}): Promise<string> {
+  const fromFlag = (opts.passwordFlag || "").trim();
+  if (fromFlag) return fromFlag;
+
+  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)"
+    );
+  }
+
+  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");
+  }
+}
+
+export async function buildInitPlan(params: {
+  database: string;
+  monitoringUser?: string;
+  monitoringPassword: string;
+  includeOptionalPermissions: boolean;
+  roleExists?: boolean;
+}): Promise<InitPlan> {
+  const monitoringUser = params.monitoringUser || "postgres_ai_mon";
+  const database = params.database;
+
+  const qRole = quoteIdent(monitoringUser);
+  const qDb = quoteIdent(database);
+
+  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.
+  }
+
+  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;`,
+    }
+  );
+
+  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};`,
+        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};`,
+        optional: true,
+      }
+    );
+  }
+
+  return { monitoringUser, database, steps };
+}
+
+export async function applyInitPlan(params: {
+  client: PgClient;
+  plan: InitPlan;
+  verbose?: boolean;
+}): Promise<{ applied: string[]; skippedOptional: string[] }> {
+  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)) {
+      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("commit;");
+  } catch (e) {
+    await params.client.query("rollback;");
+    throw e;
+  }
+
+  // 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);
+    } catch {
+      skippedOptional.push(step.name);
+      // best-effort: ignore
+    }
+  }
+
+  return { applied, skippedOptional };
+}
+
+

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "postgresai",
-  "version": "0.12.0-beta.7",
+  "version": "0.14.0-dev.7",
   "description": "postgres_ai CLI (Node.js)",
   "license": "Apache-2.0",
   "private": false,
@@ -25,7 +25,8 @@
     "build": "tsc",
     "prepare": "npm run build",
     "start": "node ./dist/bin/postgres-ai.js --help",
-    "dev": "tsc --watch"
+    "dev": "tsc --watch",
+    "test": "npm run build && node --test test/*.test.cjs"
   },
   "dependencies": {
     "@modelcontextprotocol/sdk": "^1.20.2",

package/test/init.integration.test.cjs

@@ -0,0 +1,269 @@
+const test = require("node:test");
+const assert = require("node:assert/strict");
+const fs = require("node:fs");
+const os = require("node:os");
+const path = require("node:path");
+const net = require("node:net");
+const { spawn, spawnSync } = require("node:child_process");
+
+function findOnPath(cmd) {
+  const which = spawnSync("sh", ["-lc", `command -v ${cmd}`], { encoding: "utf8" });
+  if (which.status === 0) return String(which.stdout || "").trim();
+  return null;
+}
+
+function findPgBin(cmd) {
+  const p = findOnPath(cmd);
+  if (p) return p;
+
+  // Debian/Ubuntu (GitLab CI node:*-bullseye images): binaries usually live here.
+  // We avoid filesystem globbing in JS and just ask the shell.
+  const probe = spawnSync(
+    "sh",
+    [
+      "-lc",
+      `ls -1 /usr/lib/postgresql/*/bin/${cmd} 2>/dev/null | head -n 1 || true`,
+    ],
+    { encoding: "utf8" }
+  );
+  const out = String(probe.stdout || "").trim();
+  if (out) return out;
+
+  return null;
+}
+
+function havePostgresBinaries() {
+  return !!(findPgBin("initdb") && findPgBin("postgres"));
+}
+
+async function getFreePort() {
+  return await new Promise((resolve, reject) => {
+    const srv = net.createServer();
+    srv.listen(0, "127.0.0.1", () => {
+      const addr = srv.address();
+      srv.close((err) => {
+        if (err) return reject(err);
+        resolve(addr.port);
+      });
+    });
+    srv.on("error", reject);
+  });
+}
+
+async function waitFor(fn, { timeoutMs = 10000, intervalMs = 100 } = {}) {
+  const start = Date.now();
+  // eslint-disable-next-line no-constant-condition
+  while (true) {
+    try {
+      return await fn();
+    } catch (e) {
+      if (Date.now() - start > timeoutMs) throw e;
+      await new Promise((r) => setTimeout(r, intervalMs));
+    }
+  }
+}
+
+async function withTempPostgres(t) {
+  const tmpRoot = fs.mkdtempSync(path.join(os.tmpdir(), "postgresai-init-"));
+  const dataDir = path.join(tmpRoot, "data");
+  const socketDir = path.join(tmpRoot, "sock");
+  fs.mkdirSync(socketDir, { recursive: true });
+
+  const initdb = findPgBin("initdb");
+  const postgresBin = findPgBin("postgres");
+  assert.ok(initdb && postgresBin, "PostgreSQL binaries not found (need initdb and postgres)");
+
+  const init = spawnSync(initdb, ["-D", dataDir, "-U", "postgres", "-A", "trust"], {
+    encoding: "utf8",
+  });
+  assert.equal(init.status, 0, init.stderr || init.stdout);
+
+  // Configure: local socket trust, TCP scram.
+  const hbaPath = path.join(dataDir, "pg_hba.conf");
+  fs.appendFileSync(
+    hbaPath,
+    "\n# Added by postgresai init integration tests\nlocal all all trust\nhost all all 127.0.0.1/32 scram-sha-256\nhost all all ::1/128 scram-sha-256\n",
+    "utf8"
+  );
+
+  const port = await getFreePort();
+
+  const postgresProc = spawn(postgresBin, ["-D", dataDir, "-k", socketDir, "-h", "127.0.0.1", "-p", String(port)], {
+    stdio: ["ignore", "pipe", "pipe"],
+  });
+
+  const { Client } = require("pg");
+
+  const connectLocal = async (database = "postgres") => {
+    // IMPORTANT: must match the port Postgres is started with; otherwise pg defaults to 5432 and the socket path won't exist.
+    const c = new Client({ host: socketDir, port, user: "postgres", database });
+    await c.connect();
+    return c;
+  };
+
+  await waitFor(async () => {
+    const c = await connectLocal();
+    await c.end();
+  });
+
+  const postgresPassword = "postgrespw";
+  {
+    const c = await connectLocal();
+    await c.query("alter user postgres password $1", [postgresPassword]);
+    await c.query("create database testdb");
+    await c.end();
+  }
+
+  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 });
+  });
+
+  const adminUri = `postgresql://postgres:${postgresPassword}@127.0.0.1:${port}/testdb`;
+  return { port, socketDir, adminUri, postgresPassword };
+}
+
+async function runCliInit(args, env = {}) {
+  const node = process.execPath;
+  const cliPath = path.resolve(__dirname, "..", "dist", "bin", "postgres-ai.js");
+  const res = spawnSync(node, [cliPath, "init", ...args], {
+    encoding: "utf8",
+    env: { ...process.env, ...env },
+  });
+  return res;
+}
+
+test(
+  "integration: init supports URI / conninfo / psql-like connection styles",
+  { skip: !havePostgresBinaries() },
+  async (t) => {
+    const pg = await withTempPostgres(t);
+
+    // 1) positional URI
+    {
+      const r = await runCliInit([pg.adminUri, "--password", "monpw", "--skip-optional-permissions"]);
+      assert.equal(r.status, 0, r.stderr || r.stdout);
+    }
+
+    // 2) conninfo
+    {
+      const conninfo = `dbname=testdb host=127.0.0.1 port=${pg.port} user=postgres password=${pg.postgresPassword}`;
+      const r = await runCliInit([conninfo, "--password", "monpw2", "--skip-optional-permissions"]);
+      assert.equal(r.status, 0, r.stderr || r.stdout);
+    }
+
+    // 3) psql-like options (+ PGPASSWORD)
+    {
+      const r = await runCliInit(
+        [
+          "-h",
+          "127.0.0.1",
+          "-p",
+          String(pg.port),
+          "-U",
+          "postgres",
+          "-d",
+          "testdb",
+          "--password",
+          "monpw3",
+          "--skip-optional-permissions",
+        ],
+        { PGPASSWORD: pg.postgresPassword }
+      );
+      assert.equal(r.status, 0, r.stderr || r.stdout);
+    }
+  }
+);
+
+test(
+  "integration: init fixes slightly-off permissions idempotently",
+  { skip: !havePostgresBinaries() },
+  async (t) => {
+    const pg = await withTempPostgres(t);
+    const { Client } = require("pg");
+
+    // Create monitoring role with wrong password, no grants.
+    {
+      const c = new Client({ connectionString: pg.adminUri });
+      await c.connect();
+      await c.query(
+        "do $$ begin if not exists (select 1 from pg_roles where rolname='postgres_ai_mon') then create role postgres_ai_mon login password 'wrong'; end if; end $$;"
+      );
+      await c.end();
+    }
+
+    // Run init (should grant everything).
+    {
+      const r = await runCliInit([pg.adminUri, "--password", "correctpw", "--skip-optional-permissions"]);
+      assert.equal(r.status, 0, r.stderr || r.stdout);
+    }
+
+    // Verify privileges.
+    {
+      const c = new Client({ connectionString: pg.adminUri });
+      await c.connect();
+      const dbOk = await c.query(
+        "select has_database_privilege('postgres_ai_mon', current_database(), 'CONNECT') as ok"
+      );
+      assert.equal(dbOk.rows[0].ok, true);
+      const roleOk = await c.query("select pg_has_role('postgres_ai_mon', 'pg_monitor', 'member') as ok");
+      assert.equal(roleOk.rows[0].ok, true);
+      const idxOk = await c.query(
+        "select has_table_privilege('postgres_ai_mon', 'pg_catalog.pg_index', 'SELECT') as ok"
+      );
+      assert.equal(idxOk.rows[0].ok, true);
+      const viewOk = await c.query(
+        "select has_table_privilege('postgres_ai_mon', 'public.pg_statistic', 'SELECT') as ok"
+      );
+      assert.equal(viewOk.rows[0].ok, true);
+      const sp = await c.query("select rolconfig from pg_roles where rolname='postgres_ai_mon'");
+      assert.ok(Array.isArray(sp.rows[0].rolconfig));
+      assert.ok(sp.rows[0].rolconfig.some((v) => String(v).includes("search_path=")));
+      await c.end();
+    }
+
+    // Run init again (idempotent).
+    {
+      const r = await runCliInit([pg.adminUri, "--password", "correctpw", "--skip-optional-permissions"]);
+      assert.equal(r.status, 0, r.stderr || r.stdout);
+    }
+  }
+);
+
+test("integration: init reports nicely when lacking permissions", { skip: !havePostgresBinaries() }, async (t) => {
+  const pg = await withTempPostgres(t);
+  const { Client } = require("pg");
+
+  // Create limited user that can connect but cannot create roles / grant.
+  const limitedPw = "limitedpw";
+  {
+    const c = new Client({ connectionString: pg.adminUri });
+    await c.connect();
+    await c.query(
+      "do $$ begin if not exists (select 1 from pg_roles where rolname='limited') then create role limited login password $1; end if; end $$;",
+      [limitedPw]
+    );
+    await c.query("grant connect on database testdb to limited");
+    await c.end();
+  }
+
+  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:/);
+  // Should include step context and hint.
+  assert.match(r.stderr, /Failed at step "/);
+  assert.match(r.stderr, /Hint: connect as a superuser/i);
+});
+
+