internal/codegen: add Enum.Valid and AllEnumValues (#1613)

Author: josharian

docs/reference/config.md

@@ -20,6 +20,8 @@ packages:
     emit_result_struct_pointers: false
     emit_params_struct_pointers: false
     emit_methods_with_db_argument: false
+    emit_enum_valid_method: false
+    emit_all_enum_values: false
     json_tags_case_style: "camel"
     output_db_file_name: "db.go"
     output_models_file_name: "models.go"
@@ -60,6 +62,12 @@ Each package document has the following keys:
   - If true, parameters are passed as pointers to structs. Defaults to `false`.
 - `emit_methods_with_db_argument`:
   - If true, generated methods will accept a DBTX argument instead of storing a DBTX on the `*Queries` struct. Defaults to `false`.
+- `emit_enum_valid_method`:
+  - If true, generate a Valid method on enum types,
+    indicating whether a string is a valid enum value.
+- `emit_all_enum_values`:
+  - If true, emit a function per enum type
+    that returns all valid enum values.
 - `json_tags_case_style`:
   - `camel` for camelCase, `pascal` for PascalCase, `snake` for snake_case or `none` to use the column name in the DB. Defaults to `none`.
 - `output_db_file_name`:

internal/cmd/shim.go

@@ -86,6 +86,8 @@ func pluginGoCode(s config.SQLGo) *plugin.GoCode {
 		EmitResultStructPointers:  s.EmitResultStructPointers,
 		EmitParamsStructPointers:  s.EmitParamsStructPointers,
 		EmitMethodsWithDbArgument: s.EmitMethodsWithDBArgument,
+		EmitEnumValidMethod:       s.EmitEnumValidMethod,
+		EmitAllEnumValues:         s.EmitAllEnumValues,
 		JsonTagsCaseStyle:         s.JSONTagsCaseStyle,
 		Package:                   s.Package,
 		Out:                       s.Out,

internal/codegen/golang/gen.go

@@ -32,6 +32,8 @@ type tmplCtx struct {
 	EmitInterface             bool
 	EmitEmptySlices           bool
 	EmitMethodsWithDBArgument bool
+	EmitEnumValidMethod       bool
+	EmitAllEnumValues         bool
 	UsesCopyFrom              bool
 	UsesBatch                 bool
 }
@@ -84,6 +86,8 @@ func generate(req *plugin.CodeGenRequest, enums []Enum, structs []Struct, querie
 		EmitPreparedQueries:       golang.EmitPreparedQueries,
 		EmitEmptySlices:           golang.EmitEmptySlices,
 		EmitMethodsWithDBArgument: golang.EmitMethodsWithDbArgument,
+		EmitEnumValidMethod:       golang.EmitEnumValidMethod,
+		EmitAllEnumValues:         golang.EmitAllEnumValues,
 		UsesCopyFrom:              usesCopyFrom(queries),
 		UsesBatch:                 usesBatch(queries),
 		SQLPackage:                SQLPackageFromString(golang.SqlPackage),

internal/codegen/golang/templates/template.tmpl

@@ -86,6 +86,23 @@ func (e *{{.Name}}) Scan(src interface{}) error {
 	}
 	return nil
 }
+
+{{ if $.EmitEnumValidMethod }}
+func (e {{.Name}}) Valid() bool {
+  switch e {
+  case {{ range $idx, $name := .Constants }}{{ if ne $idx 0 }},{{ "\n" }}{{ end }}{{ .Name }}{{ end }}:
+    return true
+  }
+  return false
+}
+{{ end }}
+
+{{ if $.EmitAllEnumValues }}
+func All{{ .Name }}Values() []{{ .Name }} {
+	return []{{ .Name }}{ {{ range .Constants}}{{ "\n" }}{{ .Name }},{{ end }}
+	}
+}
+{{ end }}
 {{end}}
 
 {{range .Structs}}

internal/config/config.go

@@ -128,6 +128,8 @@ type SQLGo struct {
 	EmitResultStructPointers  bool              `json:"emit_result_struct_pointers" yaml:"emit_result_struct_pointers"`
 	EmitParamsStructPointers  bool              `json:"emit_params_struct_pointers" yaml:"emit_params_struct_pointers"`
 	EmitMethodsWithDBArgument bool              `json:"emit_methods_with_db_argument,omitempty" yaml:"emit_methods_with_db_argument"`
+	EmitEnumValidMethod       bool              `json:"emit_enum_valid_method,omitempty" yaml:"emit_enum_valid_method"`
+	EmitAllEnumValues         bool              `json:"emit_all_enum_values,omitempty" yaml:"emit_all_enum_values"`
 	JSONTagsCaseStyle         string            `json:"json_tags_case_style,omitempty" yaml:"json_tags_case_style"`
 	Package                   string            `json:"package" yaml:"package"`
 	Out                       string            `json:"out" yaml:"out"`

internal/config/v_one.go

@@ -32,6 +32,8 @@ type v1PackageSettings struct {
 	EmitResultStructPointers  bool       `json:"emit_result_struct_pointers" yaml:"emit_result_struct_pointers"`
 	EmitParamsStructPointers  bool       `json:"emit_params_struct_pointers" yaml:"emit_params_struct_pointers"`
 	EmitMethodsWithDBArgument bool       `json:"emit_methods_with_db_argument" yaml:"emit_methods_with_db_argument"`
+	EmitEnumValidMethod       bool       `json:"emit_enum_valid_method,omitempty" yaml:"emit_enum_valid_method"`
+	EmitAllEnumValues         bool       `json:"emit_all_enum_values,omitempty" yaml:"emit_all_enum_values"`
 	JSONTagsCaseStyle         string     `json:"json_tags_case_style,omitempty" yaml:"json_tags_case_style"`
 	SQLPackage                string     `json:"sql_package" yaml:"sql_package"`
 	Overrides                 []Override `json:"overrides" yaml:"overrides"`
@@ -126,6 +128,8 @@ func (c *V1GenerateSettings) Translate() Config {
 					EmitResultStructPointers:  pkg.EmitResultStructPointers,
 					EmitParamsStructPointers:  pkg.EmitParamsStructPointers,
 					EmitMethodsWithDBArgument: pkg.EmitMethodsWithDBArgument,
+					EmitEnumValidMethod:       pkg.EmitEnumValidMethod,
+					EmitAllEnumValues:         pkg.EmitAllEnumValues,
 					Package:                   pkg.Name,
 					Out:                       pkg.Path,
 					SQLPackage:                pkg.SQLPackage,

internal/endtoend/testdata/codegen_json/gen/codegen_request.json

@@ -41,7 +41,9 @@
       "output_db_file_name": "",
       "output_models_file_name": "",
       "output_querier_file_name": "",
-      "output_files_suffix": ""
+      "output_files_suffix": "",
+      "emit_enum_valid_method": false,
+      "emit_all_enum_values": false
     },
     "json": {
       "out": "gen",

internal/endtoend/testdata/emit_enum_valid_and_values/go/db.go

@@ -0,0 +1,12 @@
+CREATE TYPE ip_protocol AS enum ('tcp', 'ip', 'icmp');
+
+CREATE TABLE bar_old (id_old serial not null, ip_old ip_protocol NOT NULL);
+
+-- name: ListFoo :many
+SELECT id_old as foo_old, id_old as baz_old
+FROM bar_old
+WHERE ip_old = $1 AND id_old = $2;
+
+-- name: ListBar :many
+SELECT * FROM bar_old;
+

internal/endtoend/testdata/emit_enum_valid_and_values/go/models.go

@@ -0,0 +1,22 @@
+{
+  "version": "1",
+  "packages": [
+    {
+      "path": "go",
+      "engine": "postgresql",
+      "sql_package": "pgx/v4",
+      "name": "querytest",
+      "schema": "query.sql",
+      "queries": "query.sql",
+      "emit_enum_valid_method": true,
+      "emit_all_enum_values": true
+    }
+  ],
+  "rename": {
+    "id_old": "IDNew",
+    "bar_old": "BarNew",
+    "foo_old": "FooNew",
+    "ip_protocol": "IPProtocol",
+    "ip_protocol_tcp": "IPProtocolTCP"
+  }
+}