Implement support for pgx's CopyFrom (#1352)

Author: Jille

docs/guides/development.md

@@ -57,10 +57,11 @@ MYSQL_DATABASE  dinotest
 ## Regenerate expected test output
 
 If you need to update a large number of expected test output in the
-`internal/endtoend/testdata` directory, run the `regenerate.sh` script.
+`internal/endtoend/testdata` directory, run the `regenerate` script.
 
 ```
-make regen
+go build -o ~/go/bin/sqlc-dev ./cmd/sqlc
+go run scripts/regenerate/main.go
 ```
 
 Note that this uses the `sqlc-dev` binary, not `sqlc` so make sure you have an

docs/howto/insert.md

@@ -124,15 +124,41 @@ func (q *Queries) DeleteID(ctx context.Context, id int) (int, error) {
 	return i, err
 }
 
-const deleteAuhtor = `-- name: DeleteAuthor :one
+const deleteAuthor = `-- name: DeleteAuthor :one
 DELETE FROM authors WHERE id = $1
 RETURNING id, bio
 `
 
 func (q *Queries) DeleteAuthor(ctx context.Context, id int) (Author, error) {
-	row := q.db.QueryRowContext(ctx, deleteAuhtor, id)
+	row := q.db.QueryRowContext(ctx, deleteAuthor, id)
 	var i Author
 	err := row.Scan(&i.ID, &i.Bio)
 	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

@@ -3,6 +3,7 @@ package golang
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"fmt"
 	"go/format"
 	"strings"
@@ -11,6 +12,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 +39,7 @@ type tmplCtx struct {
 	EmitInterface             bool
 	EmitEmptySlices           bool
 	EmitMethodsWithDBArgument bool
+	UsesCopyFrom              bool
 }
 
 func (t *tmplCtx) OutputQuery(sourceName string) bool {
@@ -87,6 +90,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,
@@ -95,6 +99,10 @@ func generate(settings config.CombinedSettings, enums []Enum, structs []Struct,
 		Structs:                   structs,
 	}
 
+	if tctx.UsesCopyFrom && tctx.SQLPackage != SQLPackagePGX {
+		return nil, errors.New(":copyfrom is only supported by pgx")
+	}
+
 	output := map[string]string{}
 
 	execute := func(name, templateName string) error {
@@ -135,6 +143,8 @@ func generate(settings config.CombinedSettings, enums []Enum, structs []Struct,
 	if golang.OutputQuerierFileName != "" {
 		querierFileName = golang.OutputQuerierFileName
 	}
+	copyfromFileName := "copyfrom.go"
+	// TODO(Jille): Make this configurable.
 
 	if err := execute(dbFileName, "dbFile"); err != nil {
 		return nil, err
@@ -147,6 +157,11 @@ func generate(settings config.CombinedSettings, enums []Enum, structs []Struct,
 			return nil, err
 		}
 	}
+	if tctx.UsesCopyFrom {
+		if err := execute(copyfromFileName, "copyfromFile"); err != nil {
+			return nil, err
+		}
+	}
 
 	files := map[string]struct{}{}
 	for _, gq := range queries {
@@ -160,3 +175,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/imports.go

@@ -89,6 +89,7 @@ func (i *importer) Imports(filename string) [][]ImportSpec {
 	if i.Settings.Go.OutputQuerierFileName != "" {
 		querierFileName = i.Settings.Go.OutputQuerierFileName
 	}
+	copyfromFileName := "copyfrom.go"
 
 	switch filename {
 	case dbFileName:
@@ -97,6 +98,8 @@ func (i *importer) Imports(filename string) [][]ImportSpec {
 		return mergeImports(i.modelImports())
 	case querierFileName:
 		return mergeImports(i.interfaceImports())
+	case copyfromFileName:
+		return mergeImports(i.interfaceImports())
 	default:
 		return mergeImports(i.queryImports(filename))
 	}
@@ -279,9 +282,13 @@ func sortedImports(std map[string]struct{}, pkg map[ImportSpec]struct{}) fileImp
 
 func (i *importer) queryImports(filename string) fileImports {
 	var gq []Query
+	anyNonCopyFrom := false
 	for _, query := range i.Queries {
 		if query.SourceName == filename {
 			gq = append(gq, query)
+			if query.Cmd != metadata.CmdCopyFrom {
+				anyNonCopyFrom = true
+			}
 		}
 	}
 
@@ -349,7 +356,9 @@ func (i *importer) queryImports(filename string) fileImports {
 		return false
 	}
 
-	std["context"] = struct{}{}
+	if anyNonCopyFrom {
+		std["context"] = struct{}{}
+	}
 
 	sqlpkg := SQLPackageFromString(i.Settings.Go.SQLPackage)
 	if sliceScan() && sqlpkg != SQLPackagePGX {

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)
 
@@ -295,9 +296,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/copyfromCopy.tmpl

@@ -0,0 +1,51 @@
+{{define "copyfromCodePgx"}}
+{{range .GoQueries}}
+{{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}}

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}}
     }