Repository: ariga/atlas
Pull request #2724 (2024-04-26 18:19:43)

sql: minor fixes in migration planning


Patch:
diff --git a/sql/internal/sqlx/plan.go b/sql/internal/sqlx/plan.go
index c8eeccf9..955f3702 100644
--- a/sql/internal/sqlx/plan.go
+++ b/sql/internal/sqlx/plan.go
@@ -452,13 +452,21 @@ func dependsOn(c1, c2 schema.Change) bool {
 	case *schema.DropSchema:
 		switch c2 := c2.(type) {
 		case *schema.DropFunc:
-			return c1.S.Name == c2.F.Schema.Name
+			return sameSchema(c1.S, c2.F.Schema)
 		case *schema.DropProc:
-			return c1.S.Name == c2.P.Schema.Name
+			return sameSchema(c1.S, c2.P.Schema)
 		case *schema.DropTable:
-			return c1.S.Name == c2.T.Schema.Name
+			// Schema must be dropped after all its tables and references to them.
+			return sameSchema(c1.S, c2.T.Schema) || slices.ContainsFunc(c2.T.ForeignKeys, func(fk *schema.ForeignKey) bool {
+				return sameSchema(c1.S, fk.RefTable.Schema)
+			})
+		case *schema.ModifyTable:
+			return sameSchema(c1.S, c2.T.Schema) || slices.ContainsFunc(c2.Changes, func(c schema.Change) bool {
+				fk, ok := c.(*schema.DropForeignKey)
+				return ok && sameSchema(c1.S, fk.F.RefTable.Schema)
+			})
 		case *schema.DropView:
-			return c1.S.Name == c2.V.Schema.Name
+			return sameSchema(c1.S, c2.V.Schema)
 		}
 	case *schema.AddTable:
 		switch c2 := c2.(type) {
@@ -495,7 +503,7 @@ func dependsOn(c1, c2 schema.Change) bool {
 		case *schema.ModifyTable:
 			return slices.ContainsFunc(c2.Changes, func(c schema.Change) bool {
 				fk, ok := c.(*schema.DropForeignKey)
-				return ok && fk.F.RefTable == c2.T
+				return ok && refTo([]*schema.ForeignKey{fk.F}, c1.T)
 			})
 		case *schema.DropTrigger:
 			if c2.T.Table == c1.T {
@@ -513,7 +521,7 @@ func dependsOn(c1, c2 schema.Change) bool {
 			// Tables need to be created before referencing them.
 			return slices.ContainsFunc(c1.Changes, func(c schema.Change) bool {
 				fk, ok := c.(*schema.AddForeignKey)
-				return ok && fk.F.RefTable == c2.T
+				return ok && refTo([]*schema.ForeignKey{fk.F}, c2.T)
 			})
 		case *schema.ModifyTable:
 			if c1.T != c2.T {
@@ -525,7 +533,7 @@ func dependsOn(c1, c2 schema.Change) bool {
 				}
 				return slices.ContainsFunc(c1.Changes, func(c schema.Change) bool {
 					fk, ok := c.(*schema.AddForeignKey)
-					return ok && fk.F.RefTable == c2.T && slices.ContainsFunc(fk.F.RefColumns, func(c *schema.Column) bool { return addC[c] })
+					return ok && refTo([]*schema.ForeignKey{fk.F}, c2.T) && slices.ContainsFunc(fk.F.Columns, func(c *schema.Column) bool { return addC[c] })
 				})
 			}
 		case *schema.AddObject:
@@ -810,7 +818,9 @@ func depOfAdd(refs []schema.Object, c schema.Change) bool {
 
 // refTo reports if the given foreign keys reference the given table.
 func refTo(fks []*schema.ForeignKey, to *schema.Table) bool {
-	return slices.ContainsFunc(fks, func(fk *schema.ForeignKey) bool { return fk.RefTable == to })
+	return slices.ContainsFunc(fks, func(fk *schema.ForeignKey) bool {
+		return SameTable(fk.RefTable, to)
+	})
 }
 
 // dependsOnT reports if t1 depends on t2.
@@ -820,6 +830,14 @@ func dependsOnT(t1, t2 schema.Type) bool {
 	return t1 == t2 || schema.UnderlyingType(t1) == t2
 }
 
+// SameTable reports if the two objects represent the same table.
+func SameTable(t1, t2 *schema.Table) bool {
+	if t1 == nil || t2 == nil {
+		return t1 == t2
+	}
+	return t1.Name == t2.Name && sameSchema(t1.Schema, t2.Schema)
+}
+
 // sameSchema reports if the given schemas are the same.
 // Objects can be different as they might reside in two
 // different states (current and desired).
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index 98708cb3..51c93af2 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -7,6 +7,7 @@ package postgres
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
 	"hash/fnv"
 	"math/rand"
@@ -237,7 +238,20 @@ func (d *Driver) SchemaRestoreFunc(desired *schema.Schema) migrate.RestoreFunc {
 
 // RealmRestoreFunc returns a function that restores the given realm to its desired state.
 func (d *Driver) RealmRestoreFunc(desired *schema.Realm) migrate.RestoreFunc {
-	return func(ctx context.Context) error {
+	return func(ctx context.Context) (err error) {
+		// Default behavior for Postgres is to have a single "public" schema.
+		// In that case, all other schemas are dropped, but this one is cleared
+		// object by object. To keep process faster, we drop the schema and recreate it.
+		if !d.crdb && len(desired.Schemas) == 1 && desired.Schemas[0].Name == "public" {
+			if pb := desired.Schemas[0]; len(pb.Tables)+len(pb.Views)+len(pb.Funcs)+len(pb.Procs)+len(pb.Objects) == 0 {
+				desired = &schema.Realm{Attrs: desired.Attrs, Objects: desired.Objects}
+				defer func() {
+					err = errors.Join(err, d.ApplyChanges(ctx, []schema.Change{
+						&schema.AddSchema{S: pb, Extra: []schema.Clause{&schema.IfExists{}}},
+					}))
+				}()
+			}
+		}
 		current, err := d.InspectRealm(ctx, nil)
 		if err != nil {
 			return err
diff --git a/sql/postgres/driver_oss.go b/sql/postgres/driver_oss.go
index a987e71d..f6b498af 100644
--- a/sql/postgres/driver_oss.go
+++ b/sql/postgres/driver_oss.go
@@ -345,3 +345,7 @@ func uniqueConstChanged(_, _ []schema.Attr) bool {
 	// to keep BC with old versions.
 	return false
 }
+
+func detachCycles(changes []schema.Change) ([]schema.Change, error) {
+	return sqlx.DetachCycles(changes)
+}
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 8b8e66d8..6ad94a4b 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -63,7 +63,6 @@ type state struct {
 	*conn
 	migrate.Plan
 	migrate.PlanOptions
-	droppedT []*schema.Table
 }
 
 // Exec executes the changes on the database. An error is returned
@@ -78,7 +77,7 @@ func (s *state) plan(changes []schema.Change) error {
 	if err != nil {
 		return err
 	}
-	if planned, err = sqlx.DetachCycles(planned); err != nil {
+	if planned, err = detachCycles(planned); err != nil {
 		return err
 	}
 	for _, c := range sqlx.SortChanges(planned) {
@@ -287,7 +286,6 @@ func (s *state) addTable(add *schema.AddTable) error {
 // dropTable builds and executes the query for dropping a table from a schema.
 func (s *state) dropTable(drop *schema.DropTable) error {
 	cmd := &changeGroup{}
-	s.droppedT = append(s.droppedT, drop.T)
 	rs := &state{
 		conn:        s.conn,
 		PlanOptions: s.PlanOptions,
@@ -561,7 +559,11 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 				}
 				reverse = append(reverse, &schema.DropForeignKey{F: change.F})
 			case *schema.DropForeignKey:
-				b.P("DROP CONSTRAINT").Ident(change.F.Symbol)
+				b.P("DROP CONSTRAINT")
+				if sqlx.Has(change.Extra, &schema.IfExists{}) {
+					b.P("IF EXISTS")
+				}
+				b.Ident(change.F.Symbol)
 				reverse = append(reverse, &schema.AddForeignKey{F: change.F})
 			case *schema.AddCheck:
 				check(b.P("ADD"), change.C)
diff --git a/sql/schema/migrate.go b/sql/schema/migrate.go
index 78c7ef24..6335c40e 100644
--- a/sql/schema/migrate.go
+++ b/sql/schema/migrate.go
@@ -281,6 +281,7 @@ type (
 	// DropForeignKey describes a foreign-key removal change.
 	DropForeignKey struct {
 		F *ForeignKey
+		Extra []Clause // Extra clauses and options.
 	}
 
 	// ModifyForeignKey describes a change that modifies a foreign-key.


Test patch:
diff --git a/sql/internal/sqlx/plan_test.go b/sql/internal/sqlx/plan_test.go
index 90ab3849..297e4b06 100644
--- a/sql/internal/sqlx/plan_test.go
+++ b/sql/internal/sqlx/plan_test.go
@@ -238,3 +238,16 @@ func TestSortChanges(t *testing.T) {
 	})
 	require.Equal(t, []schema.Change{&schema.DropFunc{F: f2}, &schema.AddFunc{F: f1}}, planned)
 }
+
+func TestSameTable(t *testing.T) {
+	t1 := schema.NewTable("t1")
+	require.True(t, SameTable(t1, t1))
+	require.True(t, SameTable(t1, schema.NewTable("t1")), "same copy")
+	require.False(t, SameTable(t1, schema.NewTable("t2")))
+
+	t1.SetSchema(schema.New("public"))
+	require.True(t, SameTable(t1, t1))
+	require.True(t, SameTable(t1, schema.NewTable("t1").SetSchema(schema.New("public"))), "same copy")
+	require.False(t, SameTable(t1, schema.NewTable("t1")))
+	require.False(t, SameTable(t1, schema.NewTable("t1").SetSchema(schema.New("private"))))
+}
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 470d6bcb..686dac8a 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -96,6 +96,8 @@ func TestPlanChanges(t *testing.T) {
 						Columns:    []*schema.Column{pets.Columns[1]},
 						RefColumns: []*schema.Column{users.Columns[0]},
 					}
+					fk1 := *fk
+					fk1.Symbol += "1"
 					pets.ForeignKeys = []*schema.ForeignKey{fk}
 					return &schema.ModifyTable{
 						T: pets,
@@ -103,6 +105,12 @@ func TestPlanChanges(t *testing.T) {
 							&schema.DropForeignKey{
 								F: fk,
 							},
+							&schema.DropForeignKey{
+								F: &fk1,
+								Extra: []schema.Clause{
+									&schema.IfExists{},
+								},
+							},
 						},
 					}
 				}(),
@@ -112,8 +120,8 @@ func TestPlanChanges(t *testing.T) {
 				Transactional: true,
 				Changes: []*migrate.Change{
 					{
-						Cmd:     `ALTER TABLE "pets" DROP CONSTRAINT "pets_user_id_fkey"`,
-						Reverse: `ALTER TABLE "pets" ADD CONSTRAINT "pets_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "users" ("id") ON UPDATE NO ACTION ON DELETE CASCADE`,
+						Cmd:     `ALTER TABLE "pets" DROP CONSTRAINT "pets_user_id_fkey", DROP CONSTRAINT IF EXISTS "pets_user_id_fkey1"`,
+						Reverse: `ALTER TABLE "pets" ADD CONSTRAINT "pets_user_id_fkey1" FOREIGN KEY ("user_id") REFERENCES "users" ("id") ON UPDATE NO ACTION ON DELETE CASCADE, ADD CONSTRAINT "pets_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "users" ("id") ON UPDATE NO ACTION ON DELETE CASCADE`,
 					},
 				},
 			},