Repository: ariga/atlas
Pull request #669 (2022-03-24 13:56:18)

sql/internal/sqlx: keep table self references on plan


Patch:
diff --git a/sql/internal/sqlx/plan.go b/sql/internal/sqlx/plan.go
index c3e4a92e..44cd4030 100644
--- a/sql/internal/sqlx/plan.go
+++ b/sql/internal/sqlx/plan.go
@@ -63,16 +63,21 @@ func detachReferences(changes []schema.Change) []schema.Change {
 	for _, change := range changes {
 		switch change := change.(type) {
 		case *schema.AddTable:
-			var fks []schema.Change
+			var (
+				ext  []schema.Change
+				self []*schema.ForeignKey
+			)
 			for _, fk := range change.T.ForeignKeys {
-				if fk.RefTable != change.T {
-					fks = append(fks, &schema.AddForeignKey{F: fk})
+				if fk.RefTable == change.T {
+					self = append(self, fk)
+				} else {
+					ext = append(ext, &schema.AddForeignKey{F: fk})
 				}
 			}
-			if len(fks) > 0 {
-				deferred = append(deferred, &schema.ModifyTable{T: change.T, Changes: fks})
+			if len(ext) > 0 {
+				deferred = append(deferred, &schema.ModifyTable{T: change.T, Changes: ext})
 				t := *change.T
-				t.ForeignKeys = nil
+				t.ForeignKeys = self
 				change = &schema.AddTable{T: &t, Extra: change.Extra}
 			}
 			planned = append(planned, change)


Test patch:
diff --git a/sql/internal/sqlx/plan_test.go b/sql/internal/sqlx/plan_test.go
index 5e11dd49..64fb5085 100644
--- a/sql/internal/sqlx/plan_test.go
+++ b/sql/internal/sqlx/plan_test.go
@@ -28,7 +28,7 @@ func TestDetachCycles(t *testing.T) {
 		},
 	}
 	users.ForeignKeys = []*schema.ForeignKey{
-		{Symbol: "workplace", Table: users, Columns: users.Columns[1:], RefTable: workplaces, RefColumns: workplaces.Columns[:1]},
+		{Symbol: "workplace", Table: users, Columns: users.Columns[1:2], RefTable: workplaces, RefColumns: workplaces.Columns[:1]},
 	}
 	changes := []schema.Change{&schema.AddTable{T: workplaces}, &schema.AddTable{T: users}}
 	planned, err := DetachCycles(changes)
@@ -44,11 +44,15 @@ func TestDetachCycles(t *testing.T) {
 	workplaces.ForeignKeys = []*schema.ForeignKey{
 		{Symbol: "owner", Table: workplaces, Columns: workplaces.Columns[1:], RefTable: users, RefColumns: users.Columns[:1]},
 	}
+	// Add a self-ref foreign-key.
+	users.Columns = append(users.Columns, &schema.Column{Name: "spouse_id", Type: &schema.ColumnType{Raw: "bigint", Null: true}})
+	users.ForeignKeys = append(users.ForeignKeys, &schema.ForeignKey{Symbol: "spouse", Table: users, Columns: users.Columns[2:], RefTable: users, RefColumns: users.Columns[:1]})
+
 	planned, err = DetachCycles(changes)
 	require.NoError(t, err)
 	require.Len(t, planned, 4)
 	require.Empty(t, planned[0].(*schema.AddTable).T.ForeignKeys)
-	require.Empty(t, planned[1].(*schema.AddTable).T.ForeignKeys)
+	require.NotEmpty(t, planned[1].(*schema.AddTable).T.ForeignKeys)
 	require.Equal(t, &schema.ModifyTable{
 		T: workplaces,
 		Changes: []schema.Change{
@@ -61,7 +65,7 @@ func TestDetachCycles(t *testing.T) {
 		T: users,
 		Changes: []schema.Change{
 			&schema.AddForeignKey{
-				F: &schema.ForeignKey{Symbol: "workplace", Table: users, Columns: users.Columns[1:], RefTable: workplaces, RefColumns: workplaces.Columns[:1]},
+				F: &schema.ForeignKey{Symbol: "workplace", Table: users, Columns: users.Columns[1:2], RefTable: workplaces, RefColumns: workplaces.Columns[:1]},
 			},
 		},
 	}, planned[3])
@@ -72,7 +76,7 @@ func TestDetachCycles(t *testing.T) {
 		T: users,
 		Changes: []schema.Change{
 			&schema.DropForeignKey{
-				F: &schema.ForeignKey{Symbol: "workplace", Table: users, Columns: users.Columns[1:], RefTable: workplaces, RefColumns: workplaces.Columns[:1]},
+				F: &schema.ForeignKey{Symbol: "workplace", Table: users, Columns: users.Columns[1:2], RefTable: workplaces, RefColumns: workplaces.Columns[:1]},
 			},
 		},
 	}, planned[0])