From 45d78a471aea04648840359c8badc6dc56f2501f Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 26 Dec 2025 04:18:46 +0000 Subject: [PATCH] feat(prepare-db): add helper functions and per-step transactions - Add postgres_ai schema for organizing our objects - Move pg_statistic view from public to postgres_ai schema - Add postgres_ai.explain_generic() function (SECURITY DEFINER) for collecting generic query plans with optional HypoPG index testing - Add postgres_ai.table_describe() function (SECURITY DEFINER) for collecting comprehensive table information for LLM analysis: - Table metadata (type, pages, estimated rows) - Column definitions with types, nullability, defaults - Indexes (primary key, unique, regular) - Constraints (FK, unique, check) - Foreign keys referencing this table - Vacuum/analyze statistics - Update search_path to include postgres_ai first - Change transaction model to wrap each step in its own begin/commit instead of grouping non-optional steps in a single transaction - Update verification to check for postgres_ai schema and new functions Relates to: https://gitlab.com/postgres-ai/postgres_ai/-/issues/68 --- cli/lib/init.ts | 160 ++++++----- cli/sql/02.permissions.sql | 14 +- cli/sql/05.helpers.sql | 415 +++++++++++++++++++++++++++++ cli/test/init.integration.test.cjs | 155 ++++++++++- cli/test/init.test.cjs | 63 ++++- 5 files changed, 730 insertions(+), 77 deletions(-) create mode 100644 cli/sql/05.helpers.sql diff --git a/cli/lib/init.ts b/cli/lib/init.ts index 77c97ea..92f344e 100644 --- a/cli/lib/init.ts +++ b/cli/lib/init.ts @@ -485,6 +485,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( { @@ -511,78 +517,70 @@ 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 => { + 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 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("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 { - // 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 - } + await executeStep(step); + applied.push(step.name); } catch { - // If we can't even begin/commit, treat as skipped. skippedOptional.push(step.name); + // best-effort: ignore } } @@ -642,16 +640,25 @@ export async function verifyInitSetup(params: { 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"); + // 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 public.pg_statistic exists"); + missingRequired.push("view postgres_ai.pg_statistic exists"); } else { const viewPrivRes = await params.client.query( - "select has_table_privilege($1, 'public.pg_statistic', 'SELECT') as ok", + "select has_table_privilege($1, 'postgres_ai.pg_statistic', 'SELECT') as ok", [role] ); if (!viewPrivRes.rows?.[0]?.ok) { - missingRequired.push("SELECT on view public.pg_statistic"); + missingRequired.push("SELECT on view postgres_ai.pg_statistic"); } } @@ -669,13 +676,30 @@ export async function verifyInitSetup(params: { 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. + // We accept any ordering as long as postgres_ai, 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 (!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 { diff --git a/cli/sql/02.permissions.sql b/cli/sql/02.permissions.sql index 7e760c1..d20851a 100644 --- a/cli/sql/02.permissions.sql +++ b/cli/sql/02.permissions.sql @@ -7,8 +7,12 @@ grant connect on database {{DB_IDENT}} to {{ROLE_IDENT}}; grant pg_monitor to {{ROLE_IDENT}}; grant select on pg_catalog.pg_index to {{ROLE_IDENT}}; --- Optional, for bloat analysis: expose pg_statistic via a view -create or replace view public.pg_statistic as +-- Create postgres_ai schema for our objects +create schema if not exists postgres_ai; +grant usage on schema postgres_ai to {{ROLE_IDENT}}; + +-- For bloat analysis: expose pg_statistic via a view +create or replace view postgres_ai.pg_statistic as select n.nspname as schemaname, c.relname as tablename, @@ -22,12 +26,12 @@ 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; -grant select on public.pg_statistic to {{ROLE_IDENT}}; +grant select on postgres_ai.pg_statistic to {{ROLE_IDENT}}; -- Hardened clusters sometimes revoke PUBLIC on schema public grant usage on schema public to {{ROLE_IDENT}}; --- Keep search_path predictable -alter user {{ROLE_IDENT}} set search_path = "$user", public, pg_catalog; +-- Keep search_path predictable; postgres_ai first so our objects are found +alter user {{ROLE_IDENT}} set search_path = postgres_ai, "$user", public, pg_catalog; diff --git a/cli/sql/05.helpers.sql b/cli/sql/05.helpers.sql new file mode 100644 index 0000000..427dc47 --- /dev/null +++ b/cli/sql/05.helpers.sql @@ -0,0 +1,415 @@ +-- Helper functions for postgres_ai monitoring user (template-filled by cli/lib/init.ts) +-- These functions use SECURITY DEFINER to allow the monitoring user to perform +-- operations they don't have direct permissions for. + +/* + * pgai_explain_generic + * + * Function to get generic explain plans with optional HypoPG index testing. + * Requires: PostgreSQL 16+ (for generic_plan option), HypoPG extension (optional). + * + * Usage examples: + * -- Basic generic plan + * select postgres_ai.explain_generic('select * from users where id = $1'); + * + * -- JSON format + * select postgres_ai.explain_generic('select * from users where id = $1', 'json'); + * + * -- Test a hypothetical index + * select postgres_ai.explain_generic( + * 'select * from users where email = $1', + * 'text', + * 'create index on users (email)' + * ); + */ +create or replace function postgres_ai.explain_generic( + in query text, + in format text default 'text', + in hypopg_index text default null, + out result text +) +language plpgsql +security definer +set search_path = pg_catalog, public +as $$ +declare + v_line record; + v_lines text[] := '{}'; + v_explain_query text; + v_hypo_result record; + v_version int; + v_hypopg_available boolean; +begin + -- Check PostgreSQL version (generic_plan requires 16+) + select current_setting('server_version_num')::int into v_version; + + if v_version < 160000 then + raise exception 'generic_plan requires PostgreSQL 16+, current version: %', + current_setting('server_version'); + end if; + + -- Check if HypoPG extension is available + if hypopg_index is not null then + select exists( + select 1 from pg_extension where extname = 'hypopg' + ) into v_hypopg_available; + + if not v_hypopg_available then + raise exception 'HypoPG extension is required for hypothetical index testing but is not installed'; + end if; + + -- Create hypothetical index + select * into v_hypo_result from hypopg_create_index(hypopg_index); + raise notice 'Created hypothetical index: % (oid: %)', + v_hypo_result.indexname, v_hypo_result.indexrelid; + end if; + + -- Build and execute explain query based on format + -- Output is preserved exactly as EXPLAIN returns it + begin + if lower(format) = 'json' then + v_explain_query := 'explain (verbose, settings, generic_plan, format json) ' || query; + execute v_explain_query into result; + else + v_explain_query := 'explain (verbose, settings, generic_plan) ' || query; + for v_line in execute v_explain_query loop + v_lines := array_append(v_lines, v_line."QUERY PLAN"); + end loop; + result := array_to_string(v_lines, e'\n'); + end if; + exception when others then + -- Clean up hypothetical index before re-raising + if hypopg_index is not null then + perform hypopg_reset(); + end if; + raise; + end; + + -- Clean up hypothetical index + if hypopg_index is not null then + perform hypopg_reset(); + end if; +end; +$$; + +comment on function postgres_ai.explain_generic(text, text, text) is + 'Returns generic EXPLAIN plan with optional HypoPG index testing (requires PG16+)'; + +-- Grant execute to the monitoring user +grant execute on function postgres_ai.explain_generic(text, text, text) to {{ROLE_IDENT}}; + +/* + * table_describe + * + * Collects comprehensive information about a table for LLM analysis. + * Returns a compact text format with: + * - Table metadata (type, size estimates) + * - Columns (name, type, nullable, default) + * - Indexes + * - Constraints (PK, FK, unique, check) + * - Maintenance stats (vacuum/analyze times) + * + * Usage: + * select postgres_ai.table_describe('public.users'); + * select postgres_ai.table_describe('my_table'); -- uses search_path + */ +create or replace function postgres_ai.table_describe( + in table_name text, + out result text +) +language plpgsql +security definer +set search_path = pg_catalog, public +as $$ +declare + v_oid oid; + v_schema text; + v_table text; + v_relkind char; + v_relpages int; + v_reltuples float; + v_lines text[] := '{}'; + v_line text; + v_rec record; + v_constraint_count int := 0; +begin + -- Resolve table name to OID (handles schema-qualified and search_path) + v_oid := table_name::regclass::oid; + + -- Get basic table info + select + n.nspname, + c.relname, + c.relkind, + c.relpages, + c.reltuples + into v_schema, v_table, v_relkind, v_relpages, v_reltuples + from pg_class c + join pg_namespace n on n.oid = c.relnamespace + where c.oid = v_oid; + + -- Validate object type - only tables, views, and materialized views are supported + if v_relkind not in ('r', 'p', 'v', 'm', 'f') then + raise exception 'table_describe does not support % (relkind=%)', + case v_relkind + when 'i' then 'indexes' + when 'I' then 'partitioned indexes' + when 'S' then 'sequences' + when 't' then 'TOAST tables' + when 'c' then 'composite types' + else format('objects of type "%s"', v_relkind) + end, + v_relkind; + end if; + + -- Header + v_lines := array_append(v_lines, format('Table: %I.%I', v_schema, v_table)); + v_lines := array_append(v_lines, format('Type: %s | relpages: %s | reltuples: %s', + case v_relkind + when 'r' then 'table' + when 'p' then 'partitioned table' + when 'v' then 'view' + when 'm' then 'materialized view' + when 'f' then 'foreign table' + end, + v_relpages, + case when v_reltuples < 0 then '-1' else v_reltuples::bigint::text end + )); + + -- Vacuum/analyze stats (only for tables and materialized views, not views) + if v_relkind in ('r', 'p', 'm', 'f') then + select + format('Vacuum: %s (auto: %s) | Analyze: %s (auto: %s)', + coalesce(to_char(last_vacuum at time zone 'UTC', 'YYYY-MM-DD HH24:MI:SS UTC'), 'never'), + coalesce(to_char(last_autovacuum at time zone 'UTC', 'YYYY-MM-DD HH24:MI:SS UTC'), 'never'), + coalesce(to_char(last_analyze at time zone 'UTC', 'YYYY-MM-DD HH24:MI:SS UTC'), 'never'), + coalesce(to_char(last_autoanalyze at time zone 'UTC', 'YYYY-MM-DD HH24:MI:SS UTC'), 'never') + ) + into v_line + from pg_stat_all_tables + where relid = v_oid; + + if v_line is not null then + v_lines := array_append(v_lines, v_line); + end if; + end if; + + v_lines := array_append(v_lines, ''); + + -- Columns + v_lines := array_append(v_lines, 'Columns:'); + for v_rec in + select + a.attname, + format_type(a.atttypid, a.atttypmod) as data_type, + a.attnotnull, + (select pg_get_expr(d.adbin, d.adrelid, true) + from pg_attrdef d + where d.adrelid = a.attrelid and d.adnum = a.attnum and a.atthasdef) as default_val, + a.attidentity, + a.attgenerated + from pg_attribute a + where a.attrelid = v_oid + and a.attnum > 0 + and not a.attisdropped + order by a.attnum + loop + v_line := format(' %s %s', v_rec.attname, v_rec.data_type); + + if v_rec.attnotnull then + v_line := v_line || ' NOT NULL'; + end if; + + if v_rec.attidentity = 'a' then + v_line := v_line || ' GENERATED ALWAYS AS IDENTITY'; + elsif v_rec.attidentity = 'd' then + v_line := v_line || ' GENERATED BY DEFAULT AS IDENTITY'; + elsif v_rec.attgenerated = 's' then + v_line := v_line || format(' GENERATED ALWAYS AS (%s) STORED', v_rec.default_val); + elsif v_rec.default_val is not null then + v_line := v_line || format(' DEFAULT %s', v_rec.default_val); + end if; + + v_lines := array_append(v_lines, v_line); + end loop; + + -- View definition (for views and materialized views) + if v_relkind in ('v', 'm') then + v_lines := array_append(v_lines, ''); + v_lines := array_append(v_lines, 'Definition:'); + v_line := pg_get_viewdef(v_oid, true); + if v_line is not null then + -- Indent the view definition + v_line := ' ' || replace(v_line, e'\n', e'\n '); + v_lines := array_append(v_lines, v_line); + end if; + end if; + + -- Indexes (tables, partitioned tables, and materialized views can have indexes) + if v_relkind in ('r', 'p', 'm') then + v_lines := array_append(v_lines, ''); + v_lines := array_append(v_lines, 'Indexes:'); + for v_rec in + select + i.relname as index_name, + pg_get_indexdef(i.oid) as index_def, + ix.indisprimary, + ix.indisunique + from pg_index ix + join pg_class i on i.oid = ix.indexrelid + where ix.indrelid = v_oid + order by ix.indisprimary desc, ix.indisunique desc, i.relname + loop + v_line := ' '; + if v_rec.indisprimary then + v_line := v_line || 'PRIMARY KEY: '; + elsif v_rec.indisunique then + v_line := v_line || 'UNIQUE: '; + else + v_line := v_line || 'INDEX: '; + end if; + -- Extract just the column part from index definition + v_line := v_line || v_rec.index_name || ' ' || + regexp_replace(v_rec.index_def, '^CREATE.*INDEX.*ON.*USING\s+\w+\s*', ''); + v_lines := array_append(v_lines, v_line); + end loop; + + if not exists (select 1 from pg_index where indrelid = v_oid) then + v_lines := array_append(v_lines, ' (none)'); + end if; + end if; + + -- Constraints (only tables can have constraints) + if v_relkind in ('r', 'p', 'f') then + v_lines := array_append(v_lines, ''); + v_lines := array_append(v_lines, 'Constraints:'); + v_constraint_count := 0; + + for v_rec in + select + conname, + contype, + pg_get_constraintdef(oid, true) as condef + from pg_constraint + where conrelid = v_oid + and contype != 'p' -- skip primary key (shown with indexes) + order by + case contype when 'f' then 1 when 'u' then 2 when 'c' then 3 else 4 end, + conname + loop + v_constraint_count := v_constraint_count + 1; + v_line := ' '; + case v_rec.contype + when 'f' then v_line := v_line || 'FK: '; + when 'u' then v_line := v_line || 'UNIQUE: '; + when 'c' then v_line := v_line || 'CHECK: '; + else v_line := v_line || v_rec.contype || ': '; + end case; + v_line := v_line || v_rec.conname || ' ' || v_rec.condef; + v_lines := array_append(v_lines, v_line); + end loop; + + if v_constraint_count = 0 then + v_lines := array_append(v_lines, ' (none)'); + end if; + + -- Foreign keys referencing this table + v_lines := array_append(v_lines, ''); + v_lines := array_append(v_lines, 'Referenced by:'); + v_constraint_count := 0; + + for v_rec in + select + conname, + conrelid::regclass::text as from_table, + pg_get_constraintdef(oid, true) as condef + from pg_constraint + where confrelid = v_oid + and contype = 'f' + order by conrelid::regclass::text, conname + loop + v_constraint_count := v_constraint_count + 1; + v_lines := array_append(v_lines, format(' %s.%s %s', + v_rec.from_table, v_rec.conname, v_rec.condef)); + end loop; + + if v_constraint_count = 0 then + v_lines := array_append(v_lines, ' (none)'); + end if; + end if; + + -- Partition info (if partitioned table or partition) + if v_relkind = 'p' then + -- This is a partitioned table - show partition key and partitions + v_lines := array_append(v_lines, ''); + v_lines := array_append(v_lines, 'Partitioning:'); + + select format(' %s BY %s', + case partstrat + when 'r' then 'RANGE' + when 'l' then 'LIST' + when 'h' then 'HASH' + else partstrat + end, + pg_get_partkeydef(v_oid) + ) + into v_line + from pg_partitioned_table + where partrelid = v_oid; + + if v_line is not null then + v_lines := array_append(v_lines, v_line); + end if; + + -- List partitions + v_constraint_count := 0; + for v_rec in + select + c.oid::regclass::text as partition_name, + pg_get_expr(c.relpartbound, c.oid) as partition_bound, + c.relpages, + c.reltuples + from pg_inherits i + join pg_class c on c.oid = i.inhrelid + where i.inhparent = v_oid + order by c.oid::regclass::text + loop + v_constraint_count := v_constraint_count + 1; + v_lines := array_append(v_lines, format(' %s: %s (relpages: %s, reltuples: %s)', + v_rec.partition_name, v_rec.partition_bound, + v_rec.relpages, + case when v_rec.reltuples < 0 then '-1' else v_rec.reltuples::bigint::text end + )); + end loop; + + v_lines := array_append(v_lines, format(' Total partitions: %s', v_constraint_count)); + + elsif exists (select 1 from pg_inherits where inhrelid = v_oid) then + -- This is a partition - show parent and bound + v_lines := array_append(v_lines, ''); + v_lines := array_append(v_lines, 'Partition of:'); + + select format(' %s FOR VALUES %s', + i.inhparent::regclass::text, + pg_get_expr(c.relpartbound, c.oid) + ) + into v_line + from pg_inherits i + join pg_class c on c.oid = i.inhrelid + where i.inhrelid = v_oid; + + if v_line is not null then + v_lines := array_append(v_lines, v_line); + end if; + end if; + + result := array_to_string(v_lines, e'\n'); +end; +$$; + +comment on function postgres_ai.table_describe(text) is + 'Returns comprehensive table information in compact text format for LLM analysis'; + +grant execute on function postgres_ai.table_describe(text) to {{ROLE_IDENT}}; + + diff --git a/cli/test/init.integration.test.cjs b/cli/test/init.integration.test.cjs index 25f4f66..93d984c 100644 --- a/cli/test/init.integration.test.cjs +++ b/cli/test/init.integration.test.cjs @@ -268,12 +268,21 @@ test( ); 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" + "select has_table_privilege('postgres_ai_mon', 'postgres_ai.pg_statistic', 'SELECT') as ok" ); assert.equal(viewOk.rows[0].ok, true); + const explainFnOk = await c.query( + "select has_function_privilege('postgres_ai_mon', 'postgres_ai.explain_generic(text, text, text)', 'EXECUTE') as ok" + ); + assert.equal(explainFnOk.rows[0].ok, true); + const tableDescribeFnOk = await c.query( + "select has_function_privilege('postgres_ai_mon', 'postgres_ai.table_describe(text)', 'EXECUTE') as ok" + ); + assert.equal(tableDescribeFnOk.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="))); + assert.ok(sp.rows[0].rolconfig.some((v) => String(v).includes("postgres_ai"))); await c.end(); } @@ -379,4 +388,148 @@ test("integration: prepare-db --reset-password updates the monitoring role login } }); +test("integration: table_describe works with different object types", { skip: !havePostgresBinaries() }, async (t) => { + const pg = await withTempPostgres(t); + const { Client } = require("pg"); + + // Run init first + { + const r = await runCliInit([pg.adminUri, "--password", "pw1", "--skip-optional-permissions"]); + assert.equal(r.status, 0, r.stderr || r.stdout); + } + + const c = new Client({ connectionString: pg.adminUri }); + await c.connect(); + + // Create test objects + await c.query(` + -- Regular table with various features + create table test_table ( + id serial primary key, + name text not null, + email text unique, + status text default 'active' check (status in ('active', 'inactive')), + created_at timestamptz default now() + ); + create index test_table_name_idx on test_table(name); + + -- Another table with FK + create table test_child ( + id serial primary key, + parent_id int references test_table(id) + ); + + -- Partitioned table + create table test_partitioned ( + id serial, + created_at date not null, + data text + ) partition by range (created_at); + + create table test_partitioned_2024 partition of test_partitioned + for values from ('2024-01-01') to ('2025-01-01'); + create table test_partitioned_2025 partition of test_partitioned + for values from ('2025-01-01') to ('2026-01-01'); + + -- View + create view test_view as select id, name from test_table where status = 'active'; + + -- Materialized view + create materialized view test_matview as select id, name from test_table; + create unique index test_matview_id_idx on test_matview(id); + + -- Sequence (for error test) + create sequence test_seq; + `); + + // Test regular table + { + const res = await c.query("select postgres_ai.table_describe('test_table')"); + const output = res.rows[0].result; + assert.match(output, /Table: "public"\."test_table"/); + assert.match(output, /Type: table/); + assert.match(output, /relpages:/); + assert.match(output, /reltuples:/); + assert.match(output, /Columns:/); + assert.match(output, /id integer NOT NULL/); + assert.match(output, /name text NOT NULL/); + assert.match(output, /email text/); + assert.match(output, /status text.*DEFAULT/); + assert.match(output, /Indexes:/); + assert.match(output, /PRIMARY KEY:/); + assert.match(output, /UNIQUE:/); + assert.match(output, /INDEX:.*test_table_name_idx/); + assert.match(output, /Constraints:/); + assert.match(output, /CHECK:/); + assert.match(output, /Referenced by:/); + assert.match(output, /test_child/); + } + + // Test partitioned table + { + const res = await c.query("select postgres_ai.table_describe('test_partitioned')"); + const output = res.rows[0].result; + assert.match(output, /Type: partitioned table/); + assert.match(output, /Partitioning:/); + assert.match(output, /RANGE BY/); + assert.match(output, /test_partitioned_2024/); + assert.match(output, /test_partitioned_2025/); + assert.match(output, /Total partitions: 2/); + } + + // Test partition + { + const res = await c.query("select postgres_ai.table_describe('test_partitioned_2024')"); + const output = res.rows[0].result; + assert.match(output, /Type: table/); + assert.match(output, /Partition of:/); + assert.match(output, /test_partitioned/); + assert.match(output, /FOR VALUES/); + } + + // Test view + { + const res = await c.query("select postgres_ai.table_describe('test_view')"); + const output = res.rows[0].result; + assert.match(output, /Type: view/); + assert.match(output, /Columns:/); + assert.match(output, /Definition:/); + assert.match(output, /SELECT.*FROM.*test_table/i); + // Views should NOT have Indexes or Constraints sections + assert.doesNotMatch(output, /^Indexes:/m); + assert.doesNotMatch(output, /^Constraints:/m); + } + + // Test materialized view + { + const res = await c.query("select postgres_ai.table_describe('test_matview')"); + const output = res.rows[0].result; + assert.match(output, /Type: materialized view/); + assert.match(output, /Columns:/); + assert.match(output, /Definition:/); + assert.match(output, /Indexes:/); + assert.match(output, /UNIQUE:.*test_matview_id_idx/); + // Mat views should NOT have Constraints section + assert.doesNotMatch(output, /^Constraints:/m); + } + + // Test sequence (should error) + { + await assert.rejects( + c.query("select postgres_ai.table_describe('test_seq')"), + /table_describe does not support sequences/ + ); + } + + // Test index (should error) + { + await assert.rejects( + c.query("select postgres_ai.table_describe('test_table_name_idx')"), + /table_describe does not support indexes/ + ); + } + + await c.end(); +}); + diff --git a/cli/test/init.test.cjs b/cli/test/init.test.cjs index b14e131..62fa03c 100644 --- a/cli/test/init.test.cjs +++ b/cli/test/init.test.cjs @@ -149,6 +149,23 @@ test("buildInitPlan includes optional steps when enabled", async () => { assert.ok(plan.steps.some((s) => s.optional)); }); +test("buildInitPlan includes helpers step with explain_generic and table_describe functions", async () => { + const plan = await init.buildInitPlan({ + database: "mydb", + monitoringUser: DEFAULT_MONITORING_USER, + monitoringPassword: "pw", + includeOptionalPermissions: false, + }); + + const helpersStep = plan.steps.find((s) => s.name === "05.helpers"); + assert.ok(helpersStep, "05.helpers step should exist"); + assert.match(helpersStep.sql, /postgres_ai\.explain_generic/i); + assert.match(helpersStep.sql, /postgres_ai\.table_describe/i); + assert.match(helpersStep.sql, /security definer/i); + assert.match(helpersStep.sql, /grant execute on function.*explain_generic/i); + assert.match(helpersStep.sql, /grant execute on function.*table_describe/i); +}); + test("resolveAdminConnection accepts positional URI", () => { const r = init.resolveAdminConnection({ conn: "postgresql://u:p@h:5432/d" }); assert.ok(r.clientConfig.connectionString); @@ -236,9 +253,43 @@ test("applyInitPlan preserves Postgres error fields on step failures", async () } ); + // Each step gets its own begin/commit (or rollback on failure) assert.deepEqual(calls, ["begin;", "select 1", "rollback;"]); }); +test("applyInitPlan wraps each step in its own transaction", async () => { + const plan = { + monitoringUser: DEFAULT_MONITORING_USER, + database: "mydb", + steps: [ + { name: "01.role", sql: "select 1" }, + { name: "02.permissions", sql: "select 2" }, + ], + }; + + const calls = []; + const client = { + query: async (sql) => { + calls.push(sql); + if (sql === "begin;") return { rowCount: 1 }; + if (sql === "commit;") return { rowCount: 1 }; + if (sql === "select 1") return { rowCount: 1 }; + if (sql === "select 2") return { rowCount: 1 }; + throw new Error(`unexpected sql: ${sql}`); + }, + }; + + const result = await init.applyInitPlan({ client, plan }); + assert.deepEqual(result.applied, ["01.role", "02.permissions"]); + assert.deepEqual(result.skippedOptional, []); + + // Each step should have its own begin/commit + assert.deepEqual(calls, [ + "begin;", "select 1", "commit;", + "begin;", "select 2", "commit;", + ]); +}); + test("verifyInitSetup runs inside a repeatable read snapshot and rolls back", async () => { const calls = []; const client = { @@ -252,7 +303,7 @@ test("verifyInitSetup runs inside a repeatable read snapshot and rolls back", as return { rowCount: 1, rows: [] }; } if (String(sql).includes("select rolconfig")) { - return { rowCount: 1, rows: [{ rolconfig: ['search_path="$user", public, pg_catalog'] }] }; + return { rowCount: 1, rows: [{ rolconfig: ['search_path=postgres_ai, "$user", public, pg_catalog'] }] }; } if (String(sql).includes("from pg_catalog.pg_roles")) { return { rowCount: 1, rows: [] }; @@ -266,15 +317,21 @@ test("verifyInitSetup runs inside a repeatable read snapshot and rolls back", as if (String(sql).includes("has_table_privilege") && String(sql).includes("pg_catalog.pg_index")) { return { rowCount: 1, rows: [{ ok: true }] }; } - if (String(sql).includes("to_regclass('public.pg_statistic')")) { + if (String(sql).includes("to_regclass('postgres_ai.pg_statistic')")) { return { rowCount: 1, rows: [{ ok: true }] }; } - if (String(sql).includes("has_table_privilege") && String(sql).includes("public.pg_statistic")) { + if (String(sql).includes("has_table_privilege") && String(sql).includes("postgres_ai.pg_statistic")) { return { rowCount: 1, rows: [{ ok: true }] }; } if (String(sql).includes("has_schema_privilege")) { return { rowCount: 1, rows: [{ ok: true }] }; } + if (String(sql).includes("has_function_privilege") && String(sql).includes("postgres_ai.explain_generic")) { + return { rowCount: 1, rows: [{ ok: true }] }; + } + if (String(sql).includes("has_function_privilege") && String(sql).includes("postgres_ai.table_describe")) { + return { rowCount: 1, rows: [{ ok: true }] }; + } throw new Error(`unexpected sql: ${sql} params=${JSON.stringify(params)}`); },