Repository: ariga/atlas
Pull request #1402 (2023-01-15 09:36:43)

sql: add support for diffing primary-key changes


Patch:
diff --git a/sql/internal/sqlx/diff.go b/sql/internal/sqlx/diff.go
index 39f70a3d..8281f552 100644
--- a/sql/internal/sqlx/diff.go
+++ b/sql/internal/sqlx/diff.go
@@ -156,10 +156,6 @@ func (d *Diff) TableDiff(from, to *schema.Table) ([]schema.Change, error) {
 	if from.Name != to.Name {
 		return nil, fmt.Errorf("mismatched table names: %q != %q", from.Name, to.Name)
 	}
-	// PK modification is not supported.
-	if pk1, pk2 := from.PrimaryKey, to.PrimaryKey; (pk1 != nil) != (pk2 != nil) || (pk1 != nil) && d.pkChange(pk1, pk2) != schema.NoChange {
-		return nil, fmt.Errorf("changing %q table primary key is not supported", to.Name)
-	}
 
 	// Drop or modify attributes (collations, checks, etc).
 	change, err := d.TableAttrDiff(from, to)
@@ -194,7 +190,8 @@ func (d *Diff) TableDiff(from, to *schema.Table) ([]schema.Change, error) {
 		}
 	}
 
-	// Index changes.
+	// Primary-key and index changes.
+	changes = append(changes, d.pkDiff(from, to)...)
 	changes = append(changes, d.indexDiff(from, to)...)
 
 	// Drop or modify foreign-keys.
@@ -221,6 +218,28 @@ func (d *Diff) TableDiff(from, to *schema.Table) ([]schema.Change, error) {
 	return changes, nil
 }
 
+// pkDiff returns the schema changes (if any) for migrating table
+// primary-key from current state to the desired state.
+func (d *Diff) pkDiff(from, to *schema.Table) (changes []schema.Change) {
+	switch pk1, pk2 := from.PrimaryKey, to.PrimaryKey; {
+	case pk1 == nil && pk2 != nil:
+		changes = append(changes, &schema.AddPrimaryKey{P: pk2})
+	case pk1 != nil && pk2 == nil:
+		changes = append(changes, &schema.DropPrimaryKey{P: pk1})
+	case pk1 != nil && pk2 != nil:
+		change := d.indexChange(pk1, pk2)
+		change &= ^schema.ChangeUnique
+		if change != schema.NoChange {
+			changes = append(changes, &schema.ModifyPrimaryKey{
+				From:   pk1,
+				To:     pk2,
+				Change: change,
+			})
+		}
+	}
+	return
+}
+
 // indexDiff returns the schema changes (if any) for migrating table
 // indexes from current state to the desired state.
 func (d *Diff) indexDiff(from, to *schema.Table) []schema.Change {
@@ -265,12 +284,6 @@ func (d *Diff) indexDiff(from, to *schema.Table) []schema.Change {
 	return changes
 }
 
-// pkChange returns the schema changes (if any) for migrating one primary key to the other.
-func (d *Diff) pkChange(from, to *schema.Index) schema.ChangeKind {
-	change := d.indexChange(from, to)
-	return change & ^schema.ChangeUnique
-}
-
 // indexChange returns the schema changes (if any) for migrating one index to the other.
 func (d *Diff) indexChange(from, to *schema.Index) schema.ChangeKind {
 	var change schema.ChangeKind
diff --git a/sql/mysql/inspect.go b/sql/mysql/inspect.go
index c22864d4..38ac2fbb 100644
--- a/sql/mysql/inspect.go
+++ b/sql/mysql/inspect.go
@@ -835,7 +835,7 @@ type (
 		T string // BTREE, HASH, FULLTEXT, SPATIAL, RTREE
 	}
 
-	// BitType represents a bit type.
+	// BitType represents the type bit.
 	BitType struct {
 		schema.Type
 		T    string
diff --git a/sql/schema/migrate.go b/sql/schema/migrate.go
index d41ffeb9..7bc19dd4 100644
--- a/sql/schema/migrate.go
+++ b/sql/schema/migrate.go
@@ -128,6 +128,22 @@ type (
 		From, To *Index
 	}
 
+	// AddPrimaryKey describes a primary-key creation change.
+	AddPrimaryKey struct {
+		P *Index
+	}
+
+	// DropPrimaryKey describes a primary-key removal change.
+	DropPrimaryKey struct {
+		P *Index
+	}
+
+	// ModifyPrimaryKey describes a primary-key modification.
+	ModifyPrimaryKey struct {
+		From, To *Index
+		Change   ChangeKind
+	}
+
 	// AddForeignKey describes a foreign-key creation change.
 	AddForeignKey struct {
 		F *ForeignKey
@@ -382,6 +398,9 @@ func (*AddIndex) change()         {}
 func (*DropIndex) change()        {}
 func (*ModifyIndex) change()      {}
 func (*RenameIndex) change()      {}
+func (*AddPrimaryKey) change()    {}
+func (*DropPrimaryKey) change()   {}
+func (*ModifyPrimaryKey) change() {}
 func (*AddCheck) change()         {}
 func (*DropCheck) change()        {}
 func (*ModifyCheck) change()      {}


Test patch:
diff --git a/sql/mysql/diff_test.go b/sql/mysql/diff_test.go
index 14e6e50a..495ff385 100644
--- a/sql/mysql/diff_test.go
+++ b/sql/mysql/diff_test.go
@@ -37,18 +37,6 @@ func TestDiff_TableDiff(t *testing.T) {
 			from: &schema.Table{Name: "users", Schema: &schema.Schema{Name: "public"}, Columns: []*schema.Column{{Name: "enum", Default: &schema.RawExpr{X: "'A'"}, Type: &schema.ColumnType{Type: &schema.EnumType{Values: []string{"A"}}}}}},
 			to:   &schema.Table{Name: "users", Columns: []*schema.Column{{Name: "enum", Default: &schema.RawExpr{X: `"A"`}, Type: &schema.ColumnType{Type: &schema.EnumType{Values: []string{"A"}}}}}},
 		},
-		{
-			name: "change primary key",
-			from: func() *schema.Table {
-				t := &schema.Table{Name: "users", Schema: &schema.Schema{Name: "public"}, Columns: []*schema.Column{{Name: "id", Type: &schema.ColumnType{Raw: "int", Type: &schema.IntegerType{T: "int"}}}}}
-				t.PrimaryKey = &schema.Index{
-					Parts: []*schema.IndexPart{{C: t.Columns[0]}},
-				}
-				return t
-			}(),
-			to:      &schema.Table{Name: "users"},
-			wantErr: true,
-		},
 		{
 			name: "modify counter",
 			from: &schema.Table{Name: "users", Schema: &schema.Schema{Name: "public"}, Attrs: []schema.Attr{&AutoIncrement{V: 1}}},
@@ -426,6 +414,65 @@ func TestDiff_TableDiff(t *testing.T) {
 				},
 			}
 		}(),
+		func() testcase {
+			from := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewIntColumn("id", "int"))
+			to := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewIntColumn("id", "int"))
+			to.SetPrimaryKey(schema.NewPrimaryKey(to.Columns...))
+			return testcase{
+				name: "add primary-key",
+				from: from,
+				to:   to,
+				wantChanges: []schema.Change{
+					&schema.AddPrimaryKey{P: to.PrimaryKey},
+				},
+			}
+		}(),
+		func() testcase {
+			from := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewIntColumn("id", "int"))
+			from.SetPrimaryKey(schema.NewPrimaryKey(from.Columns...))
+			to := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewIntColumn("id", "int"))
+			return testcase{
+				name: "drop primary-key",
+				from: from,
+				to:   to,
+				wantChanges: []schema.Change{
+					&schema.DropPrimaryKey{P: from.PrimaryKey},
+				},
+			}
+		}(),
+		func() testcase {
+			from := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewStringColumn("id", "varchar(255)"))
+			from.SetPrimaryKey(schema.NewPrimaryKey(from.Columns...))
+			to := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewStringColumn("id", "varchar(255)"))
+			to.SetPrimaryKey(
+				schema.NewPrimaryKey(from.Columns...).
+					AddAttrs(&IndexType{T: IndexTypeHash}),
+			)
+			return testcase{
+				name: "modify primary-key",
+				from: from,
+				to:   to,
+				wantChanges: []schema.Change{
+					&schema.ModifyPrimaryKey{
+						From:   from.PrimaryKey,
+						To:     to.PrimaryKey,
+						Change: schema.ChangeAttr,
+					},
+				},
+			}
+		}(),
 		func() testcase {
 			var (
 				ref = &schema.Table{
diff --git a/sql/postgres/diff_test.go b/sql/postgres/diff_test.go
index d438b7a7..dd320b54 100644
--- a/sql/postgres/diff_test.go
+++ b/sql/postgres/diff_test.go
@@ -27,18 +27,6 @@ func TestDiff_TableDiff(t *testing.T) {
 			from: &schema.Table{Name: "users", Schema: &schema.Schema{Name: "public"}},
 			to:   &schema.Table{Name: "users"},
 		},
-		{
-			name: "change primary key columns",
-			from: func() *schema.Table {
-				t := &schema.Table{Name: "users", Schema: &schema.Schema{Name: "public"}, Columns: []*schema.Column{{Name: "id", Type: &schema.ColumnType{Raw: "int", Type: &schema.IntegerType{T: "int"}}}}}
-				t.PrimaryKey = &schema.Index{
-					Parts: []*schema.IndexPart{{C: t.Columns[0]}},
-				}
-				return t
-			}(),
-			to:      &schema.Table{Name: "users"},
-			wantErr: true,
-		},
 		{
 			name: "change identity attributes",
 			from: func() *schema.Table {
@@ -412,6 +400,31 @@ func TestDiff_TableDiff(t *testing.T) {
 				},
 			}
 		}(),
+		func() testcase {
+			from := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewStringColumn("id", "varchar"), schema.NewBoolColumn("active", "bool"))
+			from.SetPrimaryKey(schema.NewPrimaryKey(from.Columns...))
+			to := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewStringColumn("id", "varchar"), schema.NewBoolColumn("active", "bool"))
+			to.SetPrimaryKey(
+				schema.NewPrimaryKey(from.Columns...).
+					AddAttrs(&IndexPredicate{P: "active"}),
+			)
+			return testcase{
+				name: "modify primary-key",
+				from: from,
+				to:   to,
+				wantChanges: []schema.Change{
+					&schema.ModifyPrimaryKey{
+						From:   from.PrimaryKey,
+						To:     to.PrimaryKey,
+						Change: schema.ChangeAttr,
+					},
+				},
+			}
+		}(),
 		func() testcase {
 			var (
 				ref = &schema.Table{
diff --git a/sql/sqlite/diff_test.go b/sql/sqlite/diff_test.go
index dc8e8462..ef5b1863 100644
--- a/sql/sqlite/diff_test.go
+++ b/sql/sqlite/diff_test.go
@@ -26,18 +26,6 @@ func TestDiff_TableDiff(t *testing.T) {
 			from: &schema.Table{Name: "users", Schema: &schema.Schema{Name: "public"}},
 			to:   &schema.Table{Name: "users"},
 		},
-		{
-			name: "change primary key",
-			from: func() *schema.Table {
-				t := &schema.Table{Name: "users", Schema: &schema.Schema{Name: "public"}, Columns: []*schema.Column{{Name: "id", Type: &schema.ColumnType{Raw: "int", Type: &schema.IntegerType{T: "int"}}}}}
-				t.PrimaryKey = &schema.Index{
-					Parts: []*schema.IndexPart{{C: t.Columns[0]}},
-				}
-				return t
-			}(),
-			to:      &schema.Table{Name: "users"},
-			wantErr: true,
-		},
 		{
 			name: "add attr",
 			from: &schema.Table{Name: "t1", Schema: &schema.Schema{Name: "public"}},
@@ -228,6 +216,31 @@ func TestDiff_TableDiff(t *testing.T) {
 				},
 			}
 		}(),
+		func() testcase {
+			from := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewStringColumn("id", "varchar"), schema.NewBoolColumn("active", "bool"))
+			from.SetPrimaryKey(schema.NewPrimaryKey(from.Columns...))
+			to := schema.NewTable("t1").
+				SetSchema(schema.New("test")).
+				AddColumns(schema.NewStringColumn("id", "varchar"), schema.NewBoolColumn("active", "bool"))
+			to.SetPrimaryKey(
+				schema.NewPrimaryKey(from.Columns...).
+					AddAttrs(&IndexPredicate{P: "active"}),
+			)
+			return testcase{
+				name: "modify primary-key",
+				from: from,
+				to:   to,
+				wantChanges: []schema.Change{
+					&schema.ModifyPrimaryKey{
+						From:   from.PrimaryKey,
+						To:     to.PrimaryKey,
+						Change: schema.ChangeAttr,
+					},
+				},
+			}
+		}(),
 		func() testcase {
 			var (
 				ref = &schema.Table{