fix: stub cross-schema partition parents in plan database

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_<schema>__<table>.
+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 <oldRef>" with
+// "PARTITION OF <newRef>" 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()
+}

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)
+			}
+		})
+	}
+}

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)

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 {

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)
+			}
+		})
+	}
+}