From 7e390aa6840807bd66339e2309544df02922915e Mon Sep 17 00:00:00 2001 From: tianzhou Date: Thu, 20 Aug 2026 21:47:15 -0700 Subject: [PATCH 1/8] fix: stub cross-schema partition parents in plan database (#552) When a partition child is in a different schema from its partitioned parent, the plan database fails because it tries to create the partition on the real parent table. Fix by detecting PARTITION OF references to cross-schema parents and prepending stub CREATE TABLE ... PARTITION BY statements, mirroring the existing FK stub pattern. Co-Authored-By: Claude Opus 4.6 --- cmd/plan/partition_stubs.go | 66 +++++++++++++++++++ cmd/plan/plan.go | 18 ++++++ internal/postgres/fk_refs.go | 41 ++++++++++++ internal/postgres/fk_refs_test.go | 49 ++++++++++++++ ir/stub.go | 104 ++++++++++++++++++++++++++++++ 5 files changed, 278 insertions(+) create mode 100644 cmd/plan/partition_stubs.go diff --git a/cmd/plan/partition_stubs.go b/cmd/plan/partition_stubs.go new file mode 100644 index 00000000..1f4d025a --- /dev/null +++ b/cmd/plan/partition_stubs.go @@ -0,0 +1,66 @@ +package plan + +import ( + "context" + "strings" + + "github.com/pgplex/pgschema/internal/postgres" + "github.com/pgplex/pgschema/ir" + "github.com/pgplex/pgschema/cmd/util" + "github.com/pgplex/pgschema/internal/logger" +) + +// prependPartitionParentStubs creates stub CREATE TABLE ... PARTITION BY +// statements for cross-schema partition parents referenced by PARTITION OF +// in the desired SQL. Without these stubs the plan database rejects the +// desired-state SQL because the parent table does not exist. +func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig, targetSchema, desiredSQL string) (string, error) { + refs := postgres.ExtractPartitionOfTargets(desiredSQL, targetSchema) + if len(refs) == 0 { + return desiredSQL, nil + } + + created := make(map[string]bool) + for _, name := range postgres.ExtractCreateTableNames(desiredSQL, targetSchema) { + created[name.Schema+"."+name.Table] = true + } + + var toStub []postgres.QualifiedName + seen := make(map[string]bool) + for _, ref := range refs { + if ref.Schema == targetSchema { + continue + } + key := ref.Schema + "." + ref.Table + if created[key] || seen[key] { + continue + } + seen[key] = true + toStub = append(toStub, ref) + } + if len(toStub) == 0 { + return desiredSQL, nil + } + + conn, err := util.Connect(cfg) + if err != nil { + return "", err + } + defer conn.Close() + + var stubs strings.Builder + for _, ref := range toStub { + ddl, err := ir.BuildPartitionedTableStubSQL(ctx, conn, ref.Schema, ref.Table, targetSchema) + if err != nil { + return "", err + } + if ddl == "" { + continue + } + logger.Get().Debug("prepending stub for cross-schema partition parent", + "schema", ref.Schema, "table", ref.Table) + stubs.WriteString(ddl) + } + + return stubs.String() + desiredSQL, nil +} diff --git a/cmd/plan/plan.go b/cmd/plan/plan.go index 93e6c330..84da06a1 100644 --- a/cmd/plan/plan.go +++ b/cmd/plan/plan.go @@ -316,6 +316,24 @@ func GeneratePlan(config *PlanConfig, provider postgres.DesiredStateProvider) (* } } + // Stub cross-schema partition parents so PARTITION OF references resolve + // in the plan database (issue #552). + { + connCfg := &util.ConnectionConfig{ + Host: config.Host, + Port: config.Port, + Database: config.DB, + User: config.User, + Password: config.Password, + SSLMode: config.SSLMode, + ApplicationName: config.ApplicationName, + } + desiredState, err = prependPartitionParentStubs(ctx, connCfg, config.Schema, desiredState) + if err != nil { + return nil, fmt.Errorf("failed to stub cross-schema partition parents: %w", err) + } + } + // Apply desired state SQL to the provider (embedded postgres or external database) if err := provider.ApplySchema(ctx, config.Schema, desiredState); err != nil { return nil, fmt.Errorf("failed to apply desired state: %w", err) diff --git a/internal/postgres/fk_refs.go b/internal/postgres/fk_refs.go index 215f57bc..5aac59b8 100644 --- a/internal/postgres/fk_refs.go +++ b/internal/postgres/fk_refs.go @@ -134,6 +134,47 @@ func ExtractCreateTableNames(sql, defaultSchema string) []QualifiedName { return out } +// ExtractPartitionOfTargets returns schema-qualified table names that appear +// as PARTITION OF targets in sql. Unqualified names use defaultSchema. +// String literals, comments, and dollar-quoted bodies are skipped. +func ExtractPartitionOfTargets(sql, defaultSchema string) []QualifiedName { + seen := make(map[string]bool) + var out []QualifiedName + + walkSQLCode(sql, func(code string) { + i := 0 + for i < len(code) { + idx := indexKeyword(code, i, "partition") + if idx < 0 { + return + } + i = idx + len("partition") + i = skipSpace(code, i) + if !hasKeywordAt(code, i, "of") { + continue + } + i += len("of") + schema, table, next, ok := parseQualifiedName(code, i) + if !ok { + i = next + continue + } + i = next + if schema == "" { + schema = defaultSchema + } + key := schema + "." + table + if seen[key] { + continue + } + seen[key] = true + out = append(out, QualifiedName{Schema: schema, Table: table}) + } + }) + + return out +} + func walkSQLCode(sql string, fn func(code string)) { for _, seg := range splitDollarQuotedSegments(sql) { if seg.quoted { diff --git a/internal/postgres/fk_refs_test.go b/internal/postgres/fk_refs_test.go index afbe35af..7bb27cd5 100644 --- a/internal/postgres/fk_refs_test.go +++ b/internal/postgres/fk_refs_test.go @@ -126,3 +126,52 @@ func TestExtractCreateTableNames(t *testing.T) { }) } } + +func TestExtractPartitionOfTargets(t *testing.T) { + tests := []struct { + name string + sql string + defaultSchema string + want []QualifiedName + }{ + { + name: "simple partition of", + sql: "CREATE TABLE child PARTITION OF parent FOR VALUES IN ('a');", + defaultSchema: "public", + want: []QualifiedName{{Schema: "public", Table: "parent"}}, + }, + { + name: "qualified partition of", + sql: "CREATE TABLE data.child PARTITION OF core.parent FOR VALUES FROM (1) TO (100);", + defaultSchema: "data", + want: []QualifiedName{{Schema: "core", Table: "parent"}}, + }, + { + name: "multiple partitions same parent deduped", + sql: "CREATE TABLE c1 PARTITION OF parent FOR VALUES IN ('a');\nCREATE TABLE c2 PARTITION OF parent FOR VALUES IN ('b');", + defaultSchema: "public", + want: []QualifiedName{{Schema: "public", Table: "parent"}}, + }, + { + name: "partition by not matched", + sql: "CREATE TABLE parent (id int) PARTITION BY RANGE (id);", + defaultSchema: "public", + want: nil, + }, + { + name: "quoted identifiers", + sql: `CREATE TABLE "Child" PARTITION OF "Parent" FOR VALUES IN (1);`, + defaultSchema: "public", + want: []QualifiedName{{Schema: "public", Table: "Parent"}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := ExtractPartitionOfTargets(tt.sql, tt.defaultSchema) + if !reflect.DeepEqual(got, tt.want) { + t.Errorf("ExtractPartitionOfTargets() = %#v, want %#v", got, tt.want) + } + }) + } +} diff --git a/ir/stub.go b/ir/stub.go index b793a262..a3adee0c 100644 --- a/ir/stub.go +++ b/ir/stub.go @@ -71,6 +71,110 @@ func BuildTableStubSQL(ctx context.Context, db *sql.DB, schema, table, targetSch return b.String(), nil } +// BuildPartitionedTableStubSQL returns CREATE SCHEMA / CREATE TABLE DDL for a +// partitioned parent table that is referenced by PARTITION OF from a child in +// another schema. The stub includes the PARTITION BY clause so the child can +// attach as a partition. +// +// Returns an empty string if the table does not exist or is not partitioned. +func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table, targetSchema string) (string, error) { + cols, err := queryStubColumns(ctx, db, schema, table) + if err != nil { + return "", err + } + if len(cols) == 0 { + return "", nil + } + + partStrategy, partKey, err := queryPartitionInfo(ctx, db, schema, table) + if err != nil { + return "", err + } + if partStrategy == "" { + return "", nil + } + + constraints, err := queryStubConstraints(ctx, db, schema, table) + if err != nil { + return "", err + } + + var b strings.Builder + qualified := QualifyEntityNameWithQuotesMode(schema, table, targetSchema, schema != targetSchema) + + if schema != targetSchema { + b.WriteString("CREATE SCHEMA IF NOT EXISTS ") + b.WriteString(QuoteIdentifier(schema)) + b.WriteString(";\n") + } + + b.WriteString("-- pgschema: stub for cross-schema partition parent ") + b.WriteString(sanitizeComment(schema)) + b.WriteString(".") + b.WriteString(sanitizeComment(table)) + b.WriteString("\nCREATE TABLE IF NOT EXISTS ") + b.WriteString(qualified) + b.WriteString(" (\n") + + for i, col := range cols { + b.WriteString(" ") + b.WriteString(QuoteIdentifier(col.name)) + b.WriteString(" ") + b.WriteString(col.dataType) + if col.notNull { + b.WriteString(" NOT NULL") + } + if i < len(cols)-1 || len(constraints) > 0 { + b.WriteString(",") + } + b.WriteString("\n") + } + + for i, def := range constraints { + b.WriteString(" ") + b.WriteString(def) + if i < len(constraints)-1 { + b.WriteString(",") + } + b.WriteString("\n") + } + + b.WriteString(") PARTITION BY ") + b.WriteString(partStrategy) + b.WriteString(" (") + b.WriteString(partKey) + b.WriteString(");\n") + return b.String(), nil +} + +func queryPartitionInfo(ctx context.Context, db *sql.DB, schema, table string) (strategy, key string, err error) { + const q = ` +SELECT + CASE c.relkind WHEN 'p' THEN + CASE pt.partstrat + WHEN 'h' THEN 'HASH' + WHEN 'l' THEN 'LIST' + WHEN 'r' THEN 'RANGE' + END + END AS strategy, + pg_catalog.pg_get_partkeydef(c.oid) AS partition_key +FROM pg_catalog.pg_class c +JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace +LEFT JOIN pg_catalog.pg_partitioned_table pt ON pt.partrelid = c.oid +WHERE n.nspname = $1 + AND c.relname = $2 + AND c.relkind = 'p'` + + err = db.QueryRowContext(ctx, q, schema, table).Scan(&strategy, &key) + if err == sql.ErrNoRows { + return "", "", nil + } + if err != nil { + return "", "", fmt.Errorf("query partition info for %s.%s: %w", schema, table, err) + } + return strategy, key, nil +} + type stubColumn struct { name string dataType string From e774cd8e65eda170cd23e95c6994b2cb01e3d6b7 Mon Sep 17 00:00:00 2001 From: tianzhou Date: Thu, 20 Aug 2026 21:53:41 -0700 Subject: [PATCH 2/8] fix: use pg_get_partkeydef directly to avoid duplicating partition strategy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pg_get_partkeydef returns the full partition key definition including the strategy (e.g., "RANGE (id)"), so we must not prepend the strategy separately. Simplify queryPartitionInfo → queryPartitionKeyDef. Co-Authored-By: Claude Opus 4.6 --- ir/stub.go | 32 +++++++++++--------------------- 1 file changed, 11 insertions(+), 21 deletions(-) diff --git a/ir/stub.go b/ir/stub.go index a3adee0c..ab8963cc 100644 --- a/ir/stub.go +++ b/ir/stub.go @@ -86,11 +86,11 @@ func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table return "", nil } - partStrategy, partKey, err := queryPartitionInfo(ctx, db, schema, table) + partDef, err := queryPartitionKeyDef(ctx, db, schema, table) if err != nil { return "", err } - if partStrategy == "" { + if partDef == "" { return "", nil } @@ -140,39 +140,29 @@ func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table } b.WriteString(") PARTITION BY ") - b.WriteString(partStrategy) - b.WriteString(" (") - b.WriteString(partKey) - b.WriteString(");\n") + b.WriteString(partDef) + b.WriteString(";\n") return b.String(), nil } -func queryPartitionInfo(ctx context.Context, db *sql.DB, schema, table string) (strategy, key string, err error) { +func queryPartitionKeyDef(ctx context.Context, db *sql.DB, schema, table string) (string, error) { const q = ` -SELECT - CASE c.relkind WHEN 'p' THEN - CASE pt.partstrat - WHEN 'h' THEN 'HASH' - WHEN 'l' THEN 'LIST' - WHEN 'r' THEN 'RANGE' - END - END AS strategy, - pg_catalog.pg_get_partkeydef(c.oid) AS partition_key +SELECT pg_catalog.pg_get_partkeydef(c.oid) FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace -LEFT JOIN pg_catalog.pg_partitioned_table pt ON pt.partrelid = c.oid WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind = 'p'` - err = db.QueryRowContext(ctx, q, schema, table).Scan(&strategy, &key) + var def string + err := db.QueryRowContext(ctx, q, schema, table).Scan(&def) if err == sql.ErrNoRows { - return "", "", nil + return "", nil } if err != nil { - return "", "", fmt.Errorf("query partition info for %s.%s: %w", schema, table, err) + return "", fmt.Errorf("query partition key def for %s.%s: %w", schema, table, err) } - return strategy, key, nil + return def, nil } type stubColumn struct { From 27e77796fa4e9f2cf6cf9cc0275b0543040bbe42 Mon Sep 17 00:00:00 2001 From: tianzhou Date: Thu, 20 Aug 2026 21:56:39 -0700 Subject: [PATCH 3/8] fix: drop existing cross-schema parent before stubbing to handle plan DB with pre-existing partitions When the external plan DB already contains the parent table with its partitions, CREATE TABLE IF NOT EXISTS is a no-op and the desired SQL's PARTITION OF still conflicts with existing partitions. Fix by dropping the parent CASCADE before creating the stub, so the desired SQL can attach its partition children cleanly. Co-Authored-By: Claude Opus 4.6 --- cmd/plan/partition_stubs.go | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/cmd/plan/partition_stubs.go b/cmd/plan/partition_stubs.go index 1f4d025a..7a49df36 100644 --- a/cmd/plan/partition_stubs.go +++ b/cmd/plan/partition_stubs.go @@ -2,12 +2,13 @@ package plan import ( "context" + "fmt" "strings" - "github.com/pgplex/pgschema/internal/postgres" - "github.com/pgplex/pgschema/ir" "github.com/pgplex/pgschema/cmd/util" "github.com/pgplex/pgschema/internal/logger" + "github.com/pgplex/pgschema/internal/postgres" + "github.com/pgplex/pgschema/ir" ) // prependPartitionParentStubs creates stub CREATE TABLE ... PARTITION BY @@ -59,6 +60,11 @@ func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig } logger.Get().Debug("prepending stub for cross-schema partition parent", "schema", ref.Schema, "table", ref.Table) + // Drop the parent if it already exists in the plan DB (e.g. when the + // plan DB mirrors the target). CASCADE detaches existing partitions so + // the desired SQL can re-attach its own partition children cleanly. + qualified := ir.QualifyEntityNameWithQuotesMode(ref.Schema, ref.Table, targetSchema, ref.Schema != targetSchema) + stubs.WriteString(fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE;\n", qualified)) stubs.WriteString(ddl) } From b4603b9f5fc741d4010833f8f19f65f93b312d49 Mon Sep 17 00:00:00 2001 From: tianzhou Date: Thu, 20 Aug 2026 21:58:26 -0700 Subject: [PATCH 4/8] Revert "fix: drop existing cross-schema parent before stubbing to handle plan DB with pre-existing partitions" This reverts commit 27e77796fa4e9f2cf6cf9cc0275b0543040bbe42. --- cmd/plan/partition_stubs.go | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/cmd/plan/partition_stubs.go b/cmd/plan/partition_stubs.go index 7a49df36..1f4d025a 100644 --- a/cmd/plan/partition_stubs.go +++ b/cmd/plan/partition_stubs.go @@ -2,13 +2,12 @@ package plan import ( "context" - "fmt" "strings" - "github.com/pgplex/pgschema/cmd/util" - "github.com/pgplex/pgschema/internal/logger" "github.com/pgplex/pgschema/internal/postgres" "github.com/pgplex/pgschema/ir" + "github.com/pgplex/pgschema/cmd/util" + "github.com/pgplex/pgschema/internal/logger" ) // prependPartitionParentStubs creates stub CREATE TABLE ... PARTITION BY @@ -60,11 +59,6 @@ func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig } logger.Get().Debug("prepending stub for cross-schema partition parent", "schema", ref.Schema, "table", ref.Table) - // Drop the parent if it already exists in the plan DB (e.g. when the - // plan DB mirrors the target). CASCADE detaches existing partitions so - // the desired SQL can re-attach its own partition children cleanly. - qualified := ir.QualifyEntityNameWithQuotesMode(ref.Schema, ref.Table, targetSchema, ref.Schema != targetSchema) - stubs.WriteString(fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE;\n", qualified)) stubs.WriteString(ddl) } From 4e46ffe648cc1a27613545f3890d36b68c0b366b Mon Sep 17 00:00:00 2001 From: tianzhou Date: Thu, 20 Aug 2026 22:00:15 -0700 Subject: [PATCH 5/8] fix: safely detach existing partitions instead of destructive DROP CASCADE When the plan DB already has the cross-schema parent with attached partitions, use a PL/pgSQL DO block to DETACH PARTITION each child instead of DROP TABLE CASCADE. This is non-destructive: it only removes the parent-child relationship without dropping any tables. The DO block is a no-op when the parent does not exist (clean plan DB). Co-Authored-By: Claude Opus 4.6 --- cmd/plan/partition_stubs.go | 34 ++++++++++++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/cmd/plan/partition_stubs.go b/cmd/plan/partition_stubs.go index 1f4d025a..35505bb5 100644 --- a/cmd/plan/partition_stubs.go +++ b/cmd/plan/partition_stubs.go @@ -2,12 +2,13 @@ package plan import ( "context" + "fmt" "strings" - "github.com/pgplex/pgschema/internal/postgres" - "github.com/pgplex/pgschema/ir" "github.com/pgplex/pgschema/cmd/util" "github.com/pgplex/pgschema/internal/logger" + "github.com/pgplex/pgschema/internal/postgres" + "github.com/pgplex/pgschema/ir" ) // prependPartitionParentStubs creates stub CREATE TABLE ... PARTITION BY @@ -59,8 +60,37 @@ func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig } logger.Get().Debug("prepending stub for cross-schema partition parent", "schema", ref.Schema, "table", ref.Table) + // If the parent already exists in the plan DB with attached partitions, + // detach them so the desired SQL can re-attach its own children cleanly. + // This DO block is a no-op when the parent does not exist. + stubs.WriteString(fmt.Sprintf(`DO $pgschema_detach$ +DECLARE + detach_sql text; +BEGIN + SELECT string_agg( + format('ALTER TABLE %s.%s DETACH PARTITION %%s', inhrelid::regclass), + '; ' + ) INTO detach_sql + FROM pg_catalog.pg_inherits + WHERE inhparent = ( + SELECT c.oid FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = %s AND c.relname = %s AND c.relkind = 'p' + ); + IF detach_sql IS NOT NULL THEN + EXECUTE detach_sql; + END IF; +END $pgschema_detach$; +`, + ir.QuoteIdentifier(ref.Schema), ir.QuoteIdentifier(ref.Table), + pgQuoteLiteral(ref.Schema), pgQuoteLiteral(ref.Table), + )) stubs.WriteString(ddl) } return stubs.String() + desiredSQL, nil } + +func pgQuoteLiteral(s string) string { + return "'" + strings.ReplaceAll(s, "'", "''") + "'" +} From e7590f1b0247f17220e80f92ea1df648694aee50 Mon Sep 17 00:00:00 2001 From: tianzhou Date: Thu, 20 Aug 2026 22:02:14 -0700 Subject: [PATCH 6/8] Revert "fix: safely detach existing partitions instead of destructive DROP CASCADE" This reverts commit 4e46ffe648cc1a27613545f3890d36b68c0b366b. --- cmd/plan/partition_stubs.go | 34 ++-------------------------------- 1 file changed, 2 insertions(+), 32 deletions(-) diff --git a/cmd/plan/partition_stubs.go b/cmd/plan/partition_stubs.go index 35505bb5..1f4d025a 100644 --- a/cmd/plan/partition_stubs.go +++ b/cmd/plan/partition_stubs.go @@ -2,13 +2,12 @@ package plan import ( "context" - "fmt" "strings" - "github.com/pgplex/pgschema/cmd/util" - "github.com/pgplex/pgschema/internal/logger" "github.com/pgplex/pgschema/internal/postgres" "github.com/pgplex/pgschema/ir" + "github.com/pgplex/pgschema/cmd/util" + "github.com/pgplex/pgschema/internal/logger" ) // prependPartitionParentStubs creates stub CREATE TABLE ... PARTITION BY @@ -60,37 +59,8 @@ func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig } logger.Get().Debug("prepending stub for cross-schema partition parent", "schema", ref.Schema, "table", ref.Table) - // If the parent already exists in the plan DB with attached partitions, - // detach them so the desired SQL can re-attach its own children cleanly. - // This DO block is a no-op when the parent does not exist. - stubs.WriteString(fmt.Sprintf(`DO $pgschema_detach$ -DECLARE - detach_sql text; -BEGIN - SELECT string_agg( - format('ALTER TABLE %s.%s DETACH PARTITION %%s', inhrelid::regclass), - '; ' - ) INTO detach_sql - FROM pg_catalog.pg_inherits - WHERE inhparent = ( - SELECT c.oid FROM pg_catalog.pg_class c - JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace - WHERE n.nspname = %s AND c.relname = %s AND c.relkind = 'p' - ); - IF detach_sql IS NOT NULL THEN - EXECUTE detach_sql; - END IF; -END $pgschema_detach$; -`, - ir.QuoteIdentifier(ref.Schema), ir.QuoteIdentifier(ref.Table), - pgQuoteLiteral(ref.Schema), pgQuoteLiteral(ref.Table), - )) stubs.WriteString(ddl) } return stubs.String() + desiredSQL, nil } - -func pgQuoteLiteral(s string) string { - return "'" + strings.ReplaceAll(s, "'", "''") + "'" -} From 7a7444e6e0ce2a0c4a08fa8866c9f4094cef2266 Mon Sep 17 00:00:00 2001 From: tianzhou Date: Thu, 20 Aug 2026 22:15:33 -0700 Subject: [PATCH 7/8] fix: keep partition stubs inside temp schema to avoid plan DB mutations Redesign the cross-schema partition stub approach to stay within the temp schema sandbox: - BuildPartitionedTableStubSQL now emits unqualified table names so the stub lands in the temp schema via search_path - PARTITION OF references to cross-schema parents are rewritten to strip the schema prefix, resolving to the stub via search_path - Return an error when a referenced partition parent is not found on the target database instead of silently skipping - Add unit tests for replacePartitionOfRef This avoids creating or mutating persistent objects (schemas, tables) outside the temp schema in the plan database. Co-Authored-By: Claude Opus 4.6 --- cmd/plan/partition_stubs.go | 104 ++++++++++++++++++++++++++++--- cmd/plan/partition_stubs_test.go | 58 +++++++++++++++++ ir/stub.go | 19 +++--- 3 files changed, 162 insertions(+), 19 deletions(-) create mode 100644 cmd/plan/partition_stubs_test.go diff --git a/cmd/plan/partition_stubs.go b/cmd/plan/partition_stubs.go index 1f4d025a..2b09bcb7 100644 --- a/cmd/plan/partition_stubs.go +++ b/cmd/plan/partition_stubs.go @@ -2,18 +2,21 @@ package plan import ( "context" + "fmt" "strings" - "github.com/pgplex/pgschema/internal/postgres" - "github.com/pgplex/pgschema/ir" "github.com/pgplex/pgschema/cmd/util" "github.com/pgplex/pgschema/internal/logger" + "github.com/pgplex/pgschema/internal/postgres" + "github.com/pgplex/pgschema/ir" ) // prependPartitionParentStubs creates stub CREATE TABLE ... PARTITION BY // statements for cross-schema partition parents referenced by PARTITION OF -// in the desired SQL. Without these stubs the plan database rejects the -// desired-state SQL because the parent table does not exist. +// in the desired SQL. The stubs use unqualified names so they land in the +// temp schema via search_path, and the PARTITION OF references in the +// desired SQL are rewritten to also be unqualified. This avoids creating +// or mutating persistent objects outside the temp schema in the plan database. func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig, targetSchema, desiredSQL string) (string, error) { refs := postgres.ExtractPartitionOfTargets(desiredSQL, targetSchema) if len(refs) == 0 { @@ -49,18 +52,105 @@ func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig defer conn.Close() var stubs strings.Builder + rewritten := desiredSQL for _, ref := range toStub { - ddl, err := ir.BuildPartitionedTableStubSQL(ctx, conn, ref.Schema, ref.Table, targetSchema) + ddl, err := ir.BuildPartitionedTableStubSQL(ctx, conn, ref.Schema, ref.Table) if err != nil { return "", err } if ddl == "" { - continue + logger.Get().Warn("cross-schema partition parent not found or not partitioned", + "schema", ref.Schema, "table", ref.Table) + return "", fmt.Errorf("cross-schema partition parent %s.%s not found or not a partitioned table on the target database", ref.Schema, ref.Table) } logger.Get().Debug("prepending stub for cross-schema partition parent", "schema", ref.Schema, "table", ref.Table) stubs.WriteString(ddl) + + // Rewrite PARTITION OF references to strip the cross-schema prefix so + // they resolve to the stub in the temp schema via search_path. + rewritten = stripPartitionOfSchema(rewritten, ref.Schema, ref.Table) } - return stubs.String() + desiredSQL, nil + return stubs.String() + rewritten, nil +} + +// stripPartitionOfSchema rewrites "PARTITION OF schema.table" to +// "PARTITION OF table" in SQL, handling both quoted and unquoted identifiers. +func stripPartitionOfSchema(sql, schema, table string) string { + // Build the qualified reference as it appears in SQL. + // Handle both quoted and unquoted forms. + patterns := []struct{ old, new string }{ + // unquoted: schema.table + { + old: strings.ToLower(schema) + "." + strings.ToLower(table), + new: strings.ToLower(table), + }, + } + + // Also handle quoted schema: "schema".table, "schema"."table", schema."table" + quotedSchema := ir.QuoteIdentifier(schema) + quotedTable := ir.QuoteIdentifier(table) + if quotedSchema != schema { + patterns = append(patterns, + struct{ old, new string }{quotedSchema + "." + table, table}, + struct{ old, new string }{quotedSchema + "." + quotedTable, quotedTable}, + ) + } + if quotedTable != table { + patterns = append(patterns, + struct{ old, new string }{schema + "." + quotedTable, quotedTable}, + ) + } + + result := sql + for _, p := range patterns { + result = replacePartitionOfRef(result, p.old, p.new) + } + return result +} + +// replacePartitionOfRef replaces "PARTITION OF " with +// "PARTITION OF " in SQL, case-insensitively matching the +// PARTITION OF keywords. +func replacePartitionOfRef(sql, oldRef, newRef string) string { + lower := strings.ToLower(sql) + lowerOld := strings.ToLower(oldRef) + var b strings.Builder + i := 0 + for i < len(sql) { + // Find "partition" keyword + idx := strings.Index(lower[i:], "partition") + if idx < 0 { + b.WriteString(sql[i:]) + break + } + pos := i + idx + b.WriteString(sql[i:pos]) + + // Check if it's followed by whitespace + "of" + whitespace + oldRef + rest := pos + len("partition") + j := rest + for j < len(sql) && (sql[j] == ' ' || sql[j] == '\t' || sql[j] == '\n' || sql[j] == '\r') { + j++ + } + if j+2 <= len(sql) && strings.EqualFold(sql[j:j+2], "of") { + afterOf := j + 2 + k := afterOf + for k < len(sql) && (sql[k] == ' ' || sql[k] == '\t' || sql[k] == '\n' || sql[k] == '\r') { + k++ + } + if k+len(lowerOld) <= len(sql) && strings.EqualFold(sql[k:k+len(lowerOld)], lowerOld) { + // Matched — write "PARTITION OF " preserving original case of keywords + b.WriteString(sql[pos:afterOf]) + b.WriteString(sql[afterOf:k]) + b.WriteString(newRef) + i = k + len(lowerOld) + continue + } + } + b.WriteString(sql[pos : pos+len("partition")]) + i = pos + len("partition") + } + return b.String() } diff --git a/cmd/plan/partition_stubs_test.go b/cmd/plan/partition_stubs_test.go new file mode 100644 index 00000000..5690d76c --- /dev/null +++ b/cmd/plan/partition_stubs_test.go @@ -0,0 +1,58 @@ +package plan + +import "testing" + +func TestReplacePartitionOfRef(t *testing.T) { + tests := []struct { + name string + sql string + oldRef string + newRef string + want string + }{ + { + name: "simple unqualified rewrite", + sql: "CREATE TABLE child PARTITION OF public.parent FOR VALUES IN ('a');", + oldRef: "public.parent", + newRef: "parent", + want: "CREATE TABLE child PARTITION OF parent FOR VALUES IN ('a');", + }, + { + name: "case insensitive keywords", + sql: "CREATE TABLE child partition of public.parent FOR VALUES IN ('a');", + oldRef: "public.parent", + newRef: "parent", + want: "CREATE TABLE child partition of parent FOR VALUES IN ('a');", + }, + { + name: "multiple partitions same parent", + sql: "CREATE TABLE c1 PARTITION OF public.parent FOR VALUES IN ('a');\nCREATE TABLE c2 PARTITION OF public.parent FOR VALUES IN ('b');", + oldRef: "public.parent", + newRef: "parent", + want: "CREATE TABLE c1 PARTITION OF parent FOR VALUES IN ('a');\nCREATE TABLE c2 PARTITION OF parent FOR VALUES IN ('b');", + }, + { + name: "partition by not affected", + sql: "CREATE TABLE parent (id int) PARTITION BY RANGE (id);", + oldRef: "public.parent", + newRef: "parent", + want: "CREATE TABLE parent (id int) PARTITION BY RANGE (id);", + }, + { + name: "no match leaves sql unchanged", + sql: "CREATE TABLE child PARTITION OF other.parent FOR VALUES IN ('a');", + oldRef: "public.parent", + newRef: "parent", + want: "CREATE TABLE child PARTITION OF other.parent FOR VALUES IN ('a');", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := replacePartitionOfRef(tt.sql, tt.oldRef, tt.newRef) + if got != tt.want { + t.Errorf("replacePartitionOfRef()\ngot: %s\nwant: %s", got, tt.want) + } + }) + } +} diff --git a/ir/stub.go b/ir/stub.go index ab8963cc..fa039655 100644 --- a/ir/stub.go +++ b/ir/stub.go @@ -71,13 +71,15 @@ func BuildTableStubSQL(ctx context.Context, db *sql.DB, schema, table, targetSch return b.String(), nil } -// BuildPartitionedTableStubSQL returns CREATE SCHEMA / CREATE TABLE DDL for a -// partitioned parent table that is referenced by PARTITION OF from a child in -// another schema. The stub includes the PARTITION BY clause so the child can +// BuildPartitionedTableStubSQL returns CREATE TABLE DDL for a partitioned +// parent table that is referenced by PARTITION OF from a child in another +// schema. The stub uses an unqualified table name so it lands in the temp +// schema via search_path, avoiding mutations to persistent schemas in the +// plan database. The stub includes the PARTITION BY clause so the child can // attach as a partition. // // Returns an empty string if the table does not exist or is not partitioned. -func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table, targetSchema string) (string, error) { +func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table string) (string, error) { cols, err := queryStubColumns(ctx, db, schema, table) if err != nil { return "", err @@ -100,20 +102,13 @@ func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table } var b strings.Builder - qualified := QualifyEntityNameWithQuotesMode(schema, table, targetSchema, schema != targetSchema) - - if schema != targetSchema { - b.WriteString("CREATE SCHEMA IF NOT EXISTS ") - b.WriteString(QuoteIdentifier(schema)) - b.WriteString(";\n") - } b.WriteString("-- pgschema: stub for cross-schema partition parent ") b.WriteString(sanitizeComment(schema)) b.WriteString(".") b.WriteString(sanitizeComment(table)) b.WriteString("\nCREATE TABLE IF NOT EXISTS ") - b.WriteString(qualified) + b.WriteString(QuoteIdentifier(table)) b.WriteString(" (\n") for i, col := range cols { From 4dd147fe9fe6534eae21dfc194ab2136e6824d1e Mon Sep 17 00:00:00 2001 From: tianzhou Date: Thu, 20 Aug 2026 22:23:15 -0700 Subject: [PATCH 8/8] fix: use unique prefixed stub names to avoid identity collisions Partition parent stubs now use _pgschema_partstub___ as the table name instead of the bare table name. This prevents collisions when: - the target schema has a table with the same name as the parent - two cross-schema parents share the same table name (e.g., public.events and archive.events) PARTITION OF references are rewritten to point to the unique stub name. Co-Authored-By: Claude Opus 4.6 --- cmd/plan/partition_stubs.go | 53 ++++++++++++++++++-------------- cmd/plan/partition_stubs_test.go | 40 ++++++++++++++++++------ ir/stub.go | 12 ++++---- 3 files changed, 67 insertions(+), 38 deletions(-) diff --git a/cmd/plan/partition_stubs.go b/cmd/plan/partition_stubs.go index 2b09bcb7..f027b51e 100644 --- a/cmd/plan/partition_stubs.go +++ b/cmd/plan/partition_stubs.go @@ -11,12 +11,22 @@ import ( "github.com/pgplex/pgschema/ir" ) +// partitionStubPrefix is prepended to stub table names to avoid collisions +// with real tables in the temp schema. +const partitionStubPrefix = "_pgschema_partstub_" + +// partitionStubName returns a unique unqualified name for a cross-schema +// partition parent stub: _pgschema_partstub___
. +func partitionStubName(schema, table string) string { + return partitionStubPrefix + schema + "__" + table +} + // prependPartitionParentStubs creates stub CREATE TABLE ... PARTITION BY // statements for cross-schema partition parents referenced by PARTITION OF -// in the desired SQL. The stubs use unqualified names so they land in the -// temp schema via search_path, and the PARTITION OF references in the -// desired SQL are rewritten to also be unqualified. This avoids creating -// or mutating persistent objects outside the temp schema in the plan database. +// in the desired SQL. Each stub uses a unique prefixed name to avoid +// collisions with real tables or other stubs, and PARTITION OF references +// are rewritten to point to the stub name. Everything stays inside the +// temp schema via search_path. func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig, targetSchema, desiredSQL string) (string, error) { refs := postgres.ExtractPartitionOfTargets(desiredSQL, targetSchema) if len(refs) == 0 { @@ -54,7 +64,8 @@ func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig var stubs strings.Builder rewritten := desiredSQL for _, ref := range toStub { - ddl, err := ir.BuildPartitionedTableStubSQL(ctx, conn, ref.Schema, ref.Table) + stubName := partitionStubName(ref.Schema, ref.Table) + ddl, err := ir.BuildPartitionedTableStubSQL(ctx, conn, ref.Schema, ref.Table, stubName) if err != nil { return "", err } @@ -64,42 +75,41 @@ func prependPartitionParentStubs(ctx context.Context, cfg *util.ConnectionConfig return "", fmt.Errorf("cross-schema partition parent %s.%s not found or not a partitioned table on the target database", ref.Schema, ref.Table) } logger.Get().Debug("prepending stub for cross-schema partition parent", - "schema", ref.Schema, "table", ref.Table) + "schema", ref.Schema, "table", ref.Table, "stubName", stubName) stubs.WriteString(ddl) - // Rewrite PARTITION OF references to strip the cross-schema prefix so - // they resolve to the stub in the temp schema via search_path. - rewritten = stripPartitionOfSchema(rewritten, ref.Schema, ref.Table) + // Rewrite PARTITION OF references to point to the stub name. + rewritten = rewritePartitionOfRef(rewritten, ref.Schema, ref.Table, stubName) } return stubs.String() + rewritten, nil } -// stripPartitionOfSchema rewrites "PARTITION OF schema.table" to -// "PARTITION OF table" in SQL, handling both quoted and unquoted identifiers. -func stripPartitionOfSchema(sql, schema, table string) string { - // Build the qualified reference as it appears in SQL. - // Handle both quoted and unquoted forms. +// rewritePartitionOfRef rewrites "PARTITION OF schema.table" to +// "PARTITION OF stubName" in SQL, handling both quoted and unquoted identifiers. +func rewritePartitionOfRef(sql, schema, table, stubName string) string { + // Build the qualified reference patterns as they may appear in SQL. patterns := []struct{ old, new string }{ - // unquoted: schema.table + // unquoted: schema.table → stubName { old: strings.ToLower(schema) + "." + strings.ToLower(table), - new: strings.ToLower(table), + new: ir.QuoteIdentifier(stubName), }, } - // Also handle quoted schema: "schema".table, "schema"."table", schema."table" + // Also handle quoted identifiers quotedSchema := ir.QuoteIdentifier(schema) quotedTable := ir.QuoteIdentifier(table) + quotedStub := ir.QuoteIdentifier(stubName) if quotedSchema != schema { patterns = append(patterns, - struct{ old, new string }{quotedSchema + "." + table, table}, - struct{ old, new string }{quotedSchema + "." + quotedTable, quotedTable}, + struct{ old, new string }{quotedSchema + "." + table, quotedStub}, + struct{ old, new string }{quotedSchema + "." + quotedTable, quotedStub}, ) } if quotedTable != table { patterns = append(patterns, - struct{ old, new string }{schema + "." + quotedTable, quotedTable}, + struct{ old, new string }{schema + "." + quotedTable, quotedStub}, ) } @@ -119,7 +129,6 @@ func replacePartitionOfRef(sql, oldRef, newRef string) string { var b strings.Builder i := 0 for i < len(sql) { - // Find "partition" keyword idx := strings.Index(lower[i:], "partition") if idx < 0 { b.WriteString(sql[i:]) @@ -128,7 +137,6 @@ func replacePartitionOfRef(sql, oldRef, newRef string) string { pos := i + idx b.WriteString(sql[i:pos]) - // Check if it's followed by whitespace + "of" + whitespace + oldRef rest := pos + len("partition") j := rest for j < len(sql) && (sql[j] == ' ' || sql[j] == '\t' || sql[j] == '\n' || sql[j] == '\r') { @@ -141,7 +149,6 @@ func replacePartitionOfRef(sql, oldRef, newRef string) string { k++ } if k+len(lowerOld) <= len(sql) && strings.EqualFold(sql[k:k+len(lowerOld)], lowerOld) { - // Matched — write "PARTITION OF " preserving original case of keywords b.WriteString(sql[pos:afterOf]) b.WriteString(sql[afterOf:k]) b.WriteString(newRef) diff --git a/cmd/plan/partition_stubs_test.go b/cmd/plan/partition_stubs_test.go index 5690d76c..1fd213d4 100644 --- a/cmd/plan/partition_stubs_test.go +++ b/cmd/plan/partition_stubs_test.go @@ -2,6 +2,21 @@ package plan import "testing" +func TestPartitionStubName(t *testing.T) { + tests := []struct { + schema, table, want string + }{ + {"public", "events", "_pgschema_partstub_public__events"}, + {"archive", "events", "_pgschema_partstub_archive__events"}, + } + for _, tt := range tests { + got := partitionStubName(tt.schema, tt.table) + if got != tt.want { + t.Errorf("partitionStubName(%q, %q) = %q, want %q", tt.schema, tt.table, got, tt.want) + } + } +} + func TestReplacePartitionOfRef(t *testing.T) { tests := []struct { name string @@ -11,40 +26,47 @@ func TestReplacePartitionOfRef(t *testing.T) { want string }{ { - name: "simple unqualified rewrite", + name: "simple rewrite", sql: "CREATE TABLE child PARTITION OF public.parent FOR VALUES IN ('a');", oldRef: "public.parent", - newRef: "parent", - want: "CREATE TABLE child PARTITION OF parent FOR VALUES IN ('a');", + newRef: `"_pgschema_partstub_public__parent"`, + want: `CREATE TABLE child PARTITION OF "_pgschema_partstub_public__parent" FOR VALUES IN ('a');`, }, { name: "case insensitive keywords", sql: "CREATE TABLE child partition of public.parent FOR VALUES IN ('a');", oldRef: "public.parent", - newRef: "parent", - want: "CREATE TABLE child partition of parent FOR VALUES IN ('a');", + newRef: `"_pgschema_partstub_public__parent"`, + want: `CREATE TABLE child partition of "_pgschema_partstub_public__parent" FOR VALUES IN ('a');`, }, { name: "multiple partitions same parent", sql: "CREATE TABLE c1 PARTITION OF public.parent FOR VALUES IN ('a');\nCREATE TABLE c2 PARTITION OF public.parent FOR VALUES IN ('b');", oldRef: "public.parent", - newRef: "parent", - want: "CREATE TABLE c1 PARTITION OF parent FOR VALUES IN ('a');\nCREATE TABLE c2 PARTITION OF parent FOR VALUES IN ('b');", + newRef: `"_pgschema_partstub_public__parent"`, + want: "CREATE TABLE c1 PARTITION OF \"_pgschema_partstub_public__parent\" FOR VALUES IN ('a');\nCREATE TABLE c2 PARTITION OF \"_pgschema_partstub_public__parent\" FOR VALUES IN ('b');", }, { name: "partition by not affected", sql: "CREATE TABLE parent (id int) PARTITION BY RANGE (id);", oldRef: "public.parent", - newRef: "parent", + newRef: `"_pgschema_partstub_public__parent"`, want: "CREATE TABLE parent (id int) PARTITION BY RANGE (id);", }, { name: "no match leaves sql unchanged", sql: "CREATE TABLE child PARTITION OF other.parent FOR VALUES IN ('a');", oldRef: "public.parent", - newRef: "parent", + newRef: `"_pgschema_partstub_public__parent"`, want: "CREATE TABLE child PARTITION OF other.parent FOR VALUES IN ('a');", }, + { + name: "distinct schemas do not collide", + sql: "CREATE TABLE c1 PARTITION OF public.events FOR VALUES IN ('a');\nCREATE TABLE c2 PARTITION OF archive.events FOR VALUES IN ('b');", + oldRef: "public.events", + newRef: `"_pgschema_partstub_public__events"`, + want: "CREATE TABLE c1 PARTITION OF \"_pgschema_partstub_public__events\" FOR VALUES IN ('a');\nCREATE TABLE c2 PARTITION OF archive.events FOR VALUES IN ('b');", + }, } for _, tt := range tests { diff --git a/ir/stub.go b/ir/stub.go index fa039655..d10a60e1 100644 --- a/ir/stub.go +++ b/ir/stub.go @@ -73,13 +73,13 @@ func BuildTableStubSQL(ctx context.Context, db *sql.DB, schema, table, targetSch // BuildPartitionedTableStubSQL returns CREATE TABLE DDL for a partitioned // parent table that is referenced by PARTITION OF from a child in another -// schema. The stub uses an unqualified table name so it lands in the temp -// schema via search_path, avoiding mutations to persistent schemas in the -// plan database. The stub includes the PARTITION BY clause so the child can -// attach as a partition. +// schema. The stub uses the provided stubName (unqualified) so it lands in +// the temp schema via search_path, avoiding mutations to persistent schemas +// in the plan database. The stub includes the PARTITION BY clause so the +// child can attach as a partition. // // Returns an empty string if the table does not exist or is not partitioned. -func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table string) (string, error) { +func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table, stubName string) (string, error) { cols, err := queryStubColumns(ctx, db, schema, table) if err != nil { return "", err @@ -108,7 +108,7 @@ func BuildPartitionedTableStubSQL(ctx context.Context, db *sql.DB, schema, table b.WriteString(".") b.WriteString(sanitizeComment(table)) b.WriteString("\nCREATE TABLE IF NOT EXISTS ") - b.WriteString(QuoteIdentifier(table)) + b.WriteString(QuoteIdentifier(stubName)) b.WriteString(" (\n") for i, col := range cols {