Implement support for pgx's CopyFrom

This allows type-safe bulk loading with great performance.

I didn't implement it for Python and Kotlin.

This change is fully backwards compatible, as it only changes the DBTX
interface when someone adds their first :copyFrom query (at which point
it's reasonable to require the CopyFrom method on the DBTX).

Author: Jille

docs/howto/insert.md

@@ -136,3 +136,29 @@ func (q *Queries) DeleteAuthor(ctx context.Context, id int) (Author, error) {
 	return i, err
 }
 ```
+
+## Using CopyFrom
+
+PostgreSQL supports the Copy Protocol that can insert rows a lot faster than sequential inserts. You can use this easily with sqlc:
+
+```sql
+CREATE TABLE authors (
+  id         SERIAL PRIMARY KEY,
+  name       text   NOT NULL,
+  bio        text   NOT NULL
+);
+
+-- name: CreateAuthors :copyFrom
+INSERT INTO authors (name, bio) VALUES ($1, $2);
+```
+
+```go
+type CreateAuthorsParams struct {
+  Name string
+  Bio  string
+}
+
+func (q *Queries) CreateAuthors(ctx context.Context, arg []CreateAuthorsParams) (int64, error) {
+  return q.db.CopyFrom(ctx, []string{"authors"}, []string{"name", "bio"}, &iteratorForCreateAuthors{rows: arg})
+}
+```

internal/codegen/golang/field.go

@@ -9,7 +9,8 @@ import (
 )
 
 type Field struct {
-	Name    string
+	Name    string // CamelCased name for Go
+	DBName  string // Name as used in the DB
 	Type    string
 	Tags    map[string]string
 	Comment string

internal/codegen/golang/gen.go

@@ -11,6 +11,7 @@ import (
 	"github.com/kyleconroy/sqlc/internal/codegen"
 	"github.com/kyleconroy/sqlc/internal/compiler"
 	"github.com/kyleconroy/sqlc/internal/config"
+	"github.com/kyleconroy/sqlc/internal/metadata"
 )
 
 type Generateable interface {
@@ -37,6 +38,7 @@ type tmplCtx struct {
 	EmitInterface             bool
 	EmitEmptySlices           bool
 	EmitMethodsWithDBArgument bool
+	UsesCopyFrom              bool
 }
 
 func (t *tmplCtx) OutputQuery(sourceName string) bool {
@@ -87,6 +89,7 @@ func generate(settings config.CombinedSettings, enums []Enum, structs []Struct,
 		EmitPreparedQueries:       golang.EmitPreparedQueries,
 		EmitEmptySlices:           golang.EmitEmptySlices,
 		EmitMethodsWithDBArgument: golang.EmitMethodsWithDBArgument,
+		UsesCopyFrom:              usesCopyFrom(queries),
 		SQLPackage:                SQLPackageFromString(golang.SQLPackage),
 		Q:                         "`",
 		Package:                   golang.Package,
@@ -160,3 +163,12 @@ func generate(settings config.CombinedSettings, enums []Enum, structs []Struct,
 	}
 	return output, nil
 }
+
+func usesCopyFrom(queries []Query) bool {
+	for _, q := range queries {
+		if q.Cmd == metadata.CmdCopyFrom {
+			return true
+		}
+	}
+	return false
+}

internal/codegen/golang/query.go

@@ -1,9 +1,11 @@
 package golang
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/kyleconroy/sqlc/internal/metadata"
+	"github.com/kyleconroy/sqlc/internal/sql/ast"
 )
 
 type QueryValue struct {
@@ -38,6 +40,13 @@ func (v QueryValue) Pair() string {
 	return v.Name + " " + v.DefineType()
 }
 
+func (v QueryValue) SlicePair() string {
+	if v.isEmpty() {
+		return ""
+	}
+	return v.Name + " []" + v.DefineType()
+}
+
 func (v QueryValue) Type() string {
 	if v.Typ != "" {
 		return v.Typ
@@ -105,6 +114,17 @@ func (v QueryValue) Params() string {
 	return "\n" + strings.Join(out, ",\n")
 }
 
+func (v QueryValue) ColumnNames() string {
+	if v.Struct == nil {
+		return fmt.Sprintf("[]string{%q}", v.Name)
+	}
+	escapedNames := make([]string, len(v.Struct.Fields))
+	for i, f := range v.Struct.Fields {
+		escapedNames[i] = fmt.Sprintf("%q", f.DBName)
+	}
+	return "[]string{" + strings.Join(escapedNames, ", ") + "}"
+}
+
 func (v QueryValue) Scan() string {
 	var out []string
 	if v.Struct == nil {
@@ -140,9 +160,21 @@ type Query struct {
 	SourceName   string
 	Ret          QueryValue
 	Arg          QueryValue
+	// Used for :copyFrom
+	Table *ast.TableName
 }
 
 func (q Query) hasRetType() bool {
 	scanned := q.Cmd == metadata.CmdOne || q.Cmd == metadata.CmdMany
 	return scanned && !q.Ret.isEmpty()
 }
+
+func (q Query) TableIdentifier() string {
+	escapedNames := make([]string, 0, 3)
+	for _, p := range []string{q.Table.Catalog, q.Table.Schema, q.Table.Name} {
+		if p != "" {
+			escapedNames = append(escapedNames, fmt.Sprintf("%q", p))
+		}
+	}
+	return "[]string{" + strings.Join(escapedNames, ", ") + "}"
+}

internal/codegen/golang/result.go

@@ -160,6 +160,7 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 			SourceName:   query.Filename,
 			SQL:          query.SQL,
 			Comments:     query.Comments,
+			Table:        query.InsertIntoTable,
 		}
 		sqlpkg := SQLPackageFromString(settings.Go.SQLPackage)
 
@@ -291,9 +292,10 @@ func columnsToStruct(r *compiler.Result, name string, columns []goColumn, settin
 			tags["json:"] = JSONTagName(tagName, settings)
 		}
 		gs.Fields = append(gs.Fields, Field{
-			Name: fieldName,
-			Type: goType(r, c.Column, settings),
-			Tags: tags,
+			Name:   fieldName,
+			DBName: colName,
+			Type:   goType(r, c.Column, settings),
+			Tags:   tags,
 		})
 		if _, found := seen[baseFieldName]; !found {
 			seen[baseFieldName] = []int{i}

internal/codegen/golang/templates/pgx/dbCode.tmpl

@@ -4,6 +4,9 @@ type DBTX interface {
 	Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error)
 	Query(context.Context, string, ...interface{}) (pgx.Rows, error)
 	QueryRow(context.Context, string, ...interface{}) pgx.Row
+{{- if .UsesCopyFrom }}
+	CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error)
+{{- end }}
 }
 
 {{ if .EmitMethodsWithDBArgument}}

internal/codegen/golang/templates/pgx/interfaceCode.tmpl

@@ -27,6 +27,11 @@
         {{- else if eq .Cmd ":execresult" }}
             {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (pgconn.CommandTag, error)
         {{- end}}
+        {{- if and (eq .Cmd ":copyFrom") ($dbtxParam) }}
+            {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (int64, error)
+        {{- else if eq .Cmd ":copyFrom" }}
+            {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (int64, error)
+        {{- end}}
     {{- end}}
     }
 

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

@@ -1,9 +1,11 @@
 {{define "queryCodePgx"}}
 {{range .GoQueries}}
 {{if $.OutputQuery .SourceName}}
+{{if ne .Cmd ":copyFrom"}}
 const {{.ConstantName}} = {{$.Q}}-- name: {{.MethodName}} {{.Cmd}}
 {{escape .SQL}}
 {{$.Q}}
+{{end}}
 
 {{if .Arg.EmitStruct}}
 type {{.Arg.Type}} struct { {{- range .Arg.Struct.Fields}}
@@ -112,6 +114,53 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (pgconn.Co
 }
 {{end}}
 
+{{if eq .Cmd ":copyFrom"}}
+// iteratorFor{{.MethodName}} implements pgx.CopyFromSource.
+type iteratorFor{{.MethodName}} struct {
+	rows []{{.Arg.DefineType}}
+	skippedFirstNextCall bool
+}
+
+func (r *iteratorFor{{.MethodName}}) Next() bool {
+	if len(r.rows) == 0 {
+		return false
+	}
+	if !r.skippedFirstNextCall {
+		r.skippedFirstNextCall = true
+		return true
+	}
+	r.rows = r.rows[1:]
+	return len(r.rows) > 0
+}
+
+func (r iteratorFor{{.MethodName}}) Values() ([]interface{}, error) {
+	return []interface{}{
+{{- if .Arg.Struct }}
+{{- range .Arg.Struct.Fields }}
+		r.rows[0].{{.Name}},
+{{- end }}
+{{- else }}
+		r.rows[0],
+{{- end }}
+	}, nil
+}
+
+func (r iteratorFor{{.MethodName}}) Err() error {
+	return nil
+}
+
+{{range .Comments}}//{{.}}
+{{end -}}
+{{- if $.EmitMethodsWithDBArgument}}
+func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.SlicePair}}) (int64, error) {
+	return db.CopyFrom(ctx, {{.TableIdentifier}}, {{.Arg.ColumnNames}}, &iteratorFor{{.MethodName}}{rows: {{.Arg.Name}}})
+{{- else}}
+func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.SlicePair}}) (int64, error) {
+	return q.db.CopyFrom(ctx, {{.TableIdentifier}}, {{.Arg.ColumnNames}}, &iteratorFor{{.MethodName}}{rows: {{.Arg.Name}}})
+{{- end}}
+}
+{{end}}
+
 {{end}}
 {{end}}
 {{end}}

internal/codegen/kotlin/gen.go

@@ -3,6 +3,7 @@ package kotlin
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"fmt"
 	"regexp"
 	"sort"
@@ -14,6 +15,7 @@ import (
 	"github.com/kyleconroy/sqlc/internal/config"
 	"github.com/kyleconroy/sqlc/internal/core"
 	"github.com/kyleconroy/sqlc/internal/inflection"
+	"github.com/kyleconroy/sqlc/internal/metadata"
 	"github.com/kyleconroy/sqlc/internal/sql/ast"
 	"github.com/kyleconroy/sqlc/internal/sql/catalog"
 )
@@ -458,7 +460,7 @@ func jdbcSQL(s string, engine config.Engine) string {
 	return s
 }
 
-func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs []Struct) []Query {
+func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs []Struct) ([]Query, error) {
 	qs := make([]Query, 0, len(r.Queries))
 	for _, query := range r.Queries {
 		if query.Name == "" {
@@ -467,6 +469,9 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 		if query.Cmd == "" {
 			continue
 		}
+		if query.Cmd == metadata.CmdCopyFrom {
+			return nil, errors.New("Support for CopyFrom in Kotlin is not implemented")
+		}
 
 		gq := Query{
 			Cmd:          query.Cmd,
@@ -543,7 +548,7 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 		qs = append(qs, gq)
 	}
 	sort.Slice(qs, func(i, j int) bool { return qs[i].MethodName < qs[j].MethodName })
-	return qs
+	return qs, nil
 }
 
 var ktIfaceTmpl = `// Code generated by sqlc. DO NOT EDIT.
@@ -769,7 +774,10 @@ func ktFormat(s string) string {
 func Generate(r *compiler.Result, settings config.CombinedSettings) (map[string]string, error) {
 	enums := buildEnums(r, settings)
 	structs := buildDataClasses(r, settings)
-	queries := buildQueries(r, settings, structs)
+	queries, err := buildQueries(r, settings, structs)
+	if err != nil {
+		return nil, err
+	}
 
 	i := &importer{
 		Settings:    settings,

internal/codegen/python/gen.go

@@ -1,6 +1,7 @@
 package python
 
 import (
+	"errors"
 	"fmt"
 	"log"
 	"regexp"
@@ -12,6 +13,7 @@ import (
 	"github.com/kyleconroy/sqlc/internal/config"
 	"github.com/kyleconroy/sqlc/internal/core"
 	"github.com/kyleconroy/sqlc/internal/inflection"
+	"github.com/kyleconroy/sqlc/internal/metadata"
 	pyast "github.com/kyleconroy/sqlc/internal/python/ast"
 	"github.com/kyleconroy/sqlc/internal/python/poet"
 	pyprint "github.com/kyleconroy/sqlc/internal/python/printer"
@@ -390,7 +392,7 @@ func sqlalchemySQL(s string, engine config.Engine) string {
 	return s
 }
 
-func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs []Struct) []Query {
+func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs []Struct) ([]Query, error) {
 	qs := make([]Query, 0, len(r.Queries))
 	for _, query := range r.Queries {
 		if query.Name == "" {
@@ -399,6 +401,9 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 		if query.Cmd == "" {
 			continue
 		}
+		if query.Cmd == metadata.CmdCopyFrom {
+			return nil, errors.New("Support for CopyFrom in Python is not implemented")
+		}
 
 		methodName := MethodName(query.Name)
 
@@ -490,7 +495,7 @@ func buildQueries(r *compiler.Result, settings config.CombinedSettings, structs
 		qs = append(qs, gq)
 	}
 	sort.Slice(qs, func(i, j int) bool { return qs[i].MethodName < qs[j].MethodName })
-	return qs
+	return qs, nil
 }
 
 func importNode(name string) *pyast.Node {
@@ -1052,7 +1057,10 @@ func HashComment(s string) string {
 func Generate(r *compiler.Result, settings config.CombinedSettings) (map[string]string, error) {
 	enums := buildEnums(r, settings)
 	models := buildModels(r, settings)
-	queries := buildQueries(r, settings, models)
+	queries, err := buildQueries(r, settings, models)
+	if err != nil {
+		return nil, err
+	}
 
 	i := &importer{
 		Settings: settings,