fix(codegen/golang): Refactor imports code to match templates (#2709)

Author: kyleconroy

internal/codegen/golang/gen.go

@@ -303,7 +303,7 @@ func usesBatch(queries []Query) bool {
 
 func checkNoTimesForMySQLCopyFrom(queries []Query) error {
 	for _, q := range queries {
-		for _, f := range q.Arg.Fields() {
+		for _, f := range q.Arg.CopyFromMySQLFields() {
 			if f.Type == "time.Time" {
 				return fmt.Errorf("values with a timezone are not yet supported")
 			}

internal/codegen/golang/imports.go

@@ -242,11 +242,9 @@ func (i *importer) interfaceImports() fileImports {
 					return true
 				}
 			}
-			if !q.Arg.isEmpty() {
-				for _, f := range q.Arg.Fields() {
-					if hasPrefixIgnoringSliceAndPointerPrefix(f.Type, name) {
-						return true
-					}
+			for _, f := range q.Arg.Pairs() {
+				if hasPrefixIgnoringSliceAndPointerPrefix(f.Type, name) {
+					return true
 				}
 			}
 		}
@@ -312,13 +310,20 @@ func (i *importer) queryImports(filename string) fileImports {
 					return true
 				}
 			}
-			if !q.Arg.isEmpty() {
-				for _, f := range q.Arg.Fields() {
+			// Check the fields of the argument struct if it's emitted
+			if q.Arg.EmitStruct() {
+				for _, f := range q.Arg.Struct.Fields {
 					if hasPrefixIgnoringSliceAndPointerPrefix(f.Type, name) {
 						return true
 					}
 				}
 			}
+			// Check the argument pairs inside the method definition
+			for _, f := range q.Arg.Pairs() {
+				if hasPrefixIgnoringSliceAndPointerPrefix(f.Type, name) {
+					return true
+				}
+			}
 		}
 		return false
 	})
@@ -441,15 +446,15 @@ func (i *importer) batchImports() fileImports {
 					return true
 				}
 			}
-			if !q.Arg.isEmpty() {
-				if q.Arg.EmitStruct() {
-					for _, f := range q.Arg.Struct.Fields {
-						if hasPrefixIgnoringSliceAndPointerPrefix(f.Type, name) {
-							return true
-						}
+			if q.Arg.EmitStruct() {
+				for _, f := range q.Arg.Struct.Fields {
+					if hasPrefixIgnoringSliceAndPointerPrefix(f.Type, name) {
+						return true
 					}
 				}
-				if hasPrefixIgnoringSliceAndPointerPrefix(q.Arg.Type(), name) {
+			}
+			for _, f := range q.Arg.Pairs() {
+				if hasPrefixIgnoringSliceAndPointerPrefix(f.Type, name) {
 					return true
 				}
 			}

internal/codegen/golang/query.go

@@ -38,21 +38,41 @@ func (v QueryValue) isEmpty() bool {
 	return v.Typ == "" && v.Name == "" && v.Struct == nil
 }
 
+type Argument struct {
+	Name string
+	Type string
+}
+
 func (v QueryValue) Pair() string {
-	if v.isEmpty() {
-		return ""
+	var out []string
+	for _, arg := range v.Pairs() {
+		out = append(out, arg.Name+" "+arg.Type)
 	}
+	return strings.Join(out, ",")
+}
 
-	var out []string
+// Return the argument name and type for query methods. Should only be used in
+// the context of method arguments.
+func (v QueryValue) Pairs() []Argument {
+	if v.isEmpty() {
+		return nil
+	}
 	if !v.EmitStruct() && v.IsStruct() {
+		var out []Argument
 		for _, f := range v.Struct.Fields {
-			out = append(out, toLowerCase(f.Name)+" "+f.Type)
+			out = append(out, Argument{
+				Name: toLowerCase(f.Name),
+				Type: f.Type,
+			})
 		}
-
-		return strings.Join(out, ",")
+		return out
+	}
+	return []Argument{
+		{
+			Name: v.Name,
+			Type: v.DefineType(),
+		},
 	}
-
-	return v.Name + " " + v.DefineType()
 }
 
 func (v QueryValue) SlicePair() string {
@@ -202,7 +222,11 @@ func (v QueryValue) Scan() string {
 	return "\n" + strings.Join(out, ",\n")
 }
 
-func (v QueryValue) Fields() []Field {
+// Deprecated: This method does not respect the Emit field set on the
+// QueryValue. It's used by the go-sql-driver-mysql/copyfromCopy.tmpl and should
+// not be used other places.
+func (v QueryValue) CopyFromMySQLFields() []Field {
+	// fmt.Printf("%#v\n", v)
 	if v.Struct != nil {
 		return v.Struct.Fields
 	}

internal/codegen/golang/templates/go-sql-driver-mysql/copyfromCopy.tmpl

@@ -4,16 +4,16 @@
 var readerHandlerSequenceFor{{.MethodName}} uint32 = 1
 
 func convertRowsFor{{.MethodName}}(w *io.PipeWriter, {{.Arg.SlicePair}}) {
-	e := mysqltsv.NewEncoder(w, {{ len .Arg.Fields }}, nil)
+	e := mysqltsv.NewEncoder(w, {{ len .Arg.CopyFromMySQLFields }}, nil)
 	for _, row := range {{.Arg.Name}} {
 {{- with $arg := .Arg }}
-{{- range $arg.Fields}}
+{{- range $arg.CopyFromMySQLFields}}
 {{- if eq .Type "string"}}
-	e.AppendString({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
+	e.AppendString({{if eq (len $arg.CopyFromMySQLFields) 1}}row{{else}}row.{{.Name}}{{end}})
 {{- else if eq .Type "[]byte"}}
-	e.AppendBytes({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
+	e.AppendBytes({{if eq (len $arg.CopyFromMySQLFields) 1}}row{{else}}row.{{.Name}}{{end}})
 {{- else}}
-	e.AppendValue({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
+	e.AppendValue({{if eq (len $arg.CopyFromMySQLFields) 1}}row{{else}}row.{{.Name}}{{end}})
 {{- end}}
 {{- end}}
 {{- end}}

internal/endtoend/testdata/query_parameter_limit_to_zero/postgresql/go/querier.go

@@ -7,8 +7,9 @@
         "name": "querytest",
         "schema": "query.sql",
         "queries": "query.sql",
-        "query_parameter_limit": 0
+        "query_parameter_limit": 0,
+        "emit_interface": true
       }
     ]
   }
-  
+