MySQL LOAD DATA INFILE: First version

This enables the :copyfrom query annotation for people using
go-sql-driver/mysql that transforms it into a LOAD DATA LOCAL INFILE.

issue #2179

Author: Jille

internal/codegen/golang/gen.go

@@ -138,8 +138,8 @@ 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.UsesBatch && !tctx.SQLDriver.IsPGX() {

internal/codegen/golang/imports.go

@@ -414,6 +414,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{}{}
+		std["github.com/go-sql-driver/mysql"] = struct{}{}
+		std["github.com/hexon/mysqltsv"] = struct{}{}
+	}
 
 	return sortedImports(std, pkg)
 }

internal/codegen/golang/query.go

@@ -131,7 +131,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)
 	}
@@ -189,6 +200,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,
+		},
+	}
+}
+
 // A struct used to generate methods and fields on the Queries struct
 type Query struct {
 	Cmd          string
@@ -210,7 +234,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 != "" {
@@ -219,3 +243,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,45 @@
+{{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.
+// This is a MySQL limitation, not sqlc. 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}})
+	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}}
+    {{- template "copyfromCodeGoSqlDriver" .}}
 {{end}}
 {{end}}