feat(mysql): :copyfrom support via LOAD DATA INFILE

internal/codegen/golang/driver.go

@@ -1,6 +1,6 @@
 package golang
 
-type SQLDriver int
+type SQLDriver string
 
 const (
 	SQLPackagePGXV4    string = "pgx/v4"
@@ -9,13 +9,12 @@ const (
 )
 
 const (
-	SQLDriverPGXV4 SQLDriver = iota
-	SQLDriverPGXV5
-	SQLDriverLibPQ
+	SQLDriverPGXV4            SQLDriver = "github.com/jackc/pgx/v4"
+	SQLDriverPGXV5                      = "github.com/jackc/pgx/v5"
+	SQLDriverLibPQ                      = "github.com/lib/pq"
+	SQLDriverGoSQLDriverMySQL           = "github.com/go-sql-driver/mysql"
 )
 
-const SQLDriverGoSQLDriverMySQL = "github.com/go-sql-driver/mysql"
-
 func parseDriver(sqlPackage string) SQLDriver {
 	switch sqlPackage {
 	case SQLPackagePGXV4:
@@ -31,6 +30,10 @@ func (d SQLDriver) IsPGX() bool {
 	return d == SQLDriverPGXV4 || d == SQLDriverPGXV5
 }
 
+func (d SQLDriver) IsGoSQLDriverMySQL() bool {
+	return d == SQLDriverGoSQLDriverMySQL
+}
+
 func (d SQLDriver) Package() string {
 	switch d {
 	case SQLDriverPGXV4:

internal/codegen/golang/gen.go

@@ -145,8 +145,15 @@ func generate(req *plugin.CodeGenRequest, enums []Enum, structs []Struct, querie
 		SqlcVersion:               req.SqlcVersion,
 	}
 
-	if tctx.UsesCopyFrom && !tctx.SQLDriver.IsPGX() {
-		return nil, errors.New(":copyfrom is only supported by pgx")
+	if tctx.UsesCopyFrom && !tctx.SQLDriver.IsPGX() && golang.SqlDriver != SQLDriverGoSQLDriverMySQL {
+		return nil, errors.New(":copyfrom is only supported by pgx and github.com/go-sql-driver/mysql")
+	}
+
+	if tctx.UsesCopyFrom && golang.SqlDriver == SQLDriverGoSQLDriverMySQL {
+		if err := checkNoTimesForMySQLCopyFrom(queries); err != nil {
+			return nil, err
+		}
+		tctx.SQLDriver = SQLDriverGoSQLDriverMySQL
 	}
 
 	if tctx.UsesBatch && !tctx.SQLDriver.IsPGX() {
@@ -294,6 +301,17 @@ func usesBatch(queries []Query) bool {
 	return false
 }
 
+func checkNoTimesForMySQLCopyFrom(queries []Query) error {
+	for _, q := range queries {
+		for _, f := range q.Arg.Fields() {
+			if f.Type == "time.Time" {
+				return fmt.Errorf("values with a timezone are not yet supported")
+			}
+		}
+	}
+	return nil
+}
+
 func filterUnusedStructs(enums []Enum, structs []Struct, queries []Query) ([]Enum, []Struct) {
 	keepTypes := make(map[string]struct{})
 

internal/codegen/golang/imports.go

@@ -407,6 +407,13 @@ func (i *importer) copyfromImports() fileImports {
 	})
 
 	std["context"] = struct{}{}
+	if i.Settings.Go.SqlDriver == SQLDriverGoSQLDriverMySQL {
+		std["io"] = struct{}{}
+		std["fmt"] = struct{}{}
+		std["sync/atomic"] = struct{}{}
+		pkg[ImportSpec{Path: "github.com/go-sql-driver/mysql"}] = struct{}{}
+		pkg[ImportSpec{Path: "github.com/hexon/mysqltsv"}] = struct{}{}
+	}
 
 	return sortedImports(std, pkg)
 }

internal/codegen/golang/query.go

@@ -129,7 +129,18 @@ func (v QueryValue) Params() string {
 	return "\n" + strings.Join(out, ",\n")
 }
 
-func (v QueryValue) ColumnNames() string {
+func (v QueryValue) ColumnNames() []string {
+	if v.Struct == nil {
+		return []string{v.DBName}
+	}
+	names := make([]string, len(v.Struct.Fields))
+	for i, f := range v.Struct.Fields {
+		names[i] = f.DBName
+	}
+	return names
+}
+
+func (v QueryValue) ColumnNamesAsGoSlice() string {
 	if v.Struct == nil {
 		return fmt.Sprintf("[]string{%q}", v.DBName)
 	}
@@ -187,6 +198,19 @@ func (v QueryValue) Scan() string {
 	return "\n" + strings.Join(out, ",\n")
 }
 
+func (v QueryValue) Fields() []Field {
+	if v.Struct != nil {
+		return v.Struct.Fields
+	}
+	return []Field{
+		{
+			Name:   v.Name,
+			DBName: v.DBName,
+			Type:   v.Typ,
+		},
+	}
+}
+
 func (v QueryValue) VariableForField(f Field) string {
 	if !v.IsStruct() {
 		return v.Name
@@ -218,7 +242,7 @@ func (q Query) hasRetType() bool {
 	return scanned && !q.Ret.isEmpty()
 }
 
-func (q Query) TableIdentifier() string {
+func (q Query) TableIdentifierAsGoSlice() string {
 	escapedNames := make([]string, 0, 3)
 	for _, p := range []string{q.Table.Catalog, q.Table.Schema, q.Table.Name} {
 		if p != "" {
@@ -227,3 +251,13 @@ func (q Query) TableIdentifier() string {
 	}
 	return "[]string{" + strings.Join(escapedNames, ", ") + "}"
 }
+
+func (q Query) TableIdentifierForMySQL() 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("`%s`", p))
+		}
+	}
+	return strings.Join(escapedNames, ".")
+}

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

@@ -0,0 +1,52 @@
+{{define "copyfromCodeGoSqlDriver"}}
+{{range .GoQueries}}
+{{if eq .Cmd ":copyfrom" }}
+var readerHandlerSequenceFor{{.MethodName}} uint32 = 1
+
+func convertRowsFor{{.MethodName}}(w *io.PipeWriter, {{.Arg.SlicePair}}) {
+	e := mysqltsv.NewEncoder(w, {{ len .Arg.Fields }}, nil)
+	for _, row := range {{.Arg.Name}} {
+{{- with $arg := .Arg }}
+{{- range $arg.Fields}}
+{{- if eq .Type "string"}}
+	e.AppendString({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
+{{- else if eq .Type "[]byte"}}
+	e.AppendBytes({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
+{{- else}}
+	e.AppendValue({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
+{{- end}}
+{{- end}}
+{{- end}}
+	}
+	w.CloseWithError(e.Close())
+}
+
+{{range .Comments}}//{{.}}
+{{end -}}
+// {{.MethodName}} uses MySQL's LOAD DATA LOCAL INFILE and is not atomic.
+//
+// Errors and duplicate keys are treated as warnings and insertion will
+// continue, even without an error for some cases.  Use this in a transaction
+// and use SHOW WARNINGS to check for any problems and roll back if you want to.
+//
+// Check the documentation for more information:
+// https://dev.mysql.com/doc/refman/8.0/en/load-data.html#load-data-error-handling
+func (q *Queries) {{.MethodName}}(ctx context.Context{{if $.EmitMethodsWithDBArgument}}, db DBTX{{end}}, {{.Arg.SlicePair}}) (int64, error) {
+	pr, pw := io.Pipe()
+	defer pr.Close()
+	rh := fmt.Sprintf("{{.MethodName}}_%d", atomic.AddUint32(&readerHandlerSequenceFor{{.MethodName}}, 1))
+	mysql.RegisterReaderHandler(rh, func() io.Reader { return pr })
+	defer mysql.DeregisterReaderHandler(rh)
+	go convertRowsFor{{.MethodName}}(pw, {{.Arg.Name}})
+	// The string interpolation is necessary because LOAD DATA INFILE requires
+	// the file name to be given as a literal string.
+	result, err := {{if (not $.EmitMethodsWithDBArgument)}}q.{{end}}db.ExecContext(ctx, fmt.Sprintf("LOAD DATA LOCAL INFILE '%s' INTO TABLE {{.TableIdentifierForMySQL}} %s ({{range $index, $name := .Arg.ColumnNames}}{{if gt $index 0}}, {{end}}{{$name}}{{end}})", "Reader::" + rh, mysqltsv.Escaping))
+	if err != nil {
+		return 0, err
+	}
+	return result.RowsAffected()
+}
+
+{{end}}
+{{end}}
+{{end}}

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

@@ -39,10 +39,10 @@ func (r iteratorFor{{.MethodName}}) Err() error {
 {{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}}})
+	return db.CopyFrom(ctx, {{.TableIdentifierAsGoSlice}}, {{.Arg.ColumnNamesAsGoSlice}}, &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}}})
+	return q.db.CopyFrom(ctx, {{.TableIdentifierAsGoSlice}}, {{.Arg.ColumnNamesAsGoSlice}}, &iteratorFor{{.MethodName}}{rows: {{.Arg.Name}}})
 {{- end}}
 }
 

internal/codegen/golang/templates/template.tmpl

@@ -186,6 +186,8 @@ import (
 {{define "copyfromCode"}}
 {{if .SQLDriver.IsPGX }}
     {{- template "copyfromCodePgx" .}}
+{{else if .SQLDriver.IsGoSQLDriverMySQL }}
+    {{- template "copyfromCodeGoSqlDriver" .}}
 {{end}}
 {{end}}
 

internal/endtoend/testdata/copyfrom/mysql/query.sql

@@ -0,0 +1,7 @@
+CREATE TABLE foo (a text, b integer, c DATETIME, d DATE);
+
+-- name: InsertValues :copyfrom
+INSERT INTO foo (a, b, c, d) VALUES (?, ?, ?, ?);
+
+-- name: InsertSingleValue :copyfrom
+INSERT INTO foo (a) VALUES (?);