ent/index: add indexes api

Reviewed By: alexsn

Differential Revision: D16757699

fbshipit-source-id: 6a0027b1f855721f1415b8c72b5b1be8bc2ce902

dialect/sql/schema/migrate.go

@@ -130,6 +130,13 @@ func (m *Migrate) create(ctx context.Context, tx dialect.Tx, tables ...*Table) e
 					return err
 				}
 			}
+			// indexes.
+			for _, idx := range t.Indexes {
+				query, args := idx.Builder(t.Name).Query()
+				if err := tx.Exec(ctx, query, args, new(sql.Result)); err != nil {
+					return fmt.Errorf("create index %q: %v", idx.Name, err)
+				}
+			}
 		}
 	}
 	// create foreign keys after tables were created/altered,
@@ -173,7 +180,7 @@ func (m *Migrate) apply(ctx context.Context, tx dialect.Tx, table string, change
 		for _, idx := range change.index.drop {
 			query, args := idx.DropBuilder(table).Query()
 			if err := tx.Exec(ctx, query, args, new(sql.Result)); err != nil {
-				return fmt.Errorf("drop index %q: %v", table, err)
+				return fmt.Errorf("drop index of table %q: %v", table, err)
 			}
 		}
 	}
@@ -215,8 +222,8 @@ type changes struct {
 	}
 	// index changes.
 	index struct {
-		add  []*Index
-		drop []*Index
+		add  Indexes
+		drop Indexes
 	}
 }
 
@@ -242,7 +249,7 @@ func (m *Migrate) changeSet(curr, new *Table) (*changes, error) {
 			change.column.add = append(change.column.add, c1)
 		// modify a non-unique column to unique.
 		case c1.Unique && !c2.Unique:
-			change.index.add = append(change.index.add, &Index{
+			change.index.add.append(&Index{
 				Name:    c1.Name,
 				Unique:  true,
 				Columns: []*Column{c1},
@@ -254,7 +261,7 @@ func (m *Migrate) changeSet(curr, new *Table) (*changes, error) {
 			if !ok {
 				return nil, fmt.Errorf("missing index to drop for column %q", c2.Name)
 			}
-			change.index.drop = append(change.index.drop, idx)
+			change.index.drop.append(idx)
 		// extending column types.
 		case m.cType(c1) != m.cType(c2):
 			if !c2.ConvertibleTo(c1) {
@@ -282,18 +289,18 @@ func (m *Migrate) changeSet(curr, new *Table) (*changes, error) {
 	for _, idx1 := range new.Indexes {
 		switch idx2, ok := curr.index(idx1.Name); {
 		case !ok:
-			change.index.add = append(change.index.add, idx1)
+			change.index.add.append(idx1)
 		// changing index cardinality require drop and create.
 		case idx1.Unique != idx2.Unique:
-			change.index.drop = append(change.index.drop, idx2)
-			change.index.add = append(change.index.add, idx1)
+			change.index.drop.append(idx2)
+			change.index.add.append(idx1)
 		}
 	}
 
 	// drop indexes.
 	for _, idx1 := range curr.Indexes {
-		if _, ok := new.index(idx1.Name); ok {
-			change.index.drop = append(change.index.drop, idx1)
+		if _, ok := new.index(idx1.Name); !ok {
+			change.index.drop.append(idx1)
 		}
 	}
 	return change, nil

dialect/sql/schema/mysql_test.go

@@ -186,6 +186,41 @@ func TestMySQL_Create(t *testing.T) {
 				mock.ExpectCommit()
 			},
 		},
+		{
+			name: "drop column to table",
+			tables: []*Table{
+				{
+					Name: "users",
+					Columns: []*Column{
+						{Name: "id", Type: field.TypeInt, Increment: true},
+					},
+					PrimaryKey: []*Column{
+						{Name: "id", Type: field.TypeInt, Increment: true},
+					},
+				},
+			},
+			options: []MigrateOption{WithDropColumn(true)},
+			before: func(mock sqlmock.Sqlmock) {
+				mock.ExpectBegin()
+				mock.ExpectQuery(escape("SHOW VARIABLES LIKE 'version'")).
+					WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).AddRow("version", "5.7.23"))
+				mock.ExpectQuery(escape("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE `TABLE_SCHEMA` = (SELECT DATABASE()) AND `TABLE_NAME` = ?")).
+					WithArgs("users").
+					WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
+				mock.ExpectQuery(escape("SELECT `column_name`, `column_type`, `is_nullable`, `column_key`, `column_default`, `extra`, `character_set_name`, `collation_name` FROM INFORMATION_SCHEMA.COLUMNS WHERE `TABLE_SCHEMA` = (SELECT DATABASE()) AND `TABLE_NAME` = ?")).
+					WithArgs("users").
+					WillReturnRows(sqlmock.NewRows([]string{"column_name", "column_type", "is_nullable", "column_key", "column_default", "extra", "character_set_name", "collation_name"}).
+						AddRow("id", "bigint(20)", "NO", "PRI", "NULL", "auto_increment", "", "").
+						AddRow("name", "varchar(255)", "NO", "YES", "NULL", "", "", ""))
+				mock.ExpectQuery(escape("SELECT `index_name`, `column_name`, `non_unique`, `seq_in_index` FROM INFORMATION_SCHEMA.STATISTICS WHERE `TABLE_SCHEMA` = (SELECT DATABASE()) AND `TABLE_NAME` = ?")).
+					WithArgs("users").
+					WillReturnRows(sqlmock.NewRows([]string{"index_name", "column_name", "non_unique", "seq_in_index"}).
+						AddRow("PRIMARY", "id", "0", "1"))
+				mock.ExpectExec(escape("ALTER TABLE `users` DROP COLUMN `name`")).
+					WillReturnResult(sqlmock.NewResult(0, 1))
+				mock.ExpectCommit()
+			},
+		},
 		{
 			name: "modify column",
 			tables: []*Table{
@@ -294,7 +329,7 @@ func TestMySQL_Create(t *testing.T) {
 			},
 		},
 		{
-			name: "remove uniqueness from column without option",
+			name: "remove uniqueness from column with option",
 			tables: []*Table{
 				{
 					Name: "users",

dialect/sql/schema/schema.go

@@ -9,17 +9,26 @@ import (
 	"fbc/ent/field"
 )
 
+// DefaultStringLen describes the default length for string/varchar types.
+const DefaultStringLen = 255
+
 // Table schema definition for SQL dialects.
 type Table struct {
 	Name        string
 	Columns     []*Column
+	columns     map[string]*Column
 	Indexes     []*Index
 	PrimaryKey  []*Column
 	ForeignKeys []*ForeignKey
 }
 
 // NewTable returns a new table with the given name.
-func NewTable(name string) *Table { return &Table{Name: name} }
+func NewTable(name string) *Table {
+	return &Table{
+		Name:    name,
+		columns: make(map[string]*Column),
+	}
+}
 
 // AddPrimary adds a new primary key to the table.
 func (t *Table) AddPrimary(c *Column) *Table {
@@ -36,10 +45,29 @@ func (t *Table) AddForeignKey(fk *ForeignKey) *Table {
 
 // AddColumn adds a new column to the table.
 func (t *Table) AddColumn(c *Column) *Table {
+	t.columns[c.Name] = c
 	t.Columns = append(t.Columns, c)
 	return t
 }
 
+// AddIndex creates and adds a new index to the table from the given options.
+func (t *Table) AddIndex(name string, unique bool, columns []string) *Table {
+	idx := &Index{
+		Name:    name,
+		Unique:  unique,
+		columns: columns,
+		Columns: make([]*Column, len(columns)),
+	}
+	for i, name := range columns {
+		c, ok := t.columns[name]
+		if ok {
+			idx.Columns[i] = c
+		}
+	}
+	t.Indexes = append(t.Indexes, idx)
+	return t
+}
+
 // MySQL returns the MySQL DSL query for table creation.
 func (t *Table) MySQL(version string) *sql.TableBuilder {
 	b := sql.CreateTable(t.Name).IfNotExists()
@@ -99,6 +127,10 @@ func (t *Table) index(name string) (*Index, bool) {
 			return idx, true
 		}
 	}
+	// if it is an "implicit index" (unique constraint on table creation).
+	if c, ok := t.column(name); ok && c.Unique {
+		return &Index{Name: name, Unique: c.Unique, Columns: []*Column{c}, columns: []string{c.Name}}, true
+	}
 	return nil, false
 }
 
@@ -218,7 +250,7 @@ func (c *Column) SQLiteType() (t string) {
 	case field.TypeString:
 		size := c.Size
 		if size == 0 {
-			size = 255
+			size = DefaultStringLen
 		}
 		// sqlite has no size limit on varchar.
 		t = fmt.Sprintf("varchar(%d)", size)
@@ -343,14 +375,13 @@ func (c *Column) nullable(b *sql.ColumnBuilder) {
 // defaultSize returns the default size for MySQL varchar
 // type based on column size, charset and table indexes.
 func (c *Column) defaultSize(version string) int {
-	size := 255
 	parts := strings.Split(version, ".")
 	// non-unique or invalid version.
 	if !c.Unique || len(parts) == 1 || parts[0] == "" || parts[1] == "" {
-		return size
+		return DefaultStringLen
 	}
 	if major, minor := parts[0], parts[1]; major > "5" || minor > "6" {
-		return size
+		return DefaultStringLen
 	}
 	return 191
 }
@@ -441,6 +472,16 @@ func (i *Index) DropBuilder(table string) *sql.DropIndexBuilder {
 // multiple sql rows can represent the same index (multi-columns indexes).
 type Indexes []*Index
 
+// append wraps the basic `append` function by filtering duplicates indexes.
+func (i *Indexes) append(idx1 *Index) {
+	for _, idx2 := range *i {
+		if idx2.Name == idx1.Name {
+			return
+		}
+	}
+	*i = append(*i, idx1)
+}
+
 // ScanMySQL scans sql.Rows into an Indexes list. The query for returning the rows,
 // should return the following 4 columns: INDEX_NAME, COLUMN_NAME, NON_UNIQUE, SEQ_IN_INDEX.
 // SEQ_IN_INDEX specifies the position of the column in the index columns.
@@ -459,6 +500,10 @@ func (i *Indexes) ScanMySQL(rows *sql.Rows) error {
 		idx, ok := names[name]
 		if !ok {
 			idx = &Index{Name: name, Unique: !nonuniq}
+			// ignore primary keys.
+			if idx.Primary() {
+				continue
+			}
 			*i = append(*i, idx)
 			names[name] = idx
 		}