ent: add support for enum types

Summary: Pull Request resolved: https://github.com/facebookincubator/ent/pull/45 Reviewed By: alexsn Differential Revision: D17715085 fbshipit-source-id: 7472e1bd9cf7a8a5bd98f96e6e884c0e27f36803
2026-05-22 09:31:45 +03:00 · 2019-10-02 12:54:06 -07:00
parent 55fefbfbec
commit 1ebfa489c5
61 changed files with 4269 additions and 3312 deletions
--- a/dialect/sql/schema/mysql_test.go
+++ b/dialect/sql/schema/mysql_test.go
@@ -55,6 +55,7 @@ func TestMySQL_Create(t *testing.T) {
 						{Name: "name", Type: field.TypeString, Nullable: true},
 						{Name: "age", Type: field.TypeInt},
 						{Name: "doc", Type: field.TypeJSON, Nullable: true},
+						{Name: "enums", Type: field.TypeEnum, Enums: []string{"a", "b"}},
 					},
 				},
 			},
@@ -65,7 +66,7 @@ func TestMySQL_Create(t *testing.T) {
 				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(0))
-				mock.ExpectExec(escape("CREATE TABLE IF NOT EXISTS `users`(`id` bigint AUTO_INCREMENT NOT NULL, `name` varchar(255) NULL, `age` bigint NOT NULL, `doc` json NULL, PRIMARY KEY(`id`)) CHARACTER SET utf8mb4")).
+				mock.ExpectExec(escape("CREATE TABLE IF NOT EXISTS `users`(`id` bigint AUTO_INCREMENT NOT NULL, `name` varchar(255) NULL, `age` bigint NOT NULL, `doc` json NULL, `enums` enum('a', 'b') NOT NULL, PRIMARY KEY(`id`)) CHARACTER SET utf8mb4")).
 					WillReturnResult(sqlmock.NewResult(0, 1))
 				mock.ExpectCommit()
 			},
@@ -183,8 +184,8 @@ func TestMySQL_Create(t *testing.T) {
 					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", "", "", "").
-						AddRow("text", "longtext", "NO", "YES", "NULL", "", "", ""))
+						AddRow("name", "varchar(255)", "YES", "YES", "NULL", "", "", "").
+						AddRow("text", "longtext", "YES", "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"}).
@@ -194,6 +195,47 @@ func TestMySQL_Create(t *testing.T) {
 				mock.ExpectCommit()
 			},
 		},
+		{
+			name: "enums",
+			tables: []*Table{
+				{
+					Name: "users",
+					Columns: []*Column{
+						{Name: "id", Type: field.TypeInt, Increment: true},
+						{Name: "name", Type: field.TypeString, Nullable: true},
+						{Name: "enums1", Type: field.TypeEnum, Enums: []string{"a", "b"}}, // add enum.
+						{Name: "enums2", Type: field.TypeEnum, Enums: []string{"a"}},      // remove enum.
+						{Name: "enums3", Type: field.TypeEnum, Enums: []string{"a", "b"}}, // order does not effect.
+					},
+					PrimaryKey: []*Column{
+						{Name: "id", Type: field.TypeInt, Increment: 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)", "YES", "YES", "NULL", "", "", "").
+						AddRow("enums1", "enum('a')", "YES", "NO", "NULL", "", "", "").
+						AddRow("enums2", "enum('b', 'a')", "NO", "YES", "NULL", "", "", "").
+						AddRow("enums3", "enum('b', 'a')", "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` MODIFY COLUMN `enums1` enum('a', 'b') NOT NULL, MODIFY COLUMN `enums2` enum('a') NOT NULL")).
+					WillReturnResult(sqlmock.NewResult(0, 1))
+				mock.ExpectCommit()
+			},
+		},
 		{
 			name: "datetime and timestamp",
 			tables: []*Table{
--- a/dialect/sql/schema/schema.go
+++ b/dialect/sql/schema/schema.go
@@ -7,6 +7,7 @@ package schema
 import (
 	"fmt"
 	"math"
+	"sort"
 	"strconv"
 	"strings"

@@ -182,6 +183,7 @@ type Column struct {
 	Nullable  bool        // null or not null attribute.
 	Default   interface{} // default value.
 	indexes   Indexes     // linked indexes.
+	Enums     []string    // enum values.
 }

 // UniqueKey returns boolean indicates if this column is a unique key.
@@ -275,6 +277,13 @@ func (c *Column) MySQLType(version string) (t string) {
 		// in MySQL timestamp columns are `NOT NULL by default, and assigning NULL
 		// assigns the current_timestamp(). We avoid this if not set otherwise.
 		c.Nullable = true
+	case field.TypeEnum:
+		values := make([]string, len(c.Enums))
+		for i, e := range c.Enums {
+			values[i] = fmt.Sprintf("'%s'", e)
+		}
+		sort.Strings(values)
+		t = fmt.Sprintf("enum(%s)", strings.Join(values, ", "))
 	default:
 		panic(fmt.Sprintf("unsupported type %q for column %q", c.Type.String(), c.Name))
 	}
@@ -292,7 +301,7 @@ func (c *Column) SQLiteType() (t string) {
 		t = "bigint"
 	case field.TypeBytes:
 		t = "blob"
-	case field.TypeString:
+	case field.TypeString, field.TypeEnum:
 		size := c.Size
 		if size == 0 {
 			size = DefaultStringLen
@@ -325,7 +334,7 @@ func (c *Column) ScanMySQL(rows *sql.Rows) error {
 		c.Nullable = nullable.String == "YES"
 	}
 	switch parts := strings.FieldsFunc(c.typ, func(r rune) bool {
-		return r == '(' || r == ')' || r == ' '
+		return r == '(' || r == ')' || r == ' ' || r == ','
 	}); parts[0] {
 	case "int":
 		c.Type = field.TypeInt32
@@ -380,6 +389,12 @@ func (c *Column) ScanMySQL(rows *sql.Rows) error {
 		c.Type = field.TypeString
 	case "json":
 		c.Type = field.TypeJSON
+	case "enum":
+		c.Type = field.TypeEnum
+		c.Enums = make([]string, len(parts)-1)
+		for i, e := range parts[1:] {
+			c.Enums[i] = strings.Trim(e, "'")
+		}
 	}
 	if defaults.Valid && defaults.String != Null {
 		return c.ScanDefault(defaults.String)