From 63c1fd643031651d22ced1a8410cb9669c45f69f Mon Sep 17 00:00:00 2001
From: go-jet <zer0.sub@protonmail.com>
Date: Sat, 15 Feb 2020 11:20:51 +0100
Subject: [PATCH] [bug] Fix crash on generating enum sql builder files when
 enum contains numeric values.

---
 generator/internal/template/generate.go  |  3 ++-
 generator/internal/template/templates.go |  4 ++--
 internal/utils/utils.go                  | 10 ++++++++++
 internal/utils/utils_test.go             |  5 +++++
 tests/postgres/generator_test.go         | 25 ++++++++++++++++++++++--
 5 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/generator/internal/template/generate.go b/generator/internal/template/generate.go
index a076bb1..3872f76 100644
--- a/generator/internal/template/generate.go
+++ b/generator/internal/template/generate.go
@@ -73,7 +73,8 @@ func generateGoFiles(dirPath, packageName string, template string, metaDataList
 func GenerateTemplate(templateText string, templateData interface{}, dialect jet.Dialect, params ...map[string]interface{}) ([]byte, error) {
 
 	t, err := template.New("sqlBuilderTableTemplate").Funcs(template.FuncMap{
-		"ToGoIdentifier": utils.ToGoIdentifier,
+		"ToGoIdentifier":          utils.ToGoIdentifier,
+		"ToGoEnumValueIdentifier": utils.ToGoEnumValueIdentifier,
 		"now": func() string {
 			return time.Now().Format(time.RFC850)
 		},
diff --git a/generator/internal/template/templates.go b/generator/internal/template/templates.go
index e5328a3..40e4773 100644
--- a/generator/internal/template/templates.go
+++ b/generator/internal/template/templates.go
@@ -94,11 +94,11 @@ import "github.com/go-jet/jet/{{dialect.PackageName}}"
 
 var {{ToGoIdentifier $.Name}} = &struct {
 {{- range $index, $element := .Values}}
-	{{ToGoIdentifier $element}} {{dialect.PackageName}}.StringExpression
+	{{ToGoEnumValueIdentifier $.Name $element}} {{dialect.PackageName}}.StringExpression
 {{- end}}
 } {
 {{- range $index, $element := .Values}}
-	{{ToGoIdentifier $element}}: {{dialect.PackageName}}.NewEnumValue("{{$element}}"),
+	{{ToGoEnumValueIdentifier $.Name $element}}: {{dialect.PackageName}}.NewEnumValue("{{$element}}"),
 {{- end}}
 }
 `
diff --git a/internal/utils/utils.go b/internal/utils/utils.go
index 42a5c36..5301605 100644
--- a/internal/utils/utils.go
+++ b/internal/utils/utils.go
@@ -10,6 +10,7 @@ import (
 	"reflect"
 	"strings"
 	"time"
+	"unicode"
 )
 
 // ToGoIdentifier converts database to Go identifier.
@@ -17,6 +18,15 @@ func ToGoIdentifier(databaseIdentifier string) string {
 	return snaker.SnakeToCamel(replaceInvalidChars(databaseIdentifier))
 }
 
+func ToGoEnumValueIdentifier(enumName, enumValue string) string {
+	enumValueIdentifier := ToGoIdentifier(enumValue)
+	if !unicode.IsLetter([]rune(enumValueIdentifier)[0]) {
+		return ToGoIdentifier(enumName) + enumValueIdentifier
+	}
+
+	return enumValueIdentifier
+}
+
 // ToGoFileName converts database identifier to Go file name.
 func ToGoFileName(databaseIdentifier string) string {
 	return strings.ToLower(replaceInvalidChars(databaseIdentifier))
diff --git a/internal/utils/utils_test.go b/internal/utils/utils_test.go
index ac57bbd..787b14b 100644
--- a/internal/utils/utils_test.go
+++ b/internal/utils/utils_test.go
@@ -25,6 +25,11 @@ func TestToGoIdentifier(t *testing.T) {
 	assert.Equal(t, ToGoIdentifier("My-Table"), "MyTable")
 }
 
+func TestToGoEnumValueIdentifier(t *testing.T) {
+	assert.Equal(t, ToGoEnumValueIdentifier("enum_name", "enum_value"), "EnumValue")
+	assert.Equal(t, ToGoEnumValueIdentifier("NumEnum", "100"), "NumEnum100")
+}
+
 func TestErrorCatchErr(t *testing.T) {
 	var err error
 
diff --git a/tests/postgres/generator_test.go b/tests/postgres/generator_test.go
index b98633a..52ddfac 100644
--- a/tests/postgres/generator_test.go
+++ b/tests/postgres/generator_test.go
@@ -283,14 +283,15 @@ func TestGeneratedAllTypesSQLBuilderFiles(t *testing.T) {
 	enumFiles, err := ioutil.ReadDir(enumDir)
 	assert.NoError(t, err)
 
-	testutils.AssertFileNamesEqual(t, enumFiles, "mood.go")
+	testutils.AssertFileNamesEqual(t, enumFiles, "mood.go", "level.go")
 	testutils.AssertFileContent(t, enumDir+"mood.go", "\npackage enum", moodEnumContent)
+	testutils.AssertFileContent(t, enumDir+"level.go", "\npackage enum", levelEnumContent)
 
 	modelFiles, err := ioutil.ReadDir(modelDir)
 	assert.NoError(t, err)
 
 	testutils.AssertFileNamesEqual(t, modelFiles, "all_types.go", "all_types_view.go", "employee.go", "link.go",
-		"mood.go", "person.go", "person_phone.go", "weird_names_table.go")
+		"mood.go", "person.go", "person_phone.go", "weird_names_table.go", "level.go")
 
 	testutils.AssertFileContent(t, modelDir+"all_types.go", "\npackage model", allTypesModelContent)
 
@@ -319,6 +320,26 @@ var Mood = &struct {
 }
 `
 
+var levelEnumContent = `
+package enum
+
+import "github.com/go-jet/jet/postgres"
+
+var Level = &struct {
+	Level1 postgres.StringExpression
+	Level2 postgres.StringExpression
+	Level3 postgres.StringExpression
+	Level4 postgres.StringExpression
+	Level5 postgres.StringExpression
+}{
+	Level1: postgres.NewEnumValue("1"),
+	Level2: postgres.NewEnumValue("2"),
+	Level3: postgres.NewEnumValue("3"),
+	Level4: postgres.NewEnumValue("4"),
+	Level5: postgres.NewEnumValue("5"),
+}
+`
+
 var allTypesModelContent = `
 package model