Repository: ariga/atlas
Pull request #1403 (2023-01-16 04:39:19)

sql/mysql: support planning pk changes


Patch:
diff --git a/sql/mysql/migrate.go b/sql/mysql/migrate.go
index 7bcf56f1..bc147d49 100644
--- a/sql/mysql/migrate.go
+++ b/sql/mysql/migrate.go
@@ -232,7 +232,7 @@ func (s *state) addTable(add *schema.AddTable) error {
 		})
 		if pk := add.T.PrimaryKey; pk != nil {
 			b.Comma().P("PRIMARY KEY")
-			indexParts(b, pk.Parts)
+			indexTypeParts(b, pk)
 		}
 		if len(add.T.Indexes) > 0 {
 			b.Comma()
@@ -325,6 +325,14 @@ func (s *state) modifyTable(modify *schema.ModifyTable) error {
 			changes[1] = append(changes[1], &schema.AddIndex{
 				I: change.To,
 			})
+		// Primary-key modification requires rebuilding the index.
+		case *schema.ModifyPrimaryKey:
+			changes[0] = append(changes[0], &schema.DropPrimaryKey{
+				P: change.From,
+			})
+			changes[1] = append(changes[1], &schema.AddPrimaryKey{
+				P: change.To,
+			})
 		case *schema.DropAttr:
 			return fmt.Errorf("unsupported change type: %v", change.A)
 		default:
@@ -394,6 +402,13 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 			case *schema.DropIndex:
 				b.P("DROP INDEX").Ident(change.I.Name)
 				reverse = append(reverse, &schema.AddIndex{I: change.I})
+			case *schema.AddPrimaryKey:
+				b.P("ADD PRIMARY KEY")
+				indexTypeParts(b, change.P)
+				reverse = append(reverse, &schema.DropPrimaryKey{P: change.P})
+			case *schema.DropPrimaryKey:
+				b.P("DROP PRIMARY KEY")
+				reverse = append(reverse, &schema.AddPrimaryKey{P: change.P})
 			case *schema.AddForeignKey:
 				b.P("ADD")
 				if err := s.fks(b, change.F); err != nil {
@@ -555,41 +570,37 @@ func (s *state) column(b *sqlx.Builder, t *schema.Table, c *schema.Column) error
 }
 
 func index(b *sqlx.Builder, idx *schema.Index) {
-	var t IndexType
-	if sqlx.Has(idx.Attrs, &t) {
-		t.T = strings.ToUpper(t.T)
-	}
-	switch {
+	switch t := indexType(idx.Attrs); {
 	case idx.Unique:
 		b.P("UNIQUE")
 	case t.T == IndexTypeFullText || t.T == IndexTypeSpatial:
 		b.P(t.T)
 	}
 	b.P("INDEX").Ident(idx.Name)
-	// Skip BTREE as it is the default type.
-	if t.T == IndexTypeHash {
-		b.P("USING", t.T)
-	}
-	indexParts(b, idx.Parts)
+	indexTypeParts(b, idx)
 	if c := (schema.Comment{}); sqlx.Has(idx.Attrs, &c) {
 		b.P("COMMENT", quote(c.Text))
 	}
 }
 
-func indexParts(b *sqlx.Builder, parts []*schema.IndexPart) {
+func indexTypeParts(b *sqlx.Builder, idx *schema.Index) {
+	// Skip BTREE as it is the default type.
+	if t := indexType(idx.Attrs); t.T == IndexTypeHash {
+		b.P("USING", t.T)
+	}
 	b.Wrap(func(b *sqlx.Builder) {
-		b.MapComma(parts, func(i int, b *sqlx.Builder) {
-			switch part := parts[i]; {
+		b.MapComma(idx.Parts, func(i int, b *sqlx.Builder) {
+			switch part := idx.Parts[i]; {
 			case part.C != nil:
 				b.Ident(part.C.Name)
 			case part.X != nil:
 				b.WriteString(sqlx.MayWrap(part.X.(*schema.RawExpr).X))
 			}
-			if s := (&SubPart{}); sqlx.Has(parts[i].Attrs, s) {
+			if s := (&SubPart{}); sqlx.Has(idx.Parts[i].Attrs, s) {
 				b.WriteString(fmt.Sprintf("(%d)", s.Len))
 			}
 			// Ignore default collation (i.e. "ASC")
-			if parts[i].Desc {
+			if idx.Parts[i].Desc {
 				b.P("DESC")
 			}
 		})
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index 1f327b19..8d91b951 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -80,6 +80,7 @@ var (
 	hclState = schemahcl.New(
 		schemahcl.WithTypes(TypeRegistry.Specs()),
 		schemahcl.WithScopedEnums("table.index.type", IndexTypeBTree, IndexTypeHash, IndexTypeFullText, IndexTypeSpatial),
+		schemahcl.WithScopedEnums("table.primary_key.type", IndexTypeBTree, IndexTypeHash, IndexTypeFullText, IndexTypeSpatial),
 		schemahcl.WithScopedEnums("table.column.as.type", stored, persistent, virtual),
 		schemahcl.WithScopedEnums("table.foreign_key.on_update", specutil.ReferenceVars...),
 		schemahcl.WithScopedEnums("table.foreign_key.on_delete", specutil.ReferenceVars...),
@@ -100,7 +101,7 @@ var (
 // ForeignKeySpecs into ForeignKeys, as the target tables do not necessarily exist in the schema
 // at this point. Instead, the linking is done by the convertSchema function.
 func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, error) {
-	t, err := specutil.Table(spec, parent, convertColumn, specutil.PrimaryKey, convertIndex, convertCheck)
+	t, err := specutil.Table(spec, parent, convertColumn, convertPK, convertIndex, convertCheck)
 	if err != nil {
 		return nil, err
 	}
@@ -119,20 +120,39 @@ func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, er
 	return t, err
 }
 
+// convertPK converts a sqlspec.PrimaryKey into a schema.Index.
+func convertPK(spec *sqlspec.PrimaryKey, parent *schema.Table) (*schema.Index, error) {
+	idx, err := specutil.PrimaryKey(spec, parent)
+	if err != nil {
+		return nil, err
+	}
+	if err := convertIndexType(idx, spec); err != nil {
+		return nil, err
+	}
+	return idx, nil
+}
+
 // convertIndex converts a sqlspec.Index into a schema.Index.
 func convertIndex(spec *sqlspec.Index, parent *schema.Table) (*schema.Index, error) {
 	idx, err := specutil.Index(spec, parent, convertPart)
 	if err != nil {
 		return nil, err
 	}
+	if err := convertIndexType(idx, spec); err != nil {
+		return nil, err
+	}
+	return idx, nil
+}
+
+func convertIndexType(idx *schema.Index, spec specutil.Attrer) error {
 	if attr, ok := spec.Attr("type"); ok {
 		t, err := attr.String()
 		if err != nil {
-			return nil, err
+			return err
 		}
 		idx.AddAttrs(&IndexType{T: t})
 	}
-	return idx, nil
+	return nil
 }
 
 func convertPart(spec *sqlspec.IndexPart, part *schema.IndexPart) error {
@@ -221,7 +241,7 @@ func tableSpec(t *schema.Table) (*sqlspec.Table, error) {
 	ts, err := specutil.FromTable(
 		t,
 		columnSpec,
-		specutil.FromPrimaryKey,
+		pkSpec,
 		indexSpec,
 		specutil.FromForeignKey,
 		checkSpec,
@@ -238,16 +258,30 @@ func tableSpec(t *schema.Table) (*sqlspec.Table, error) {
 	return ts, nil
 }
 
+func pkSpec(idx *schema.Index) (*sqlspec.PrimaryKey, error) {
+	spec, err := specutil.FromPrimaryKey(idx)
+	if err != nil {
+		return nil, err
+	}
+	spec.Extra.Attrs = indexTypeSpec(idx, spec.Extra.Attrs)
+	return spec, nil
+}
+
 func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {
 	spec, err := specutil.FromIndex(idx, partAttr)
 	if err != nil {
 		return nil, err
 	}
+	spec.Extra.Attrs = indexTypeSpec(idx, spec.Extra.Attrs)
+	return spec, nil
+}
+
+func indexTypeSpec(idx *schema.Index, attrs []*schemahcl.Attr) []*schemahcl.Attr {
 	// Avoid printing the index type if it is the default.
 	if i := (IndexType{}); sqlx.Has(idx.Attrs, &i) && i.T != IndexTypeBTree {
-		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("type", strings.ToUpper(i.T)))
+		attrs = append(attrs, specutil.VarAttr("type", strings.ToUpper(i.T)))
 	}
-	return spec, nil
+	return attrs
 }
 
 func partAttr(_ *schema.Index, part *schema.IndexPart, spec *sqlspec.IndexPart) error {


Test patch:
diff --git a/internal/integration/testdata/mysql/primary-key.txt b/internal/integration/testdata/mysql/primary-key.txt
new file mode 100644
index 00000000..1f054b29
--- /dev/null
+++ b/internal/integration/testdata/mysql/primary-key.txt
@@ -0,0 +1,99 @@
+only mysql8
+
+# Create table.
+apply 1.hcl
+cmpshow users 1.sql
+
+# Add a primary-key.
+apply 2.hcl
+cmpshow users 2.sql
+
+# Modify the primary-key column and type.
+apply 3.hcl
+cmpshow users 3.sql
+
+# Drop the primary-key.
+apply 1.hcl
+cmpshow users 1.sql
+
+-- 1.hcl --
+schema "$db" {
+  charset = "utf8mb4"
+  collate = "utf8mb4_general_ci"
+}
+
+table "users" {
+  schema = schema.$db
+  column "id" {
+    null = false
+    type = varchar(128)
+  }
+  column "oid" {
+    null = false
+    type = varchar(128)
+  }
+}
+
+-- 1.sql --
+CREATE TABLE `users` (
+  `id` varchar(128) COLLATE utf8mb4_general_ci NOT NULL,
+  `oid` varchar(128) COLLATE utf8mb4_general_ci NOT NULL
+)
+
+-- 2.hcl --
+schema "$db" {
+  charset = "utf8mb4"
+  collate = "utf8mb4_general_ci"
+}
+
+table "users" {
+  schema = schema.$db
+  column "id" {
+    null = false
+    type = varchar(128)
+  }
+  column "oid" {
+    null = false
+    type = varchar(128)
+  }
+  primary_key {
+    columns = [column.id]
+  }
+}
+
+-- 2.sql --
+CREATE TABLE `users` (
+  `id` varchar(128) COLLATE utf8mb4_general_ci NOT NULL,
+  `oid` varchar(128) COLLATE utf8mb4_general_ci NOT NULL,
+  PRIMARY KEY (`id`)
+)
+
+
+-- 3.hcl --
+schema "$db" {
+  charset = "utf8mb4"
+  collate = "utf8mb4_general_ci"
+}
+
+table "users" {
+  schema = schema.$db
+  column "id" {
+    null = false
+    type = varchar(128)
+  }
+  column "oid" {
+    null = false
+    type = varchar(128)
+  }
+  primary_key {
+    columns = [column.oid]
+    type    = HASH
+  }
+}
+
+-- 3.sql --
+CREATE TABLE `users` (
+  `id` varchar(128) COLLATE utf8mb4_general_ci NOT NULL,
+  `oid` varchar(128) COLLATE utf8mb4_general_ci NOT NULL,
+  PRIMARY KEY (`oid`)
+)
diff --git a/sql/mysql/migrate_test.go b/sql/mysql/migrate_test.go
index 8e2c9533..3ab08f6f 100644
--- a/sql/mysql/migrate_test.go
+++ b/sql/mysql/migrate_test.go
@@ -264,6 +264,109 @@ func TestPlanChanges(t *testing.T) {
 				},
 			},
 		},
+		// Drop a primary key.
+		{
+			changes: []schema.Change{
+				func() schema.Change {
+					users := schema.NewTable("users").
+						SetSchema(schema.New("test")).
+						AddColumns(
+							schema.NewIntColumn("id", "bigint"),
+						)
+					return &schema.ModifyTable{
+						T: users,
+						Changes: []schema.Change{
+							&schema.DropPrimaryKey{
+								P: schema.NewPrimaryKey(users.Columns...),
+							},
+						},
+					}
+				}(),
+			},
+			wantPlan: &migrate.Plan{
+				Reversible: true,
+				Changes: []*migrate.Change{
+					{
+						Cmd:     "ALTER TABLE `users` DROP PRIMARY KEY",
+						Reverse: "ALTER TABLE `users` ADD PRIMARY KEY (`id`)",
+					},
+				},
+			},
+			options: []migrate.PlanOption{
+				func(o *migrate.PlanOptions) { o.SchemaQualifier = new(string) },
+			},
+		},
+		// Add a primary key.
+		{
+			changes: []schema.Change{
+				func() schema.Change {
+					users := schema.NewTable("users").
+						SetSchema(schema.New("test")).
+						AddColumns(
+							schema.NewIntColumn("id", "bigint"),
+						)
+					users.SetPrimaryKey(schema.NewPrimaryKey(users.Columns...))
+					return &schema.ModifyTable{
+						T: users,
+						Changes: []schema.Change{
+							&schema.AddPrimaryKey{
+								P: users.PrimaryKey,
+							},
+						},
+					}
+				}(),
+			},
+			wantPlan: &migrate.Plan{
+				Reversible: true,
+				Changes: []*migrate.Change{
+					{
+						Cmd:     "ALTER TABLE `dev`.`users` ADD PRIMARY KEY (`id`)",
+						Reverse: "ALTER TABLE `dev`.`users` DROP PRIMARY KEY",
+					},
+				},
+			},
+			options: []migrate.PlanOption{
+				func(o *migrate.PlanOptions) {
+					dev := "dev"
+					o.SchemaQualifier = &dev
+				},
+			},
+		},
+		// Modify a primary key.
+		{
+			changes: []schema.Change{
+				func() schema.Change {
+					users := schema.NewTable("users").
+						AddColumns(
+							schema.NewStringColumn("id", "varchar(255)"),
+						)
+					users.SetPrimaryKey(schema.NewPrimaryKey(users.Columns...))
+					return &schema.ModifyTable{
+						T: users,
+						Changes: []schema.Change{
+							&schema.ModifyPrimaryKey{
+								From: schema.NewPrimaryKey(users.Columns...).
+									AddAttrs(&IndexType{T: IndexTypeHash}),
+								To: users.PrimaryKey,
+							},
+						},
+					}
+				}(),
+			},
+			wantPlan: &migrate.Plan{
+				Reversible: true,
+				Changes: []*migrate.Change{
+					{
+						Cmd:     "ALTER TABLE `users` DROP PRIMARY KEY",
+						Reverse: "ALTER TABLE `users` ADD PRIMARY KEY USING HASH (`id`)",
+					},
+					{
+						Cmd:     "ALTER TABLE `users` ADD PRIMARY KEY (`id`)",
+						Reverse: "ALTER TABLE `users` DROP PRIMARY KEY",
+					},
+				},
+			},
+		},
 		{
 			changes: []schema.Change{
 				func() schema.Change {
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index 5451f79f..c5effdeb 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -583,6 +583,73 @@ table "users" {
 	require.EqualValues(t, exp, &s)
 }
 
+func TestMarshalSpec_PrimaryKeyType(t *testing.T) {
+	s := schema.New("test").
+		AddTables(
+			schema.NewTable("users").
+				AddColumns(
+					schema.NewStringColumn("id", "varchar(255)"),
+				),
+		)
+	s.Tables[0].SetPrimaryKey(
+		schema.NewPrimaryKey(s.Tables[0].Columns...).
+			AddAttrs(&IndexType{T: IndexTypeHash}),
+	)
+	buf, err := MarshalHCL(s)
+	require.NoError(t, err)
+	exp := `table "users" {
+  schema = schema.test
+  column "id" {
+    null = false
+    type = sql("varchar(255)")
+  }
+  primary_key {
+    columns = [column.id]
+    type    = HASH
+  }
+}
+schema "test" {
+}
+`
+	require.EqualValues(t, exp, string(buf))
+}
+
+func TestUnmarshalSpec_PrimaryKeyType(t *testing.T) {
+	var s schema.Schema
+	err := EvalHCLBytes([]byte(`table "users" {
+  schema = schema.test
+  column "id" {
+    null = false
+    type = sql("varchar(255)")
+  }
+  primary_key {
+    columns = [column.id]
+    type    = HASH
+  }
+}
+schema "test" {
+}`), &s, nil)
+	require.NoError(t, err)
+	exp := schema.New("test").
+		AddTables(
+			schema.NewTable("users").
+				AddColumns(&schema.Column{
+					Name: "id",
+					Type: &schema.ColumnType{
+						Type: &schema.StringType{
+							T:    "varchar",
+							Size: 255,
+						},
+					},
+				}),
+		)
+	exp.Tables[0].SetPrimaryKey(&schema.Index{
+		Parts: []*schema.IndexPart{{C: exp.Tables[0].Columns[0]}},
+		Attrs: []schema.Attr{&IndexType{T: IndexTypeHash}},
+	})
+	require.EqualValues(t, exp, &s)
+}
+
 func TestMarshalSpec_IndexParts(t *testing.T) {
 	c := schema.NewStringColumn("name", "text")
 	c2 := schema.NewStringColumn("Full Name", "text")