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

package/lib/init.ts

@@ -0,0 +1,452 @@
+import * as readline from "readline";
+import { randomBytes } from "crypto";
+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.",
+        "",
+        "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")
+    );
+  }
+
+  const cfg: PgClientConfig = {};
+  if (opts.host) cfg.host = opts.host;
+  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) };
+}
+
+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");
+  }
+
+  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();
+  });
+}
+
+function generateMonitoringPassword(): string {
+  // URL-safe and easy to copy/paste; length ~32 chars.
+  return randomBytes(24).toString("base64url");
+}
+
+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();
+  if (fromFlag) return { password: fromFlag, generated: false };
+
+  const fromEnv = (opts.passwordEnv || "").trim();
+  if (fromEnv) return { password: fromEnv, generated: false };
+
+  // Default: auto-generate (safer than prompting; works in non-interactive mode).
+  return { password: generateMonitoringPassword(), generated: true };
+}
+
+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.8",
   "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",