sql/mysql: add pk part attributes to spec (#3180)

internal/integration/testdata/mysql/primary-key-parts.txt

@@ -0,0 +1,103 @@
+only mysql8
+
+# Version 1.
+atlas schema inspect --url file://schema.v1.sql --dev-url URL > got
+cmp got schema.v1.hcl
+atlas schema inspect --url file://schema.v1.hcl --dev-url URL > got
+cmp got schema.v1.hcl
+atlas migrate diff v1 --to file://schema.v1.hcl --dev-url URL
+cmpmig 0 migration.v1.sql
+atlas migrate diff v1-check --to file://schema.v1.hcl --dev-url URL
+stdout 'The migration directory is synced with the desired state, no changes to be made'
+
+# Version 2.
+atlas migrate diff v2 --to file://schema.v2.hcl --dev-url URL
+cmpmig 1 migration.v2.sql
+atlas migrate diff v2-check --to file://schema.v2.hcl --dev-url URL
+stdout 'The migration directory is synced with the desired state, no changes to be made'
+
+-- schema.v1.sql --
+CREATE TABLE `t1` (`id` tinytext NOT NULL, PRIMARY KEY (`id` (7)));
+CREATE TABLE `t2` (`id` tinytext NOT NULL, PRIMARY KEY (`id` (7) DESC));
+-- schema.v1.hcl --
+table "t1" {
+  schema = schema.script_primary_key_parts
+  column "id" {
+    null = false
+    type = tinytext
+  }
+  primary_key {
+    on {
+      column = column.id
+      prefix = 7
+    }
+  }
+}
+table "t2" {
+  schema = schema.script_primary_key_parts
+  column "id" {
+    null = false
+    type = tinytext
+  }
+  primary_key {
+    on {
+      desc   = true
+      column = column.id
+      prefix = 7
+    }
+  }
+}
+schema "script_primary_key_parts" {
+  charset = "utf8mb4"
+  collate = "utf8mb4_0900_ai_ci"
+}
+-- migration.v1.sql --
+-- Create "t1" table
+CREATE TABLE `t1` (`id` tinytext NOT NULL, PRIMARY KEY (`id` (7))) CHARSET utf8mb4 COLLATE utf8mb4_0900_ai_ci;
+-- Create "t2" table
+CREATE TABLE `t2` (`id` tinytext NOT NULL, PRIMARY KEY (`id` (7) DESC)) CHARSET utf8mb4 COLLATE utf8mb4_0900_ai_ci;
+-- schema.v2.hcl --
+table "t1" {
+  schema = schema.script_primary_key_parts
+  column "id" {
+    null = false
+    type = tinytext
+  }
+  column "id2" {
+    null = false
+    type = tinytext
+  }
+  primary_key {
+    on {
+      column = column.id
+      prefix = 7
+    }
+    on {
+      column = column.id2
+      prefix = 1
+    }
+  }
+}
+table "t2" {
+  schema = schema.script_primary_key_parts
+  column "id" {
+    null = false
+    type = tinytext
+  }
+  primary_key {
+    on {
+      desc   = false
+      column = column.id
+      prefix = 6
+    }
+  }
+}
+schema "script_primary_key_parts" {
+  charset = "utf8mb4"
+  collate = "utf8mb4_0900_ai_ci"
+}
+-- migration.v2.sql --
+-- Modify "t1" table
+ALTER TABLE `t1` ADD COLUMN `id2` tinytext NOT NULL, DROP PRIMARY KEY, ADD PRIMARY KEY (`id` (7), `id2` (1));
+-- Modify "t2" table
+ALTER TABLE `t2` DROP PRIMARY KEY, ADD PRIMARY KEY (`id` (6));

sql/mysql/inspect.go

@@ -330,17 +330,6 @@ func (i *inspect) addColumn(s *schema.Schema, rows *sql.Rows) error {
 		c.SetCollation(collation.String)
 	}
 	t.AddColumns(c)
-	// From MySQL doc: A UNIQUE index may be displayed as "PRI" if it is NOT NULL
-	// and there is no PRIMARY KEY in the table. We detect this in `addIndexes`.
-	if key.String == "PRI" {
-		if t.PrimaryKey == nil {
-			t.PrimaryKey = &schema.Index{Table: t, Name: key.String}
-		}
-		t.PrimaryKey.Parts = append(t.PrimaryKey.Parts, &schema.IndexPart{
-			C:     c,
-			SeqNo: len(t.PrimaryKey.Parts),
-		})
-	}
 	return nil
 }
 
@@ -360,7 +349,6 @@ func (i *inspect) indexes(ctx context.Context, s *schema.Schema) error {
 
 // addIndexes scans the rows and adds the indexes to the table.
 func (i *inspect) addIndexes(s *schema.Schema, rows *sql.Rows) error {
-	hasPK := make(map[*schema.Table]bool)
 	for rows.Next() {
 		var (
 			seqno                          int
@@ -375,23 +363,28 @@ func (i *inspect) addIndexes(s *schema.Schema, rows *sql.Rows) error {
 		if !ok {
 			return fmt.Errorf("table %q was not found in schema", table)
 		}
-		// Ignore primary keys.
-		if name == "PRIMARY" {
-			hasPK[t] = true
-			continue
-		}
-		idx, ok := t.Index(name)
-		if !ok {
-			idx = schema.NewIndex(name).
-				SetUnique(!nonuniq.Bool).
-				AddAttrs(&IndexType{T: indexType})
-			if indexType == IndexTypeFullText {
-				putShow(t).addFullText(idx)
+		var idx *schema.Index
+		switch {
+		// Primary key.
+		case name == "PRIMARY":
+			if idx = t.PrimaryKey; idx == nil {
+				idx = schema.NewIndex("PRI").
+					AddAttrs(&IndexType{T: indexType})
+				t.PrimaryKey = idx
 			}
-			if sqlx.ValidString(comment) {
-				idx.SetComment(comment.String)
+		default:
+			if idx, ok = t.Index(name); !ok {
+				idx = schema.NewIndex(name).
+					SetUnique(!nonuniq.Bool).
+					AddAttrs(&IndexType{T: indexType})
+				if indexType == IndexTypeFullText {
+					putShow(t).addFullText(idx)
+				}
+				if sqlx.ValidString(comment) {
+					idx.SetComment(comment.String)
+				}
+				t.AddIndexes(idx)
 			}
-			t.AddIndexes(idx)
 		}
 		// Rows are ordered by SEQ_IN_INDEX that specifies the
 		// position of the column in the index definition.
@@ -419,11 +412,6 @@ func (i *inspect) addIndexes(s *schema.Schema, rows *sql.Rows) error {
 		}
 		idx.Parts = append(idx.Parts, part)
 	}
-	for _, t := range s.Tables {
-		if !hasPK[t] && t.PrimaryKey != nil {
-			t.PrimaryKey = nil
-		}
-	}
 	return nil
 }
 

sql/mysql/inspect_test.go

@@ -92,7 +92,7 @@ func TestDriver_InspectTable(t *testing.T) {
 				require.Len(t.PrimaryKey.Parts, 1)
 				require.True(t.PrimaryKey.Parts[0].C == t.Columns[0])
 				require.EqualValues([]*schema.Column{
-					{Name: "id", Type: &schema.ColumnType{Raw: "bigint(20)", Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoIncrement{V: 55834574848}}},
+					{Name: "id", Type: &schema.ColumnType{Raw: "bigint(20)", Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoIncrement{V: 55834574848}}, Indexes: []*schema.Index{t.PrimaryKey}},
 				}, t.Columns)
 			},
 		},
@@ -581,7 +581,7 @@ func TestDriver_InspectTable(t *testing.T) {
 					{Name: "unique_index", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "BTREE"}}},
 				}
 				columns := []*schema.Column{
-					{Name: "id", Type: &schema.ColumnType{Raw: "int", Type: &schema.IntegerType{T: "int"}}},
+					{Name: "id", Type: &schema.ColumnType{Raw: "int", Type: &schema.IntegerType{T: "int"}}, Indexes: []*schema.Index{t.PrimaryKey}},
 					{Name: "nickname", Type: &schema.ColumnType{Raw: "varchar(255)", Type: &schema.StringType{T: "varchar", Size: 255}}, Indexes: indexes[0:1], Attrs: []schema.Attr{&schema.Charset{V: "utf8mb4"}, &schema.Collation{V: "utf8mb4_0900_ai_ci"}}},
 					{Name: "oid", Type: &schema.ColumnType{Raw: "int", Type: &schema.IntegerType{T: "int"}}, Indexes: indexes[2:]},
 					{Name: "uid", Type: &schema.ColumnType{Raw: "int", Type: &schema.IntegerType{T: "int"}}, Indexes: indexes[2:]},

sql/mysql/sqlspec.go

@@ -157,17 +157,11 @@ func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, er
 
 // 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(spec, idx); err != nil {
-		return nil, err
-	}
-	if err := convertIndexParser(spec, idx); err != nil {
-		return nil, err
-	}
-	return idx, nil
+	return convertIndex(&sqlspec.Index{
+		Parts:            spec.Parts,
+		Columns:          spec.Columns,
+		DefaultExtension: spec.DefaultExtension,
+	}, parent)
 }
 
 // convertIndex converts a sqlspec.Index into a schema.Index.
@@ -323,12 +317,16 @@ func tableSpec(t *schema.Table) (*sqlspec.Table, error) {
 }
 
 func pkSpec(idx *schema.Index) (*sqlspec.PrimaryKey, error) {
-	spec, err := specutil.FromPrimaryKey(idx)
+	spec, err := indexSpec(idx)
 	if err != nil {
 		return nil, err
 	}
-	spec.Extra.Attrs = indexTypeSpec(idx, spec.Extra.Attrs)
-	return spec, nil
+	for _, p := range spec.Parts {
+		if p.Expr != "" {
+			return nil, fmt.Errorf("primary key %q cannot have functional part", idx.Name)
+		}
+	}
+	return &sqlspec.PrimaryKey{Parts: spec.Parts, Columns: spec.Columns, DefaultExtension: spec.DefaultExtension}, nil
 }
 
 func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {

sql/sqlspec/sqlspec.go

@@ -58,6 +58,7 @@ type (
 
 	// PrimaryKey holds a specification for the primary key of a table.
 	PrimaryKey struct {
+		Parts   []*IndexPart     `spec:"on"`
 		Columns []*schemahcl.Ref `spec:"columns"`
 		schemahcl.DefaultExtension
 	}