diff --git a/cmd/plan/partition_stubs.go b/cmd/plan/partition_stubs.go new file mode 100644 index 00000000..f027b51e --- /dev/null +++ b/cmd/plan/partition_stubs.go @@ -0,0 +1,163 @@ +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" +) + +// 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. 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 { + 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 + rewritten := desiredSQL + for _, ref := range toStub { + stubName := partitionStubName(ref.Schema, ref.Table) + ddl, err := ir.BuildPartitionedTableStubSQL(ctx, conn, ref.Schema, ref.Table, stubName) + if err != nil { + return "", err + } + if ddl == "" { + 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, "stubName", stubName) + stubs.WriteString(ddl) + + // Rewrite PARTITION OF references to point to the stub name. + rewritten = rewritePartitionOfRef(rewritten, ref.Schema, ref.Table, stubName) + } + + return stubs.String() + rewritten, nil +} + +// 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 → stubName + { + old: strings.ToLower(schema) + "." + strings.ToLower(table), + new: ir.QuoteIdentifier(stubName), + }, + } + + // 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, quotedStub}, + struct{ old, new string }{quotedSchema + "." + quotedTable, quotedStub}, + ) + } + if quotedTable != table { + patterns = append(patterns, + struct{ old, new string }{schema + "." + quotedTable, quotedStub}, + ) + } + + 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) { + idx := strings.Index(lower[i:], "partition") + if idx < 0 { + b.WriteString(sql[i:]) + break + } + pos := i + idx + b.WriteString(sql[i:pos]) + + 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) { + 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..1fd213d4 --- /dev/null +++ b/cmd/plan/partition_stubs_test.go @@ -0,0 +1,80 @@ +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 + sql string + oldRef string + newRef string + want string + }{ + { + name: "simple rewrite", + sql: "CREATE TABLE child PARTITION OF public.parent FOR VALUES IN ('a');", + oldRef: "public.parent", + 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: `"_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: `"_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: `"_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: `"_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 { + 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/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..d10a60e1 100644 --- a/ir/stub.go +++ b/ir/stub.go @@ -71,6 +71,95 @@ func BuildTableStubSQL(ctx context.Context, db *sql.DB, schema, table, targetSch return b.String(), nil } +// 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 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, stubName string) (string, error) { + cols, err := queryStubColumns(ctx, db, schema, table) + if err != nil { + return "", err + } + if len(cols) == 0 { + return "", nil + } + + partDef, err := queryPartitionKeyDef(ctx, db, schema, table) + if err != nil { + return "", err + } + if partDef == "" { + return "", nil + } + + constraints, err := queryStubConstraints(ctx, db, schema, table) + if err != nil { + return "", err + } + + var b strings.Builder + + 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(QuoteIdentifier(stubName)) + 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(partDef) + b.WriteString(";\n") + return b.String(), nil +} + +func queryPartitionKeyDef(ctx context.Context, db *sql.DB, schema, table string) (string, error) { + const q = ` +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 +WHERE n.nspname = $1 + AND c.relname = $2 + AND c.relkind = 'p'` + + var def string + err := db.QueryRowContext(ctx, q, schema, table).Scan(&def) + if err == sql.ErrNoRows { + return "", nil + } + if err != nil { + return "", fmt.Errorf("query partition key def for %s.%s: %w", schema, table, err) + } + return def, nil +} + type stubColumn struct { name string dataType string