fix: defer new table inline FK when it depends on a new unique constraint (#506)

internal/diff/table.go

@@ -620,7 +620,27 @@ func planFKRecreationForReplacedConstraints(modifiedTables []*tableDiff, addedTa
 			}
 		}
 	}
-	if len(replaced) == 0 {
+
+	// Unique/PK constraints and unique indexes newly added to existing tables
+	// in this migration. These are emitted in the modify phase, so FKs in new
+	// tables that depend on them must be deferred past CREATE TABLE (issue #506).
+	addedConstraints := make(map[string][]*ir.Constraint)
+	addedUniqueIndexes := make(map[string][]*ir.Index)
+	for _, td := range modifiedTables {
+		key := td.Table.Schema + "." + td.Table.Name
+		for _, c := range td.AddedConstraints {
+			if c.Type == ir.ConstraintTypeUnique || c.Type == ir.ConstraintTypePrimaryKey {
+				addedConstraints[key] = append(addedConstraints[key], c)
+			}
+		}
+		for _, idx := range td.AddedIndexes {
+			if idx.Type == ir.IndexTypeUnique && !idx.IsPartial {
+				addedUniqueIndexes[key] = append(addedUniqueIndexes[key], idx)
+			}
+		}
+	}
+
+	if len(replaced) == 0 && len(addedConstraints) == 0 && len(addedUniqueIndexes) == 0 {
 		return nil, nil, nil
 	}
 
@@ -654,17 +674,20 @@ func planFKRecreationForReplacedConstraints(modifiedTables []*tableDiff, addedTa
 		}
 	}
 
-	// FKs on newly added tables that target a replaced constraint: keep them
-	// out of CREATE TABLE (create phase runs before the modify-phase swap) and
-	// add them with the other recreated FKs instead.
+	// FKs on newly added tables that target a replaced or newly-added constraint
+	// (or unique index): keep them out of CREATE TABLE (create phase runs before
+	// the modify phase) and add them after the modify phase instead (#439, #506).
 	suppressedInlineFKs = make(map[string]bool)
 	for _, table := range addedTables {
 		for _, name := range sortedKeys(table.Constraints) {
 			fk := table.Constraints[name]
 			if fk.Type != ir.ConstraintTypeForeignKey {
 				continue
 			}
-			if !fkReferencesAnyConstraint(fk, replaced[fkReferencedTableKey(fk)]) {
+			refKey := fkReferencedTableKey(fk)
+			if !fkReferencesAnyConstraint(fk, replaced[refKey]) &&
+				!fkReferencesAnyConstraint(fk, addedConstraints[refKey]) &&
+				!fkReferencesAnyUniqueIndex(fk, addedUniqueIndexes[refKey]) {
 				continue
 			}
 			suppressedInlineFKs[constraintPathKey(fk)] = true
@@ -720,6 +743,31 @@ func fkReferencesAnyConstraint(fk *ir.Constraint, candidates []*ir.Constraint) b
 	return false
 }
 
+// fkReferencesAnyUniqueIndex reports whether the FK's referenced columns match
+// the column set of any of the given unique indexes.
+func fkReferencesAnyUniqueIndex(fk *ir.Constraint, candidates []*ir.Index) bool {
+	refCols := make(map[string]bool, len(fk.ReferencedColumns))
+	for _, col := range fk.ReferencedColumns {
+		refCols[col.Name] = true
+	}
+	for _, idx := range candidates {
+		if len(idx.Columns) != len(refCols) {
+			continue
+		}
+		matched := true
+		for _, col := range idx.Columns {
+			if !refCols[col.Name] {
+				matched = false
+				break
+			}
+		}
+		if matched {
+			return true
+		}
+	}
+	return false
+}
+
 // generateDropRecreatedFKsSQL drops foreign keys bound to unique/PK constraints
 // being replaced. Must run before the table modifications. (#439)
 func generateDropRecreatedFKsSQL(fks []*ir.Constraint, targetSchema string, collector *diffCollector) {

testdata/diff/create_table/issue_506_fk_depends_on_new_unique/diff.sql

@@ -0,0 +1,24 @@
+CREATE TABLE IF NOT EXISTS child (
+    id uuid,
+    parent_id uuid,
+    tenant text NOT NULL,
+    CONSTRAINT child_pkey PRIMARY KEY (id)
+);
+
+CREATE TABLE IF NOT EXISTS child2 (
+    id uuid,
+    parent2_id uuid,
+    tenant text NOT NULL,
+    CONSTRAINT child2_pkey PRIMARY KEY (id)
+);
+
+ALTER TABLE parent
+ADD CONSTRAINT parent_id_tenant_key UNIQUE (id, tenant);
+
+CREATE UNIQUE INDEX IF NOT EXISTS parent2_id_tenant_key ON parent2 (id, tenant);
+
+ALTER TABLE child
+ADD CONSTRAINT child_parent_id_tenant_fkey FOREIGN KEY (parent_id, tenant) REFERENCES parent (id, tenant);
+
+ALTER TABLE child2
+ADD CONSTRAINT child2_parent2_id_tenant_fkey FOREIGN KEY (parent2_id, tenant) REFERENCES parent2 (id, tenant);

testdata/diff/create_table/issue_506_fk_depends_on_new_unique/new.sql

@@ -0,0 +1,26 @@
+CREATE TABLE parent (
+    id uuid PRIMARY KEY,
+    tenant text NOT NULL,
+    UNIQUE (id, tenant)
+);
+
+CREATE TABLE child (
+    id uuid PRIMARY KEY,
+    parent_id uuid,
+    tenant text NOT NULL,
+    FOREIGN KEY (parent_id, tenant) REFERENCES parent (id, tenant)
+);
+
+CREATE TABLE parent2 (
+    id uuid PRIMARY KEY,
+    tenant text NOT NULL
+);
+
+CREATE UNIQUE INDEX parent2_id_tenant_key ON parent2 (id, tenant);
+
+CREATE TABLE child2 (
+    id uuid PRIMARY KEY,
+    parent2_id uuid,
+    tenant text NOT NULL,
+    FOREIGN KEY (parent2_id, tenant) REFERENCES parent2 (id, tenant)
+);

testdata/diff/create_table/issue_506_fk_depends_on_new_unique/old.sql

@@ -0,0 +1,9 @@
+CREATE TABLE parent (
+    id uuid PRIMARY KEY,
+    tenant text NOT NULL
+);
+
+CREATE TABLE parent2 (
+    id uuid PRIMARY KEY,
+    tenant text NOT NULL
+);

testdata/diff/create_table/issue_506_fk_depends_on_new_unique/plan.json

@@ -0,0 +1,72 @@
+{
+  "version": "1.0.0",
+  "pgschema_version": "1.12.0",
+  "created_at": "1970-01-01T00:00:00Z",
+  "source_fingerprint": {
+    "hash": "a2e9edaed4b13aa629d2ed847ffd01634bf02251aaa187e810dc64ecc17497e1"
+  },
+  "groups": [
+    {
+      "steps": [
+        {
+          "sql": "CREATE TABLE IF NOT EXISTS child (\n    id uuid,\n    parent_id uuid,\n    tenant text NOT NULL,\n    CONSTRAINT child_pkey PRIMARY KEY (id)\n);",
+          "type": "table",
+          "operation": "create",
+          "path": "public.child"
+        },
+        {
+          "sql": "CREATE TABLE IF NOT EXISTS child2 (\n    id uuid,\n    parent2_id uuid,\n    tenant text NOT NULL,\n    CONSTRAINT child2_pkey PRIMARY KEY (id)\n);",
+          "type": "table",
+          "operation": "create",
+          "path": "public.child2"
+        },
+        {
+          "sql": "ALTER TABLE parent\nADD CONSTRAINT parent_id_tenant_key UNIQUE (id, tenant);",
+          "type": "table.constraint",
+          "operation": "create",
+          "path": "public.parent.parent_id_tenant_key"
+        }
+      ]
+    },
+    {
+      "steps": [
+        {
+          "sql": "CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS parent2_id_tenant_key ON parent2 (id, tenant);",
+          "type": "table.index",
+          "operation": "create",
+          "path": "public.parent2.parent2_id_tenant_key"
+        }
+      ]
+    },
+    {
+      "steps": [
+        {
+          "sql": "SELECT \n    COALESCE(i.indisvalid, false) as done,\n    CASE \n        WHEN p.blocks_total > 0 THEN p.blocks_done * 100 / p.blocks_total\n        ELSE 0\n    END as progress\nFROM pg_class c\nLEFT JOIN pg_index i ON c.oid = i.indexrelid\nLEFT JOIN pg_stat_progress_create_index p ON c.oid = p.index_relid\nWHERE c.relname = 'parent2_id_tenant_key';",
+          "directive": {
+            "type": "wait",
+            "message": "Creating index parent2_id_tenant_key"
+          },
+          "type": "table.index",
+          "operation": "create",
+          "path": "public.parent2.parent2_id_tenant_key"
+        }
+      ]
+    },
+    {
+      "steps": [
+        {
+          "sql": "ALTER TABLE child\nADD CONSTRAINT child_parent_id_tenant_fkey FOREIGN KEY (parent_id, tenant) REFERENCES parent (id, tenant);",
+          "type": "table.constraint",
+          "operation": "create",
+          "path": "public.child.child_parent_id_tenant_fkey"
+        },
+        {
+          "sql": "ALTER TABLE child2\nADD CONSTRAINT child2_parent2_id_tenant_fkey FOREIGN KEY (parent2_id, tenant) REFERENCES parent2 (id, tenant);",
+          "type": "table.constraint",
+          "operation": "create",
+          "path": "public.child2.child2_parent2_id_tenant_fkey"
+        }
+      ]
+    }
+  ]
+}

testdata/diff/create_table/issue_506_fk_depends_on_new_unique/plan.sql

@@ -0,0 +1,36 @@
+CREATE TABLE IF NOT EXISTS child (
+    id uuid,
+    parent_id uuid,
+    tenant text NOT NULL,
+    CONSTRAINT child_pkey PRIMARY KEY (id)
+);
+
+CREATE TABLE IF NOT EXISTS child2 (
+    id uuid,
+    parent2_id uuid,
+    tenant text NOT NULL,
+    CONSTRAINT child2_pkey PRIMARY KEY (id)
+);
+
+ALTER TABLE parent
+ADD CONSTRAINT parent_id_tenant_key UNIQUE (id, tenant);
+
+CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS parent2_id_tenant_key ON parent2 (id, tenant);
+
+-- pgschema:wait
+SELECT 
+    COALESCE(i.indisvalid, false) as done,
+    CASE 
+        WHEN p.blocks_total > 0 THEN p.blocks_done * 100 / p.blocks_total
+        ELSE 0
+    END as progress
+FROM pg_class c
+LEFT JOIN pg_index i ON c.oid = i.indexrelid
+LEFT JOIN pg_stat_progress_create_index p ON c.oid = p.index_relid
+WHERE c.relname = 'parent2_id_tenant_key';
+
+ALTER TABLE child
+ADD CONSTRAINT child_parent_id_tenant_fkey FOREIGN KEY (parent_id, tenant) REFERENCES parent (id, tenant);
+
+ALTER TABLE child2
+ADD CONSTRAINT child2_parent2_id_tenant_fkey FOREIGN KEY (parent2_id, tenant) REFERENCES parent2 (id, tenant);

testdata/diff/create_table/issue_506_fk_depends_on_new_unique/plan.txt

@@ -0,0 +1,58 @@
+Plan: 2 to add, 2 to modify.
+
+Summary by type:
+  tables: 2 to add, 2 to modify
+
+Tables:
+  + child
+    + child_parent_id_tenant_fkey (constraint)
+  + child2
+    + child2_parent2_id_tenant_fkey (constraint)
+  ~ parent
+    + parent_id_tenant_key (constraint)
+  ~ parent2
+    + parent2_id_tenant_key (index)
+
+DDL to be executed:
+--------------------------------------------------
+
+-- Transaction Group #1
+CREATE TABLE IF NOT EXISTS child (
+    id uuid,
+    parent_id uuid,
+    tenant text NOT NULL,
+    CONSTRAINT child_pkey PRIMARY KEY (id)
+);
+
+CREATE TABLE IF NOT EXISTS child2 (
+    id uuid,
+    parent2_id uuid,
+    tenant text NOT NULL,
+    CONSTRAINT child2_pkey PRIMARY KEY (id)
+);
+
+ALTER TABLE parent
+ADD CONSTRAINT parent_id_tenant_key UNIQUE (id, tenant);
+
+-- Transaction Group #2
+CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS parent2_id_tenant_key ON parent2 (id, tenant);
+
+-- Transaction Group #3
+-- pgschema:wait
+SELECT 
+    COALESCE(i.indisvalid, false) as done,
+    CASE 
+        WHEN p.blocks_total > 0 THEN p.blocks_done * 100 / p.blocks_total
+        ELSE 0
+    END as progress
+FROM pg_class c
+LEFT JOIN pg_index i ON c.oid = i.indexrelid
+LEFT JOIN pg_stat_progress_create_index p ON c.oid = p.index_relid
+WHERE c.relname = 'parent2_id_tenant_key';
+
+-- Transaction Group #4
+ALTER TABLE child
+ADD CONSTRAINT child_parent_id_tenant_fkey FOREIGN KEY (parent_id, tenant) REFERENCES parent (id, tenant);
+
+ALTER TABLE child2
+ADD CONSTRAINT child2_parent2_id_tenant_fkey FOREIGN KEY (parent2_id, tenant) REFERENCES parent2 (id, tenant);