upstream-sqlc/internal/codegen/golang/templates/pgx/batchCode.tmpl

{{define "batchCodePgx"}}

var (
	ErrBatchAlreadyClosed = errors.New("batch already closed")
)

{{range .GoQueries}}
{{if eq (hasPrefix .Cmd ":batch") true }}
const {{.ConstantName}} = {{$.Q}}-- name: {{.MethodName}} {{.Cmd}}
{{escape .SQL}}
{{$.Q}}

type {{.MethodName}}BatchResults struct {
    br pgx.BatchResults
    tot int
    closed bool
}

{{if .Arg.Struct}}
type {{.Arg.Type}} struct { {{- range .Arg.Struct.Fields}}
  {{.Name}} {{.Type}} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
  {{- end}}
}
{{end}}

{{if .Ret.EmitStruct}}
type {{.Ret.Type}} struct { {{- range .Ret.Struct.Fields}}
  {{.Name}} {{.Type}} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
  {{- end}}
}
{{end}}

{{range .Comments}}//{{.}}
{{end -}}
func (q *Queries) {{.MethodName}}(ctx context.Context, {{ if $.EmitMethodsWithDBArgument}}db DBTX,{{end}} {{.Arg.SlicePair}}) *{{.MethodName}}BatchResults {
    batch := &pgx.Batch{}
    for _, a := range {{index .Arg.Name}} {
        vals := []any{
        {{- if .Arg.Struct }}
        {{- range .Arg.Struct.Fields }}
            a.{{.Name}},
        {{- end }}
        {{- else }}
            a,
        {{- end }}
        }
        batch.Queue({{.ConstantName}}, vals...)
    }
    br := {{if not $.EmitMethodsWithDBArgument}}q.{{end}}db.SendBatch(ctx, batch)
    return &{{.MethodName}}BatchResults{br,len({{.Arg.Name}}),false}
}

{{if eq .Cmd ":batchexec"}}
func (b *{{.MethodName}}BatchResults) Exec(f func(int, error)) {
	defer b.br.Close()
   for t := 0; t < b.tot; t++ {
     if b.closed {
       if f != nil {
         f(t, ErrBatchAlreadyClosed)
       }
       continue
     }
     _, err := b.br.Exec()
     if f != nil {
        f(t, err)
     }
   }
}
{{end}}

{{if eq .Cmd ":batchmany"}}
func (b *{{.MethodName}}BatchResults) Query(f func(int, []{{.Ret.DefineType}}, error)) {
	defer b.br.Close()
   for t := 0; t < b.tot; t++ {
     {{- if $.EmitEmptySlices}}
     items := []{{.Ret.DefineType}}{}
     {{else}}
     var items []{{.Ret.DefineType}}
     {{end -}}
     if b.closed {
        if f != nil {
          f(t, items, ErrBatchAlreadyClosed)
        }
        continue
     }
     err := func() error {
       rows, err := b.br.Query()
       if err != nil {
         return err
       }
       defer rows.Close()
       for rows.Next() {
           var {{.Ret.Name}} {{.Ret.Type}}
           if err := rows.Scan({{.Ret.Scan}}); err != nil {
             return err
           }
           items = append(items, {{.Ret.ReturnName}})
        }
        return rows.Err()
      }()
      if f != nil {
        f(t, items, err)
      }
   }
}
{{end}}

{{if eq .Cmd ":batchone"}}
func (b *{{.MethodName}}BatchResults) QueryRow(f func(int, {{.Ret.DefineType}}, error)) {
	defer b.br.Close()
   for t := 0; t < b.tot; t++ {
     var {{.Ret.Name}} {{.Ret.Type}}
     if b.closed {
        if f != nil {
          f(t, {{if .Ret.IsPointer}}nil{{else}}{{.Ret.Name}}{{end}}, ErrBatchAlreadyClosed)
        }
        continue
     }
     row := b.br.QueryRow()
	  err := row.Scan({{.Ret.Scan}})
     if f != nil {
       f(t, {{.Ret.ReturnName}}, err)
     }
   }
}
{{end}}

func (b *{{.MethodName}}BatchResults) Close() error {
    b.closed = true
    return b.br.Close()
}
{{end}}
{{end}}
{{end}}