Skip to content

Commit e3c1471

Browse files
claudeclaude
committed
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 - 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 function Relates to: https://gitlab.com/postgres-ai/postgres_ai/-/issues/68
1 parent 08f1e7d commit e3c1471

4 files changed

Lines changed: 257 additions & 76 deletions

File tree

‎cli/lib/init.ts‎

Lines changed: 84 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -485,6 +485,12 @@ end $$;`;
485485
sql: applyTemplate(loadSqlTemplate("02.permissions.sql"), vars),
486486
});
487487

488+
// Helper functions (SECURITY DEFINER) for plan analysis and table info
489+
steps.push({
490+
name: "05.helpers",
491+
sql: applyTemplate(loadSqlTemplate("05.helpers.sql"), vars),
492+
});
493+
488494
if (params.includeOptionalPermissions) {
489495
steps.push(
490496
{
@@ -511,78 +517,70 @@ export async function applyInitPlan(params: {
511517
const applied: string[] = [];
512518
const skippedOptional: string[] = [];
513519

514-
// Apply non-optional steps in a single transaction.
515-
await params.client.query("begin;");
516-
try {
517-
for (const step of params.plan.steps.filter((s) => !s.optional)) {
520+
// Helper to wrap a step execution in begin/commit
521+
const executeStep = async (step: InitStep): Promise<void> => {
522+
await params.client.query("begin;");
523+
try {
524+
await params.client.query(step.sql, step.params as any);
525+
await params.client.query("commit;");
526+
} catch (e) {
527+
// Rollback errors should never mask the original failure.
518528
try {
519-
await params.client.query(step.sql, step.params as any);
520-
applied.push(step.name);
521-
} catch (e) {
522-
const msg = e instanceof Error ? e.message : String(e);
523-
const errAny = e as any;
524-
const wrapped: any = new Error(`Failed at step "${step.name}": ${msg}`);
525-
// Preserve useful Postgres error fields so callers can provide better hints / diagnostics.
526-
const pgErrorFields = [
527-
"code",
528-
"detail",
529-
"hint",
530-
"position",
531-
"internalPosition",
532-
"internalQuery",
533-
"where",
534-
"schema",
535-
"table",
536-
"column",
537-
"dataType",
538-
"constraint",
539-
"file",
540-
"line",
541-
"routine",
542-
] as const;
543-
if (errAny && typeof errAny === "object") {
544-
for (const field of pgErrorFields) {
545-
if (errAny[field] !== undefined) wrapped[field] = errAny[field];
546-
}
547-
}
548-
if (e instanceof Error && e.stack) {
549-
wrapped.stack = e.stack;
550-
}
551-
throw wrapped;
529+
await params.client.query("rollback;");
530+
} catch {
531+
// ignore
552532
}
533+
throw e;
553534
}
554-
await params.client.query("commit;");
555-
} catch (e) {
556-
// Rollback errors should never mask the original failure.
535+
};
536+
537+
// Apply non-optional steps, each in its own transaction
538+
for (const step of params.plan.steps.filter((s) => !s.optional)) {
557539
try {
558-
await params.client.query("rollback;");
559-
} catch {
560-
// ignore
540+
await executeStep(step);
541+
applied.push(step.name);
542+
} catch (e) {
543+
const msg = e instanceof Error ? e.message : String(e);
544+
const errAny = e as any;
545+
const wrapped: any = new Error(`Failed at step "${step.name}": ${msg}`);
546+
// Preserve useful Postgres error fields so callers can provide better hints / diagnostics.
547+
const pgErrorFields = [
548+
"code",
549+
"detail",
550+
"hint",
551+
"position",
552+
"internalPosition",
553+
"internalQuery",
554+
"where",
555+
"schema",
556+
"table",
557+
"column",
558+
"dataType",
559+
"constraint",
560+
"file",
561+
"line",
562+
"routine",
563+
] as const;
564+
if (errAny && typeof errAny === "object") {
565+
for (const field of pgErrorFields) {
566+
if (errAny[field] !== undefined) wrapped[field] = errAny[field];
567+
}
568+
}
569+
if (e instanceof Error && e.stack) {
570+
wrapped.stack = e.stack;
571+
}
572+
throw wrapped;
561573
}
562-
throw e;
563574
}
564575

565-
// Apply optional steps outside of the transaction so a failure doesn't abort everything.
576+
// Apply optional steps, each in its own transaction (failure doesn't abort)
566577
for (const step of params.plan.steps.filter((s) => s.optional)) {
567578
try {
568-
// Run each optional step in its own mini-transaction to avoid partial application.
569-
await params.client.query("begin;");
570-
try {
571-
await params.client.query(step.sql, step.params as any);
572-
await params.client.query("commit;");
573-
applied.push(step.name);
574-
} catch {
575-
try {
576-
await params.client.query("rollback;");
577-
} catch {
578-
// ignore rollback errors
579-
}
580-
skippedOptional.push(step.name);
581-
// best-effort: ignore
582-
}
579+
await executeStep(step);
580+
applied.push(step.name);
583581
} catch {
584-
// If we can't even begin/commit, treat as skipped.
585582
skippedOptional.push(step.name);
583+
// best-effort: ignore
586584
}
587585
}
588586

@@ -642,16 +640,25 @@ export async function verifyInitSetup(params: {
642640
missingRequired.push("SELECT on pg_catalog.pg_index");
643641
}
644642

645-
const viewExistsRes = await params.client.query("select to_regclass('public.pg_statistic') is not null as ok");
643+
// Check postgres_ai schema exists and is usable
644+
const schemaExistsRes = await params.client.query(
645+
"select has_schema_privilege($1, 'postgres_ai', 'USAGE') as ok",
646+
[role]
647+
);
648+
if (!schemaExistsRes.rows?.[0]?.ok) {
649+
missingRequired.push("USAGE on schema postgres_ai");
650+
}
651+
652+
const viewExistsRes = await params.client.query("select to_regclass('postgres_ai.pg_statistic') is not null as ok");
646653
if (!viewExistsRes.rows?.[0]?.ok) {
647-
missingRequired.push("view public.pg_statistic exists");
654+
missingRequired.push("view postgres_ai.pg_statistic exists");
648655
} else {
649656
const viewPrivRes = await params.client.query(
650-
"select has_table_privilege($1, 'public.pg_statistic', 'SELECT') as ok",
657+
"select has_table_privilege($1, 'postgres_ai.pg_statistic', 'SELECT') as ok",
651658
[role]
652659
);
653660
if (!viewPrivRes.rows?.[0]?.ok) {
654-
missingRequired.push("SELECT on view public.pg_statistic");
661+
missingRequired.push("SELECT on view postgres_ai.pg_statistic");
655662
}
656663
}
657664

@@ -669,13 +676,22 @@ export async function verifyInitSetup(params: {
669676
if (typeof spLine !== "string" || !spLine) {
670677
missingRequired.push("role search_path is set");
671678
} else {
672-
// We accept any ordering as long as public and pg_catalog are included.
679+
// We accept any ordering as long as postgres_ai, public, and pg_catalog are included.
673680
const sp = spLine.toLowerCase();
674-
if (!sp.includes("public") || !sp.includes("pg_catalog")) {
675-
missingRequired.push("role search_path includes public and pg_catalog");
681+
if (!sp.includes("postgres_ai") || !sp.includes("public") || !sp.includes("pg_catalog")) {
682+
missingRequired.push("role search_path includes postgres_ai, public and pg_catalog");
676683
}
677684
}
678685

686+
// Check for explain_generic helper function
687+
const explainFnRes = await params.client.query(
688+
"select has_function_privilege($1, 'postgres_ai.explain_generic(text, text, text)', 'EXECUTE') as ok",
689+
[role]
690+
);
691+
if (!explainFnRes.rows?.[0]?.ok) {
692+
missingRequired.push("EXECUTE on postgres_ai.explain_generic(text, text, text)");
693+
}
694+
679695
if (params.includeOptionalPermissions) {
680696
// Optional RDS/Aurora extras
681697
{

‎cli/sql/02.permissions.sql‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,12 @@ grant connect on database {{DB_IDENT}} to {{ROLE_IDENT}};
77
grant pg_monitor to {{ROLE_IDENT}};
88
grant select on pg_catalog.pg_index to {{ROLE_IDENT}};
99

10-
-- Optional, for bloat analysis: expose pg_statistic via a view
11-
create or replace view public.pg_statistic as
10+
-- Create postgres_ai schema for our objects
11+
create schema if not exists postgres_ai;
12+
grant usage on schema postgres_ai to {{ROLE_IDENT}};
13+
14+
-- For bloat analysis: expose pg_statistic via a view
15+
create or replace view postgres_ai.pg_statistic as
1216
select
1317
n.nspname as schemaname,
1418
c.relname as tablename,
@@ -22,12 +26,12 @@ join pg_catalog.pg_namespace n on n.oid = c.relnamespace
2226
join pg_catalog.pg_attribute a on a.attrelid = s.starelid and a.attnum = s.staattnum
2327
where a.attnum > 0 and not a.attisdropped;
2428

25-
grant select on public.pg_statistic to {{ROLE_IDENT}};
29+
grant select on postgres_ai.pg_statistic to {{ROLE_IDENT}};
2630

2731
-- Hardened clusters sometimes revoke PUBLIC on schema public
2832
grant usage on schema public to {{ROLE_IDENT}};
2933

30-
-- Keep search_path predictable
31-
alter user {{ROLE_IDENT}} set search_path = "$user", public, pg_catalog;
34+
-- Keep search_path predictable; postgres_ai first so our objects are found
35+
alter user {{ROLE_IDENT}} set search_path = postgres_ai, "$user", public, pg_catalog;
3236

3337

‎cli/sql/05.helpers.sql‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
-- Helper functions for postgres_ai monitoring user (template-filled by cli/lib/init.ts)
2+
-- These functions use SECURITY DEFINER to allow the monitoring user to perform
3+
-- operations they don't have direct permissions for.
4+
5+
/*
6+
* pgai_explain_generic
7+
*
8+
* Function to get generic explain plans with optional HypoPG index testing.
9+
* Requires: PostgreSQL 16+ (for generic_plan option), HypoPG extension (optional).
10+
*
11+
* Usage examples:
12+
* -- Basic generic plan
13+
* select postgres_ai.explain_generic('select * from users where id = $1');
14+
*
15+
* -- JSON format
16+
* select postgres_ai.explain_generic('select * from users where id = $1', 'json');
17+
*
18+
* -- Test a hypothetical index
19+
* select postgres_ai.explain_generic(
20+
* 'select * from users where email = $1',
21+
* 'text',
22+
* 'create index on users (email)'
23+
* );
24+
*/
25+
create or replace function postgres_ai.explain_generic(
26+
in query text,
27+
in format text default 'text',
28+
in hypopg_index text default null,
29+
out result text
30+
)
31+
language plpgsql
32+
security definer
33+
set search_path = pg_catalog, public
34+
as $$
35+
declare
36+
v_line record;
37+
v_lines text[] := '{}';
38+
v_explain_query text;
39+
v_hypo_result record;
40+
v_version int;
41+
v_hypopg_available boolean;
42+
begin
43+
-- Check PostgreSQL version (generic_plan requires 16+)
44+
select current_setting('server_version_num')::int into v_version;
45+
46+
if v_version < 160000 then
47+
raise exception 'generic_plan requires PostgreSQL 16+, current version: %',
48+
current_setting('server_version');
49+
end if;
50+
51+
-- Check if HypoPG extension is available
52+
if hypopg_index is not null then
53+
select exists(
54+
select 1 from pg_extension where extname = 'hypopg'
55+
) into v_hypopg_available;
56+
57+
if not v_hypopg_available then
58+
raise exception 'HypoPG extension is required for hypothetical index testing but is not installed';
59+
end if;
60+
61+
-- Create hypothetical index
62+
select * into v_hypo_result from hypopg_create_index(hypopg_index);
63+
raise notice 'Created hypothetical index: % (oid: %)',
64+
v_hypo_result.indexname, v_hypo_result.indexrelid;
65+
end if;
66+
67+
-- Build and execute explain query based on format
68+
begin
69+
if lower(format) = 'json' then
70+
v_explain_query := 'explain (verbose, settings, generic_plan, format json) ' || query;
71+
execute v_explain_query into result;
72+
else
73+
v_explain_query := 'explain (verbose, settings, generic_plan) ' || query;
74+
75+
for v_line in execute v_explain_query loop
76+
v_lines := array_append(v_lines, v_line."QUERY PLAN");
77+
end loop;
78+
79+
result := array_to_string(v_lines, e'\n');
80+
end if;
81+
exception when others then
82+
-- Clean up hypothetical index before re-raising
83+
if hypopg_index is not null then
84+
perform hypopg_reset();
85+
end if;
86+
raise;
87+
end;
88+
89+
-- Clean up hypothetical index
90+
if hypopg_index is not null then
91+
perform hypopg_reset();
92+
end if;
93+
end;
94+
$$;
95+
96+
comment on function postgres_ai.explain_generic(text, text, text) is
97+
'Returns generic EXPLAIN plan with optional HypoPG index testing (requires PG16+)';
98+
99+
-- Grant execute to the monitoring user
100+
grant execute on function postgres_ai.explain_generic(text, text, text) to {{ROLE_IDENT}};
101+
102+
-- Placeholder for table info collection function (to be implemented separately)
103+
-- This will collect information about a specified table including:
104+
-- - Table structure (columns, types, constraints)
105+
-- - Index information
106+
-- - Statistics
107+
-- - Size information
108+
109+

0 commit comments

Comments
 (0)