From ede97dea5b6ea4fb96a01dd4b3dfb672b2ceb565 Mon Sep 17 00:00:00 2001 From: Aaron L Date: Fri, 11 Nov 2016 01:01:09 -0800 Subject: [PATCH] Add enum const generation - Make postgres name its enums - Add way to filter columns by whether or not they're an enum - Split parsing of enums into name & values - Add strmangle check functions: IsEnumNormal, ShouldTitleCaseEnum - Add new strmangle enum functions to template test - Implement a set type called "once" inside the templates so that we can ensure certain things only generate one time via some unique name. --- README.md | 1 + bdb/column.go | 19 ++++++++- bdb/column_test.go | 20 +++++++++ bdb/drivers/postgres.go | 2 +- strmangle/strmangle.go | 55 +++++++++++++++++++++---- strmangle/strmangle_test.go | 66 ++++++++++++++++++++++++++---- templates.go | 32 +++++++++++++++ templates/singleton/boil_types.tpl | 49 ++++++++++++++++++++++ 8 files changed, 228 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index a00bc1f..04e5fca 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,7 @@ Table of Contents - Debug logging - Schemas support - 1d arrays, json, hstore & more +- Enum types ### Supported Databases diff --git a/bdb/column.go b/bdb/column.go index a760933..b5c408e 100644 --- a/bdb/column.go +++ b/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 +} diff --git a/bdb/column_test.go b/bdb/column_test.go index 1d144a6..0b3bb3b 100644 --- a/bdb/column_test.go +++ b/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) + } +} diff --git a/bdb/drivers/postgres.go b/bdb/drivers/postgres.go index 59c4cdc..2997117 100644 --- a/bdb/drivers/postgres.go +++ b/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 diff --git a/strmangle/strmangle.go b/strmangle/strmangle.go index 911ffdc..c8d85b0 100644 --- a/strmangle/strmangle.go +++ b/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) +} diff --git a/strmangle/strmangle_test.go b/strmangle/strmangle_test.go index b9efe1e..0d4d8cc 100644 --- a/strmangle/strmangle_test.go +++ b/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) } } } diff --git a/templates.go b/templates.go index a56c81e..9baae19 100644 --- a/templates.go +++ b/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, diff --git a/templates/singleton/boil_types.tpl b/templates/singleton/boil_types.tpl index ebb8891..9bf13e8 100644 --- a/templates/singleton/boil_types.tpl +++ b/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 -}}