fix: Broken interaction between query_parameter_limit and pq.Array() (#2383)

* Understand sqlc.slice() in more expression types

This was needed to understand PostgreSQL IN() which doens't surface as
an *ast.In

* Fix interaction between query_parameter_limit and sqlc.slice()

The sqlc.slice() templating code needed to understand how to refer to a variable when query_parameter_limit prevented struct emission.

issue #2268

* Fix imports if a file has both sqlc.slice() and regular arrays

The pq import was missing in those cases. I ran into this while writing
a test case for #2268 comment 2.

* Fix broken interaction between query_parameter_limit and pq.Array()

issue #2268

Author: Jille

internal/codegen/golang/imports.go

@@ -344,12 +344,12 @@ func (i *importer) queryImports(filename string) fileImports {
 			if !q.Arg.isEmpty() {
 				if q.Arg.IsStruct() {
 					for _, f := range q.Arg.Struct.Fields {
-						if strings.HasPrefix(f.Type, "[]") && f.Type != "[]byte" {
+						if strings.HasPrefix(f.Type, "[]") && f.Type != "[]byte" && !f.HasSqlcSlice() {
 							return true
 						}
 					}
 				} else {
-					if strings.HasPrefix(q.Arg.Type(), "[]") && q.Arg.Type() != "[]byte" {
+					if strings.HasPrefix(q.Arg.Type(), "[]") && q.Arg.Type() != "[]byte" && !q.Arg.HasSqlcSlices() {
 						return true
 					}
 				}
@@ -375,7 +375,8 @@ func (i *importer) queryImports(filename string) fileImports {
 	sqlpkg := parseDriver(i.Settings.Go.SqlPackage)
 	if sqlcSliceScan() {
 		std["strings"] = struct{}{}
-	} else if sliceScan() && !sqlpkg.IsPGX() {
+	}
+	if sliceScan() && !sqlpkg.IsPGX() {
 		pkg[ImportSpec{Path: "github.com/lib/pq"}] = struct{}{}
 	}
 

internal/codegen/golang/query.go

@@ -116,11 +116,9 @@ func (v QueryValue) Params() string {
 	} else {
 		for _, f := range v.Struct.Fields {
 			if !f.HasSqlcSlice() && strings.HasPrefix(f.Type, "[]") && f.Type != "[]byte" && !v.SQLDriver.IsPGX() {
-				out = append(out, "pq.Array("+v.Name+"."+f.Name+")")
-			} else if !v.EmitStruct() && v.IsStruct() {
-				out = append(out, toLowerCase(f.Name))
+				out = append(out, "pq.Array("+v.VariableForField(f)+")")
 			} else {
-				out = append(out, v.Name+"."+f.Name)
+				out = append(out, v.VariableForField(f))
 			}
 		}
 	}
@@ -189,6 +187,16 @@ func (v QueryValue) Scan() string {
 	return "\n" + strings.Join(out, ",\n")
 }
 
+func (v QueryValue) VariableForField(f Field) string {
+	if !v.IsStruct() {
+		return v.Name
+	}
+	if !v.EmitStruct() {
+		return toLowerCase(f.Name)
+	}
+	return v.Name + "." + f.Name
+}
+
 // A struct used to generate methods and fields on the Queries struct
 type Query struct {
 	Cmd          string

internal/codegen/golang/templates/stdlib/queryCode.tmpl

@@ -126,16 +126,16 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}
             {{- $arg := .Arg }}
             {{- range .Arg.Struct.Fields }}
                 {{- if .HasSqlcSlice }}
-                    if len({{$arg.Name}}.{{.Name}}) > 0 {
-                      for _, v := range {{$arg.Name}}.{{.Name}} {
+                    if len({{$arg.VariableForField .}}) > 0 {
+                      for _, v := range {{$arg.VariableForField .}} {
                         queryParams = append(queryParams, v)
                       }
-                      query = strings.Replace(query, "/*SLICE:{{.Column.Name}}*/?", strings.Repeat(",?", len({{$arg.Name}}.{{.Name}}))[1:], 1)
+                      query = strings.Replace(query, "/*SLICE:{{.Column.Name}}*/?", strings.Repeat(",?", len({{$arg.VariableForField .}}))[1:], 1)
                     } else {
                       query = strings.Replace(query, "/*SLICE:{{.Column.Name}}*/?", "NULL", 1)
                     }
                 {{- else }}
-                  queryParams = append(queryParams, {{$arg.Name}}.{{.Name}})
+                  queryParams = append(queryParams, {{$arg.VariableForField .}})
                 {{- end }}
             {{- end }}
         {{- else }}

internal/compiler/resolve.go

@@ -153,6 +153,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 						DataType:     dataType,
 						IsNamedParam: isNamed,
 						NotNull:      p.NotNull(),
+						IsSqlcSlice:  p.IsSqlcSlice(),
 					},
 				})
 				continue
@@ -220,6 +221,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 								Length:       c.Length,
 								Table:        table,
 								IsNamedParam: isNamed,
+								IsSqlcSlice:  p.IsSqlcSlice(),
 							},
 						})
 					}
@@ -284,6 +286,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 							IsArray:      c.IsArray,
 							Table:        table,
 							IsNamedParam: isNamed,
+							IsSqlcSlice:  p.IsSqlcSlice(),
 						},
 					})
 				}
@@ -352,6 +355,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 							DataType:     "any",
 							IsNamedParam: isNamed,
 							NotNull:      p.NotNull(),
+							IsSqlcSlice:  p.IsSqlcSlice(),
 						},
 					})
 					continue
@@ -392,6 +396,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 						DataType:     dataType(paramType),
 						NotNull:      p.NotNull(),
 						IsNamedParam: isNamed,
+						IsSqlcSlice:  p.IsSqlcSlice(),
 					},
 				})
 			}
@@ -457,6 +462,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 						Table:        &ast.TableName{Schema: schema, Name: rel},
 						Length:       c.Length,
 						IsNamedParam: isNamed,
+						IsSqlcSlice:  p.IsSqlcSlice(),
 					},
 				})
 			} else {

internal/endtoend/testdata/query_parameter_limit_to_two/postgresql/go/models.go

@@ -3,7 +3,8 @@ CREATE TABLE authors (
   id   BIGSERIAL PRIMARY KEY,
   name text      NOT NULL,
   bio  text,
-  country_code CHAR(2) NOT NULL
+  country_code CHAR(2) NOT NULL,
+  titles TEXT[]
 );
 
 -- name: GetAuthor :one
@@ -16,12 +17,19 @@ ORDER BY name;
 
 -- name: CreateAuthor :one
 INSERT INTO authors (
-  name, bio, country_code
+  name, bio, country_code, titles
 ) VALUES (
-  $1, $2, $3
+  $1, $2, $3, $4
 )
 RETURNING *;
 
 -- name: DeleteAuthor :exec
 DELETE FROM authors
 WHERE id = $1;
+
+-- name: DeleteAuthors :exec
+DELETE FROM authors
+WHERE id IN (sqlc.slice(ids)) AND name = $1;
+
+-- name: CreateAuthorOnlyTitles :one
+INSERT INTO authors (name, titles) VALUES ($1, $2) RETURNING *;