Repository: ariga/atlas
Pull request #1566 (2023-03-30 09:52:22)

sql/postgres: drop constraints first before adding new
Fixed https://github.com/ariga/atlas/issues/1551


Patch:
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 3d32ab97..54c9a191 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -8,6 +8,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sort"
 	"strconv"
 	"strings"
 
@@ -283,6 +284,14 @@ func (s *state) modifyTable(ctx context.Context, modify *schema.ModifyTable) err
 			} else {
 				dropI = append(dropI, change)
 			}
+		case *schema.ModifyPrimaryKey:
+			// Primary key modification needs to be split into "Drop" and "Add"
+			// because the new key may include columns that have not been added yet.
+			alter = append(alter, &schema.DropPrimaryKey{
+				P: change.From,
+			}, &schema.AddPrimaryKey{
+				P: change.To,
+			})
 		case *schema.ModifyIndex:
 			k := change.Change
 			if change.Change.Is(schema.ChangeComment) {
@@ -391,6 +400,10 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 		reverse    []schema.Change
 		reversible = true
 	)
+	// Constraints drop should be executed first.
+	sort.SliceStable(changes, func(i, j int) bool {
+		return dropConst(changes[i]) && !dropConst(changes[j])
+	})
 	build := func(alter *changeGroup, changes []schema.Change) (string, error) {
 		b := s.Build("ALTER TABLE").Table(t)
 		err := b.MapCommaErr(changes, func(i int, b *sqlx.Builder) error {
@@ -448,13 +461,6 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 			case *schema.DropPrimaryKey:
 				b.P("DROP CONSTRAINT").Ident(pkName(t, change.P))
 				reverse = append(reverse, &schema.AddPrimaryKey{P: change.P})
-			case *schema.ModifyPrimaryKey:
-				b.P("DROP CONSTRAINT").Ident(pkName(t, change.From))
-				b.P(", ADD PRIMARY KEY")
-				if err := s.index(b, change.To); err != nil {
-					return err
-				}
-				reverse = append(reverse, &schema.ModifyPrimaryKey{From: change.To, To: change.From, Change: change.Change})
 			case *schema.AddForeignKey:
 				b.P("ADD")
 				s.fks(b, change.F)
@@ -1325,3 +1331,13 @@ func pkName(t *schema.Table, pk *schema.Index) string {
 	// See: the ChooseIndexName function in PostgreSQL for more reference.
 	return t.Name + "_pkey"
 }
+
+// dropConst indicates if the given change is a constraint drop.
+func dropConst(c schema.Change) bool {
+	switch c.(type) {
+	case *schema.DropIndex, *schema.DropPrimaryKey, *schema.DropCheck, *schema.DropForeignKey:
+		return true
+	default:
+		return false
+	}
+}


Test patch:
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 96c67ac0..9ce97daa 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -440,18 +440,18 @@ func TestPlanChanges(t *testing.T) {
 		{
 			changes: []schema.Change{
 				func() schema.Change {
+					dropC, addC := schema.NewStringColumn("email", "text"), schema.NewStringColumn("name", "text")
 					users := schema.NewTable("users").
-						AddColumns(
-							schema.NewStringColumn("id", "varchar(255)"),
-							schema.NewStringColumn("name", "varchar(255)"),
-						)
-					users.SetPrimaryKey(schema.NewPrimaryKey(users.Columns[:1]...))
+						AddColumns(schema.NewStringColumn("id", "text"), addC, schema.NewStringColumn("last", "text"))
+					users.SetPrimaryKey(schema.NewPrimaryKey(addC))
 					return &schema.ModifyTable{
 						T: users,
 						Changes: []schema.Change{
+							&schema.DropColumn{C: dropC},
+							&schema.AddColumn{C: addC},
 							&schema.ModifyPrimaryKey{
-								From: schema.NewPrimaryKey(users.Columns[:1]...).
-									AddAttrs(&IndexInclude{Columns: users.Columns[1:]}),
+								From: schema.NewPrimaryKey(dropC).
+									AddAttrs(&IndexInclude{Columns: users.Columns[2:]}),
 								To: users.PrimaryKey,
 							},
 						},
@@ -463,8 +463,8 @@ func TestPlanChanges(t *testing.T) {
 				Transactional: true,
 				Changes: []*migrate.Change{
 					{
-						Cmd:     `ALTER TABLE "users" DROP CONSTRAINT "users_pkey" , ADD PRIMARY KEY ("id")`,
-						Reverse: `ALTER TABLE "users" DROP CONSTRAINT "users_pkey" , ADD PRIMARY KEY ("id") INCLUDE ("name")`,
+						Cmd:     `ALTER TABLE "users" DROP CONSTRAINT "users_pkey", DROP COLUMN "email", ADD COLUMN "name" text NOT NULL, ADD PRIMARY KEY ("name")`,
+						Reverse: `ALTER TABLE "users" DROP CONSTRAINT "users_pkey", DROP COLUMN "name", ADD COLUMN "email" text NOT NULL, ADD PRIMARY KEY ("email") INCLUDE ("last")`,
 					},
 				},
 			},
@@ -580,8 +580,8 @@ func TestPlanChanges(t *testing.T) {
 						Reverse: `CREATE INDEX CONCURRENTLY "drop_con" ON "users" ("id")`,
 					},
 					{
-						Cmd:     `ALTER TABLE "users" ADD COLUMN "name" character varying(255) NOT NULL DEFAULT 'logged_in', ADD COLUMN "last" character varying(255) NOT NULL DEFAULT 'logged_in', ADD CONSTRAINT "name_not_empty" CHECK ("name" <> ''), DROP CONSTRAINT "id_nonzero", DROP CONSTRAINT "id_iseven", ADD CONSTRAINT "id_iseven" CHECK (("id") % 2 = 0)`,
-						Reverse: `ALTER TABLE "users" DROP CONSTRAINT "id_iseven", ADD CONSTRAINT "id_iseven" CHECK ("id" % 2 = 0), ADD CONSTRAINT "id_nonzero" CHECK ("id" <> 0), DROP CONSTRAINT "name_not_empty", DROP COLUMN "last", DROP COLUMN "name"`,
+						Cmd:     `ALTER TABLE "users" DROP CONSTRAINT "id_nonzero", ADD COLUMN "name" character varying(255) NOT NULL DEFAULT 'logged_in', ADD COLUMN "last" character varying(255) NOT NULL DEFAULT 'logged_in', ADD CONSTRAINT "name_not_empty" CHECK ("name" <> ''), DROP CONSTRAINT "id_iseven", ADD CONSTRAINT "id_iseven" CHECK (("id") % 2 = 0)`,
+						Reverse: `ALTER TABLE "users" DROP CONSTRAINT "id_iseven", ADD CONSTRAINT "id_iseven" CHECK ("id" % 2 = 0), DROP CONSTRAINT "name_not_empty", DROP COLUMN "last", DROP COLUMN "name", ADD CONSTRAINT "id_nonzero" CHECK ("id" <> 0)`,
 					},
 					{
 						Cmd:     `CREATE INDEX "id_key" ON "users" ("id" DESC) WHERE success`,