From c2aaf39c44fa21fc818a45b4c74c24354567aeb4 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyle@sqlc.dev>
Date: Wed, 12 Jul 2023 17:31:29 -0700
Subject: [PATCH] fix(golang): Prevent duplicate enum output

---
 internal/codegen/golang/gen.go                         |  7 +++----
 .../omit_unused_structs/postgresql/stdlib/go/models.go |  6 ++++++
 .../postgresql/stdlib/go/query.sql.go                  | 10 +++++-----
 .../omit_unused_structs/postgresql/stdlib/query.sql    |  2 +-
 .../omit_unused_structs/postgresql/stdlib/schema.sql   |  1 +
 5 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/internal/codegen/golang/gen.go b/internal/codegen/golang/gen.go
index 1d3ba6e13d..260295d876 100644
--- a/internal/codegen/golang/gen.go
+++ b/internal/codegen/golang/gen.go
@@ -318,10 +318,9 @@ func filterUnusedStructs(enums []Enum, structs []Struct, queries []Query) ([]Enu
 
 	keepEnums := make([]Enum, 0, len(enums))
 	for _, enum := range enums {
-		if _, ok := keepTypes[enum.Name]; ok {
-			keepEnums = append(keepEnums, enum)
-		}
-		if _, ok := keepTypes["Null"+enum.Name]; ok {
+		_, keep := keepTypes[enum.Name]
+		_, keepNull := keepTypes["Null"+enum.Name]
+		if keep || keepNull {
 			keepEnums = append(keepEnums, enum)
 		}
 	}
diff --git a/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/go/models.go b/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/go/models.go
index d8a488d9a0..3881f0c784 100644
--- a/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/go/models.go
+++ b/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/go/models.go
@@ -219,6 +219,12 @@ func (ns NullQueryReturnStructEnumTableEnum) Value() (driver.Value, error) {
 	return string(ns.QueryReturnStructEnumTableEnum), nil
 }
 
+type QueryParamEnumTable struct {
+	ID    int32
+	Other QueryParamEnumTableEnum
+	Value NullQueryParamEnumTableEnum
+}
+
 type QueryReturnFullTable struct {
 	ID    int32
 	Value NullQueryReturnFullTableEnum
diff --git a/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/go/query.sql.go b/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/go/query.sql.go
index bb648be939..eabf22aac8 100644
--- a/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/go/query.sql.go
+++ b/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/go/query.sql.go
@@ -11,14 +11,14 @@ import (
 )
 
 const query_param_enum_table = `-- name: query_param_enum_table :one
-SELECT id FROM query_param_enum_table WHERE value = $1
+SELECT id, other, value FROM query_param_enum_table WHERE value = $1
 `
 
-func (q *Queries) query_param_enum_table(ctx context.Context, value NullQueryParamEnumTableEnum) (int32, error) {
+func (q *Queries) query_param_enum_table(ctx context.Context, value NullQueryParamEnumTableEnum) (QueryParamEnumTable, error) {
 	row := q.db.QueryRowContext(ctx, query_param_enum_table, value)
-	var id int32
-	err := row.Scan(&id)
-	return id, err
+	var i QueryParamEnumTable
+	err := row.Scan(&i.ID, &i.Other, &i.Value)
+	return i, err
 }
 
 const query_param_struct_enum_table = `-- name: query_param_struct_enum_table :one
diff --git a/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/query.sql b/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/query.sql
index 7105779d27..1420d869d3 100644
--- a/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/query.sql
+++ b/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/query.sql
@@ -2,7 +2,7 @@
 SELECT * FROM query_return_full_table;
 
 -- name: query_param_enum_table :one
-SELECT id FROM query_param_enum_table WHERE value = $1;
+SELECT * FROM query_param_enum_table WHERE value = $1;
 
 -- name: query_param_struct_enum_table :one
 SELECT id FROM query_param_struct_enum_table WHERE id = $1 AND value = $2;
diff --git a/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/schema.sql b/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/schema.sql
index e51554f040..cc2ae14b28 100644
--- a/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/schema.sql
+++ b/internal/endtoend/testdata/omit_unused_structs/postgresql/stdlib/schema.sql
@@ -23,6 +23,7 @@ CREATE TYPE query_param_enum_table_enum AS ENUM (
 );
 CREATE TABLE query_param_enum_table (
     id      INTEGER PRIMARY KEY,
+    other   query_param_enum_table_enum NOT NULL,
     value   query_param_enum_table_enum
 );