Merge branch 'dev' of github.com:vattle/sqlboiler into dev

README.md

@@ -97,6 +97,7 @@ Table of Contents
 - Debug logging
 - Schemas support
 - 1d arrays, json, hstore & more
+- Enum types
 
 ### Supported Databases
 

bdb/column.go

@@ -1,6 +1,10 @@
 package bdb
 
-import "github.com/vattle/sqlboiler/strmangle"
+import (
+	"strings"
+
+	"github.com/vattle/sqlboiler/strmangle"
+)
 
 // Column holds information about a database column.
 // Types are Go types, converted by TranslateColumnType.
@@ -54,3 +58,16 @@ func FilterColumnsByDefault(defaults bool, columns []Column) []Column {
 
 	return cols
 }
+
+// FilterColumnsByEnum generates the list of columns that are enum values.
+func FilterColumnsByEnum(columns []Column) []Column {
+	var cols []Column
+
+	for _, c := range columns {
+		if strings.HasPrefix(c.DBType, "enum") {
+			cols = append(cols, c)
+		}
+	}
+
+	return cols
+}

bdb/column_test.go

@@ -66,3 +66,23 @@ func TestFilterColumnsByDefault(t *testing.T) {
 		t.Errorf("Invalid result: %#v", res)
 	}
 }
+
+func TestFilterColumnsByEnum(t *testing.T) {
+	t.Parallel()
+
+	cols := []Column{
+		{Name: "col1", DBType: "enum('hello')"},
+		{Name: "col2", DBType: "enum('hello','there')"},
+		{Name: "col3", DBType: "enum"},
+		{Name: "col4", DBType: ""},
+		{Name: "col5", DBType: "int"},
+	}
+
+	res := FilterColumnsByEnum(cols)
+	if res[0].Name != `col1` {
+		t.Errorf("Invalid result: %#v", res)
+	}
+	if res[1].Name != `col2` {
+		t.Errorf("Invalid result: %#v", res)
+	}
+}

bdb/drivers/postgres.go

@@ -130,7 +130,7 @@ func (p *PostgresDriver) Columns(schema, tableName string) ([]bdb.Column, error)
 			case when c.data_type = 'USER-DEFINED' and c.udt_name <> 'hstore'
 			then
 			(
-				select 'enum(''' || string_agg(labels.label, ''',''') || ''')'
+				select 'enum.' || c.udt_name || '(''' || string_agg(labels.label, ''',''') || ''')'
 				from (
 					select pg_enum.enumlabel as label
 					from pg_enum

strmangle/strmangle.go

@@ -17,7 +17,9 @@ var (
 	idAlphabet    = []byte("abcdefghijklmnopqrstuvwxyz")
 	smartQuoteRgx = regexp.MustCompile(`^(?i)"?[a-z_][_a-z0-9]*"?(\."?[_a-z][_a-z0-9]*"?)*(\.\*)?$`)
 
-	rgxEnum = regexp.MustCompile(`^enum\((,?'[^']+')+\)$`)
+	rgxEnum            = regexp.MustCompile(`^enum(\.[a-z_]+)?\((,?'[^']+')+\)$`)
+	rgxEnumIsOK        = regexp.MustCompile(`^(?i)[a-z][a-z0-9_]*$`)
+	rgxEnumShouldTitle = regexp.MustCompile(`^[a-z][a-z0-9_]*$`)
 )
 
 var uppercaseWords = map[string]struct{}{
@@ -577,15 +579,54 @@ func GenerateIgnoreTags(tags []string) string {
 	return buf.String()
 }
 
-// ParseEnum takes a string that looks like:
-// enum('one','two') and returns the strings one, two
-func ParseEnum(s string) []string {
+// ParseEnumVals returns the values from an enum string
+//
+// Postgres and MySQL drivers return different values
+// psql:  enum.enum_name('values'...)
+// mysql: enum('values'...)
+func ParseEnumVals(s string) []string {
 	if !rgxEnum.MatchString(s) {
 		return nil
 	}
 
-	s = strings.TrimPrefix(s, "enum('")
-	s = strings.TrimSuffix(s, "')")
-
+	startIndex := strings.IndexByte(s, '(')
+	s = s[startIndex+2 : len(s)-2]
 	return strings.Split(s, "','")
 }
+
+// ParseEnumName returns the name portion of an enum if it exists
+//
+// Postgres and MySQL drivers return different values
+// psql:  enum.enum_name('values'...)
+// mysql: enum('values'...)
+// In the case of mysql, the name will never return anything
+func ParseEnumName(s string) string {
+	if !rgxEnum.MatchString(s) {
+		return ""
+	}
+
+	endIndex := strings.IndexByte(s, '(')
+	s = s[:endIndex]
+	startIndex := strings.IndexByte(s, '.')
+	if startIndex < 0 {
+		return ""
+	}
+
+	return s[startIndex+1:]
+}
+
+// IsEnumNormal checks a set of eval values to see if they're "normal"
+func IsEnumNormal(values []string) bool {
+	for _, v := range values {
+		if !rgxEnumIsOK.MatchString(v) {
+			return false
+		}
+	}
+
+	return true
+}
+
+// ShouldTitleCaseEnum checks a value to see if it's title-case-able
+func ShouldTitleCaseEnum(value string) bool {
+	return rgxEnumShouldTitle.MatchString(value)
+}

strmangle/strmangle_test.go

@@ -1,7 +1,6 @@
 package strmangle
 
 import (
-	"fmt"
 	"strings"
 	"testing"
 )
@@ -518,13 +517,66 @@ func TestGenerateIgnoreTags(t *testing.T) {
 func TestParseEnum(t *testing.T) {
 	t.Parallel()
 
-	vals := []string{"one", "two", "three"}
-	toParse := fmt.Sprintf("enum('%s')", strings.Join(vals, "','"))
+	tests := []struct {
+		Enum string
+		Name string
+		Vals []string
+	}{
+		{"enum('one')", "", []string{"one"}},
+		{"enum('one','two')", "", []string{"one", "two"}},
+		{"enum.working('one')", "working", []string{"one"}},
+		{"enum.wor_king('one','two')", "wor_king", []string{"one", "two"}},
+	}
 
-	gotVals := ParseEnum(toParse)
-	for i, v := range vals {
-		if gotVals[i] != v {
-			t.Errorf("%d) want: %s, got %s", i, v, gotVals[i])
+	for i, test := range tests {
+		name := ParseEnumName(test.Enum)
+		vals := ParseEnumVals(test.Enum)
+		if name != test.Name {
+			t.Errorf("%d) name was wrong, want: %s got: %s (%s)", i, test.Name, name, test.Enum)
+		}
+		for j, v := range test.Vals {
+			if v != vals[j] {
+				t.Errorf("%d.%d) value was wrong, want: %s got: %s (%s)", i, j, v, vals[j], test.Enum)
+			}
+		}
+	}
+}
+
+func TestIsEnumNormal(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		Vals []string
+		Ok   bool
+	}{
+		{[]string{"o1ne", "two2"}, true},
+		{[]string{"one", "t#wo2"}, false},
+		{[]string{"1one", "two2"}, false},
+	}
+
+	for i, test := range tests {
+		if got := IsEnumNormal(test.Vals); got != test.Ok {
+			t.Errorf("%d) want: %t got: %t, %#v", i, test.Ok, got, test.Vals)
+		}
+	}
+}
+
+func TestShouldTitleCaseEnum(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		Val string
+		Ok  bool
+	}{
+		{"hello_there0", true},
+		{"hEllo", false},
+		{"_hello", false},
+		{"0hello", false},
+	}
+
+	for i, test := range tests {
+		if got := ShouldTitleCaseEnum(test.Val); got != test.Ok {
+			t.Errorf("%d) want: %t got: %t, %v", i, test.Ok, got, test.Val)
 		}
 	}
 }

templates.go

@@ -121,6 +121,28 @@ func loadTemplate(dir string, filename string) (*template.Template, error) {
 	return tpl.Lookup(filename), err
 }
 
+// set is to stop duplication from named enums, allowing a template loop
+// to keep some state
+type once map[string]struct{}
+
+func newOnce() once {
+	return make(once)
+}
+
+func (o once) Has(s string) bool {
+	_, ok := o[s]
+	return ok
+}
+
+func (o once) Put(s string) bool {
+	if _, ok := o[s]; ok {
+		return false
+	}
+
+	o[s] = struct{}{}
+	return true
+}
+
 // templateStringMappers are placed into the data to make it easy to use the
 // stringMap function.
 var templateStringMappers = map[string]func(string) string{
@@ -157,6 +179,15 @@ var templateFunctions = template.FuncMap{
 	"generateTags":       strmangle.GenerateTags,
 	"generateIgnoreTags": strmangle.GenerateIgnoreTags,
 
+	// Enum ops
+	"parseEnumName":       strmangle.ParseEnumName,
+	"parseEnumVals":       strmangle.ParseEnumVals,
+	"isEnumNormal":        strmangle.IsEnumNormal,
+	"shouldTitleCaseEnum": strmangle.ShouldTitleCaseEnum,
+	"onceNew":             newOnce,
+	"oncePut":             once.Put,
+	"onceHas":             once.Has,
+
 	// String Map ops
 	"makeStringMap": strmangle.MakeStringMap,
 
@@ -173,6 +204,7 @@ var templateFunctions = template.FuncMap{
 
 	// dbdrivers ops
 	"filterColumnsByDefault": bdb.FilterColumnsByDefault,
+	"filterColumnsByEnum":    bdb.FilterColumnsByEnum,
 	"sqlColDefinitions":      bdb.SQLColDefinitions,
 	"columnNames":            bdb.ColumnNames,
 	"columnDBTypes":          bdb.ColumnDBTypes,

templates/singleton/boil_types.tpl

@@ -35,3 +35,52 @@ func makeCacheKey(wl, nzDefaults []string) string {
 	strmangle.PutBuffer(buf)
 	return str
 }
+
+{{/*
+The following is a little bit of black magic and deserves some explanation
+
+Because postgres and mysql define enums completely differently (one at the
+database level as a custom datatype, and one at the table column level as
+a unique thing per table)... There's a chance the enum is named (postgres)
+and not (mysql). So we can't do this per table so this code is here.
+
+We loop through each table and column looking for enums. If it's named, we
+then use some disgusting magic to write state during the template compile to
+the "once" map. This lets named enums only be defined once if they're referenced
+multiple times in many (or even the same) tables.
+
+Then we check if all it's values are normal, if they are we create the enum
+output, if not we output a friendly error message as a comment to aid in
+debugging.
+
+Postgres output looks like: EnumNameEnumValue = "enumvalue"
+MySQL output looks like:    TableNameColNameEnumValue = "enumvalue"
+
+It only titlecases the EnumValue portion if it's snake-cased.
+*/}}
+{{$dot := . -}}
+{{$once := onceNew}}
+{{- range $table := .Tables -}}
+	{{- range $col := $table.Columns | filterColumnsByEnum -}}
+		{{- $name := parseEnumName $col.DBType -}}
+		{{- $vals := parseEnumVals $col.DBType -}}
+		{{- $isNamed := ne (len $name) 0}}
+		{{- if and $isNamed (onceHas $once $name) -}}
+		{{- else -}}
+			{{- if $isNamed -}}
+				{{$_ := oncePut $once $name}}
+			{{- end -}}
+{{- if and (gt (len $vals) 0) (isEnumNormal $vals)}}
+// Enum values for {{if $isNamed}}{{$name}}{{else}}{{$table.Name}}.{{$col.Name}}{{end}}
+const (
+	{{- range $val := $vals -}}
+	{{- if $isNamed}}{{titleCase $name}}{{else}}{{titleCase $table.Name}}{{titleCase $col.Name}}{{end -}}
+	{{if shouldTitleCaseEnum $val}}{{titleCase $val}}{{else}}{{$val}}{{end}} = "{{$val}}"
+	{{end -}}
+)
+{{- else}}
+// Enum values for {{if $isNamed}}{{$name}}{{else}}{{$table.Name}}.{{$col.Name}}{{end}} are not proper Go identifiers, cannot emit constants
+{{- end -}}
+		{{- end -}}
+	{{- end -}}
+{{- end -}}