feat: add support for specifying string enum values (#165)

* feat: add support for specifying string enum values * add Quote handling
abice · Jan 30, 2023 · b9e7d1a · b9e7d1a
1 parent 1e5efaa
commit b9e7d1a
Show file tree

Hide file tree

Showing 9 changed files with 300 additions and 44 deletions.
diff --git a/_example/sql_str_int.go b/_example/sql_str_int.go
@@ -4,3 +4,6 @@ package example
 
 // ENUM(_,zeus, apollo, athena=20, ares)
 type GreekGod string
+
+// ENUM(_,zeus, apollo, _=19, athena="20", ares)
+type GreekGodCustom string
diff --git a/_example/sql_str_int_enum.go b/_example/sql_str_int_enum.go
diff --git a/_example/strings_only.go b/_example/strings_only.go
@@ -2,5 +2,5 @@
 
 package example
 
-// ENUM(pending, running, completed, failed)
+// ENUM(pending, running, completed, failed=error)
 type StrState string
diff --git a/_example/strings_only_enum.go b/_example/strings_only_enum.go
diff --git a/_example/strings_only_test.go b/_example/strings_only_test.go
@@ -20,7 +20,7 @@ func TestStrState(t *testing.T) {
 func TestStrStateMustParse(t *testing.T) {
 	x := `avocado`
 
-	assert.PanicsWithError(t, x+" is not a valid StrState, try [pending, running, completed, failed]", func() { MustParseStrState(x) })
+	assert.PanicsWithError(t, x+" is not a valid StrState, try [pending, running, completed, error]", func() { MustParseStrState(x) })
 	assert.NotPanics(t, func() { MustParseStrState(StrStateFailed.String()) })
 }
 
@@ -78,14 +78,14 @@ func TestStrStateUnmarshal(t *testing.T) {
 		},
 		{
 			name:          "failed",
-			input:         `{"state":"Failed"}`,
+			input:         `{"state":"Error"}`,
 			output:        &testData{StrStateX: StrStateFailed},
 			errorExpected: false,
 			err:           nil,
 		},
 		{
 			name:          "failedlower",
-			input:         `{"state":"failed"}`,
+			input:         `{"state":"error"}`,
 			output:        &testData{StrStateX: StrStateFailed},
 			errorExpected: false,
 			err:           nil,
@@ -141,7 +141,7 @@ func TestStrStateMarshal(t *testing.T) {
 		},
 		{
 			name:          "green",
-			output:        `{"state":"failed"}`,
+			output:        `{"state":"error"}`,
 			input:         &testData{StrStateX: StrStateFailed},
 			errorExpected: false,
 			err:           nil,

diff --git a/generator/enum_string.tmpl b/generator/enum_string.tmpl
@@ -10,7 +10,7 @@ const (
 	{{- if $value.Comment}}
 	// {{$value.Comment}}
 	{{- end}}
-    {{$value.PrefixedName}} {{$enumName}} = "{{$value.RawName}}"
+    {{$value.PrefixedName}} {{$enumName}} = "{{$value.ValueStr}}"
 {{- end}}
 )
 {{if .names -}}
@@ -100,7 +100,7 @@ func (x *{{.enum.Name}}) UnmarshalText(text []byte) error {
 }
 {{end}}
 
-{{ if or .sql .sqlnullstr .sqlint .sqlnullint }}
+{{ if .anySQLEnabled }}
 var err{{.enum.Name}}NilPtr = errors.New("value pointer is nil") // one per type for package clashes
 {{ end }}
 
@@ -136,8 +136,8 @@ func (x *{{.enum.Name}}) Scan(value interface{}) (err error) {
 	default:
 		return errors.New("invalid type for {{.enum.Name}}")
 	}
-	
-	return 
+
+	return
 }
 
 // Value implements the driver Valuer interface.
@@ -149,12 +149,12 @@ func (x {{.enum.Name}}) Value() (driver.Value, error) {
 {{/* SQL stored as an integer value */}}
 {{ if or .sqlint .sqlnullint }}
 var sqlInt{{.enum.Name}}Map = map[int64]{{.enum.Name}}{ {{ range $rIndex, $value := .enum.Values }}{{ if ne $value.Name "_"}}
-{{ $value.Value }}: {{ $value.PrefixedName }},{{end}}
+{{ $value.ValueInt }}: {{ $value.PrefixedName }},{{end}}
 {{- end}}
 }
 
 var sqlInt{{.enum.Name}}Value = map[{{.enum.Name}}]int64{ {{ range $rIndex, $value := .enum.Values }}{{ if ne $value.Name "_"}}
-	{{ $value.PrefixedName }}: {{ $value.Value }},{{end}}
+	{{ $value.PrefixedName }}: {{ $value.ValueInt }},{{end}}
 {{- end}}
 }
 
@@ -170,7 +170,7 @@ func lookupSqlInt{{.enum.Name}}(val int64) ({{.enum.Name}}, error){
 func (x *{{.enum.Name}}) Scan(value interface{}) (err error) {
 	if value == nil {
 		*x = {{.enum.Name}}("")
-		return 
+		return
 	}
 
 	// A wider range of scannable types.
@@ -235,8 +235,8 @@ func (x *{{.enum.Name}}) Scan(value interface{}) (err error) {
 	default:
 		return errors.New("invalid type for {{.enum.Name}}")
 	}
-	
-	return 
+
+	return
 }
 
 // Value implements the driver Valuer interface.