From 55fb51afc48710304a7c63d21ae5474551aa4ea0 Mon Sep 17 00:00:00 2001
From: Ryu Yamada <a2sin2a2ko@kyoto-y.work>
Date: Sat, 5 Aug 2023 17:19:25 +0900
Subject: [PATCH] fix: compare define type in codegen (#2263)

---
 .../codegen/golang/templates/pgx/queryCode.tmpl  |  2 +-
 .../golang/templates/stdlib/queryCode.tmpl       |  2 +-
 .../postgresql/go/models.go                      |  5 +++++
 .../postgresql/go/query.sql.go                   | 16 ++++++++++++++++
 .../postgresql/query.sql                         | 13 +++++++++++++
 5 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/internal/codegen/golang/templates/pgx/queryCode.tmpl b/internal/codegen/golang/templates/pgx/queryCode.tmpl
index 1736fa11f7..18de5db2ba 100644
--- a/internal/codegen/golang/templates/pgx/queryCode.tmpl
+++ b/internal/codegen/golang/templates/pgx/queryCode.tmpl
@@ -33,7 +33,7 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
 	row := q.db.QueryRow(ctx, {{.ConstantName}}, {{.Arg.Params}})
 {{- end}}
-	{{- if ne .Arg.Pair .Ret.Pair }}
+	{{- if or (ne .Arg.Pair .Ret.Pair) (ne .Arg.DefineType .Ret.DefineType) }}
 	var {{.Ret.Name}} {{.Ret.Type}}
 	{{- end}}
 	err := row.Scan({{.Ret.Scan}})
diff --git a/internal/codegen/golang/templates/stdlib/queryCode.tmpl b/internal/codegen/golang/templates/stdlib/queryCode.tmpl
index 6eca49cd17..cf56000ec6 100644
--- a/internal/codegen/golang/templates/stdlib/queryCode.tmpl
+++ b/internal/codegen/golang/templates/stdlib/queryCode.tmpl
@@ -24,7 +24,7 @@ type {{.Ret.Type}} struct { {{- range .Ret.Struct.Fields}}
 {{end -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
     {{- template "queryCodeStdExec" . }}
-	{{- if ne .Arg.Pair .Ret.Pair }}
+	{{- if or (ne .Arg.Pair .Ret.Pair) (ne .Arg.DefineType .Ret.DefineType) }}
 	var {{.Ret.Name}} {{.Ret.Type}}
 	{{- end}}
 	err := row.Scan({{.Ret.Scan}})
diff --git a/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/models.go b/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/models.go
index 7fbbf5fbc7..683ff1177d 100644
--- a/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/models.go
+++ b/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/models.go
@@ -15,3 +15,8 @@ type Author struct {
 	CountryCode string
 	Titles      []string
 }
+
+type Client struct {
+	ID   int32
+	Name string
+}
diff --git a/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/query.sql.go b/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/query.sql.go
index 867290c049..19a289d2f6 100644
--- a/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/query.sql.go
+++ b/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/query.sql.go
@@ -13,6 +13,22 @@ import (
 	"github.com/lib/pq"
 )
 
+const addNewClient = `-- name: AddNewClient :one
+INSERT INTO clients (
+  id, name
+) VALUES (
+  $1, $2
+)
+RETURNING id, name
+`
+
+func (q *Queries) AddNewClient(ctx context.Context, iD int32, name string) (Client, error) {
+	row := q.db.QueryRowContext(ctx, addNewClient, iD, name)
+	var i Client
+	err := row.Scan(&i.ID, &i.Name)
+	return i, err
+}
+
 const createAuthor = `-- name: CreateAuthor :one
 INSERT INTO authors (
   name, bio, country_code, titles
diff --git a/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/query.sql b/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/query.sql
index 8edd6ae9b9..13ee235bd6 100644
--- a/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/query.sql
+++ b/internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/query.sql
@@ -33,3 +33,16 @@ WHERE id IN (sqlc.slice(ids)) AND name = $1;
 
 -- name: CreateAuthorOnlyTitles :one
 INSERT INTO authors (name, titles) VALUES ($1, $2) RETURNING *;
+
+CREATE TABLE clients (
+  id INT PRIMARY KEY,
+  name TEXT NOT NULL
+);
+
+-- name: AddNewClient :one
+INSERT INTO clients (
+  id, name
+) VALUES (
+  $1, $2
+)
+RETURNING *;
\ No newline at end of file