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

schema/field/field.go

@@ -26,6 +26,7 @@ type Descriptor struct {
 	UpdateDefault interface{}   // default value on update.
 	Validators    []interface{} // validator functions.
 	StorageKey    string        // sql column or gremlin property.
+	Enums         []string      // enum values.
 }
 
 // String returns a new Field with type string.
@@ -113,6 +114,22 @@ func Floats(name string) *jsonsBuilder {
 	return JSON(name, []float64{})
 }
 
+// Enum returns a new Field with type enum. An example for defining enum is as follows:
+//
+//	field.Enum("state").
+//		Values(
+//			"on",
+//			"off",
+//		).
+//		Default("on")
+//
+func Enum(name string) *enumBuilder {
+	return &enumBuilder{&Descriptor{
+		Name: name,
+		Info: &TypeInfo{Type: TypeEnum},
+	}}
+}
+
 // stringBuilder is the builder for string fields.
 type stringBuilder struct {
 	desc *Descriptor
@@ -438,7 +455,67 @@ func (b *jsonsBuilder) Comment(c string) *jsonsBuilder {
 	return b
 }
 
+// StructTag sets the struct tag of the field.
+func (b *jsonsBuilder) StructTag(s string) *jsonsBuilder {
+	b.desc.Tag = s
+	return b
+}
+
 // Descriptor implements the ent.Field interface by returning its descriptor.
 func (b *jsonsBuilder) Descriptor() *Descriptor {
 	return b.desc
 }
+
+// enumBuilder is the builder for enum fields.
+type enumBuilder struct {
+	desc *Descriptor
+}
+
+// Value sets the numeric value of the enum. Defaults to its index in the declaration.
+func (b *enumBuilder) Values(values ...string) *enumBuilder {
+	b.desc.Enums = values
+	return b
+}
+
+// StorageKey sets the storage key of the field.
+// In SQL dialects is the column name and Gremlin is the property.
+func (b *enumBuilder) StorageKey(key string) *enumBuilder {
+	b.desc.StorageKey = key
+	return b
+}
+
+// Optional indicates that this field is optional on create.
+// Unlike edges, fields are required by default.
+func (b *enumBuilder) Optional() *enumBuilder {
+	b.desc.Optional = true
+	return b
+}
+
+// Immutable indicates that this field cannot be updated.
+func (b *enumBuilder) Immutable() *enumBuilder {
+	b.desc.Immutable = true
+	return b
+}
+
+// Comment sets the comment of the field.
+func (b *enumBuilder) Comment(c string) *enumBuilder {
+	return b
+}
+
+// Nillable indicates that this field is a nillable.
+// Unlike "Optional" only fields, "Nillable" fields are pointers in the generated field.
+func (b *enumBuilder) Nillable() *enumBuilder {
+	b.desc.Nillable = true
+	return b
+}
+
+// StructTag sets the struct tag of the field.
+func (b *enumBuilder) StructTag(s string) *enumBuilder {
+	b.desc.Tag = s
+	return b
+}
+
+// Descriptor implements the ent.Field interface by returning its descriptor.
+func (b *enumBuilder) Descriptor() *Descriptor {
+	return b.desc
+}

schema/field/field_test.go

@@ -141,3 +141,15 @@ func TestField_Tag(t *testing.T) {
 		Descriptor()
 	assert.Equal(t, `json:"expired,omitempty"`, fd.Tag)
 }
+
+func TestField_Enums(t *testing.T) {
+	fd := field.Enum("role").
+		Values(
+			"user",
+			"admin",
+			"master",
+		).
+		Descriptor()
+	require.Equal(t, "role", fd.Name)
+	require.Equal(t, []string{"user", "admin", "master"}, fd.Enums)
+}

schema/field/gen/numeric.tmpl

@@ -263,4 +263,4 @@ func (b *{{ $builder }}) Descriptor() *Descriptor {
 }
 {{ end }}
 
-{{ end }}
+{{ end }}

schema/field/type.go

@@ -10,6 +10,7 @@ const (
 	TypeTime
 	TypeJSON
 	TypeBytes
+	TypeEnum
 	TypeString
 	TypeInt8
 	TypeInt16
@@ -39,16 +40,19 @@ func (t Type) Numeric() bool {
 	return t >= TypeInt8 && t < endTypes
 }
 
+// Valid reports if the given type if known type.
+func (t Type) Valid() bool {
+	return t > TypeInvalid && t < endTypes
+}
+
 // ConstName returns the constant name of a info type.
 // It's used by entc for printing the constant name in templates.
 func (t Type) ConstName() string {
-	switch t {
-	case TypeJSON:
-		return "TypeJSON"
-	case TypeTime:
-		return "TypeTime"
-	case TypeBytes:
-		return "TypeBytes"
+	switch {
+	case !t.Valid():
+		return typeNames[TypeInvalid]
+	case int(t) < len(constNames) && constNames[t] != "":
+		return constNames[t]
 	default:
 		return "Type" + strings.Title(typeNames[t])
 	}
@@ -75,7 +79,7 @@ func (t TypeInfo) String() string {
 
 // Valid reports if the given type if known type.
 func (t TypeInfo) Valid() bool {
-	return t.Type > TypeInvalid && t.Type < endTypes
+	return t.Type.Valid()
 }
 
 // Numeric reports if the given type is a numeric type.
@@ -83,23 +87,32 @@ func (t TypeInfo) Numeric() bool {
 	return t.Type.Numeric()
 }
 
-var typeNames = [...]string{
-	TypeInvalid: "invalid",
-	TypeBool:    "bool",
-	TypeTime:    "time.Time",
-	TypeJSON:    "json.RawMessage",
-	TypeBytes:   "[]byte",
-	TypeString:  "string",
-	TypeInt:     "int",
-	TypeInt8:    "int8",
-	TypeInt16:   "int16",
-	TypeInt32:   "int32",
-	TypeInt64:   "int64",
-	TypeUint:    "uint",
-	TypeUint8:   "uint8",
-	TypeUint16:  "uint16",
-	TypeUint32:  "uint32",
-	TypeUint64:  "uint64",
-	TypeFloat32: "float32",
-	TypeFloat64: "float64",
-}
+var (
+	typeNames = [...]string{
+		TypeInvalid: "invalid",
+		TypeBool:    "bool",
+		TypeTime:    "time.Time",
+		TypeJSON:    "json.RawMessage",
+		TypeBytes:   "[]byte",
+		TypeEnum:    "string",
+		TypeString:  "string",
+		TypeInt:     "int",
+		TypeInt8:    "int8",
+		TypeInt16:   "int16",
+		TypeInt32:   "int32",
+		TypeInt64:   "int64",
+		TypeUint:    "uint",
+		TypeUint8:   "uint8",
+		TypeUint16:  "uint16",
+		TypeUint32:  "uint32",
+		TypeUint64:  "uint64",
+		TypeFloat32: "float32",
+		TypeFloat64: "float64",
+	}
+	constNames = [...]string{
+		TypeJSON:  "TypeJSON",
+		TypeTime:  "TypeTime",
+		TypeEnum:  "TypeEnum",
+		TypeBytes: "TypeBytes",
+	}
+)