Skip to content

Commit 759860f

Browse files
JilleJille
committed
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
1 parent a8a9ff9 commit 759860f

File tree

14 files changed

+267
-6
lines changed

14 files changed

+267
-6
lines changed

‎internal/codegen/golang/gen.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,8 +138,8 @@ func generate(req *plugin.CodeGenRequest, enums []Enum, structs []Struct, querie
138138
SqlcVersion: req.SqlcVersion,
139139
}
140140

141-
if tctx.UsesCopyFrom && !tctx.SQLDriver.IsPGX() {
142-
return nil, errors.New(":copyfrom is only supported by pgx")
141+
if tctx.UsesCopyFrom && !tctx.SQLDriver.IsPGX() && golang.SqlDriver != SQLDriverGoSQLDriverMySQL {
142+
return nil, errors.New(":copyfrom is only supported by pgx and github.com/go-sql-driver/mysql")
143143
}
144144

145145
if tctx.UsesBatch && !tctx.SQLDriver.IsPGX() {

‎internal/codegen/golang/imports.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -414,6 +414,13 @@ func (i *importer) copyfromImports() fileImports {
414414
})
415415

416416
std["context"] = struct{}{}
417+
if i.Settings.Go.SqlDriver == SQLDriverGoSQLDriverMySQL {
418+
std["io"] = struct{}{}
419+
std["fmt"] = struct{}{}
420+
std["sync/atomic"] = struct{}{}
421+
pkg[ImportSpec{Path: "github.com/go-sql-driver/mysql"}] = struct{}{}
422+
pkg[ImportSpec{Path: "github.com/hexon/mysqltsv"}] = struct{}{}
423+
}
417424

418425
return sortedImports(std, pkg)
419426
}

‎internal/codegen/golang/query.go

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,18 @@ func (v QueryValue) Params() string {
131131
return "\n" + strings.Join(out, ",\n")
132132
}
133133

134-
func (v QueryValue) ColumnNames() string {
134+
func (v QueryValue) ColumnNames() []string {
135+
if v.Struct == nil {
136+
return []string{v.DBName}
137+
}
138+
names := make([]string, len(v.Struct.Fields))
139+
for i, f := range v.Struct.Fields {
140+
names[i] = f.DBName
141+
}
142+
return names
143+
}
144+
145+
func (v QueryValue) ColumnNamesAsGoSlice() string {
135146
if v.Struct == nil {
136147
return fmt.Sprintf("[]string{%q}", v.DBName)
137148
}
@@ -189,6 +200,19 @@ func (v QueryValue) Scan() string {
189200
return "\n" + strings.Join(out, ",\n")
190201
}
191202

203+
func (v QueryValue) Fields() []Field {
204+
if v.Struct != nil {
205+
return v.Struct.Fields
206+
}
207+
return []Field{
208+
{
209+
Name: v.Name,
210+
DBName: v.DBName,
211+
Type: v.Typ,
212+
},
213+
}
214+
}
215+
192216
// A struct used to generate methods and fields on the Queries struct
193217
type Query struct {
194218
Cmd string
@@ -210,7 +234,7 @@ func (q Query) hasRetType() bool {
210234
return scanned && !q.Ret.isEmpty()
211235
}
212236

213-
func (q Query) TableIdentifier() string {
237+
func (q Query) TableIdentifierAsGoSlice() string {
214238
escapedNames := make([]string, 0, 3)
215239
for _, p := range []string{q.Table.Catalog, q.Table.Schema, q.Table.Name} {
216240
if p != "" {
@@ -219,3 +243,13 @@ func (q Query) TableIdentifier() string {
219243
}
220244
return "[]string{" + strings.Join(escapedNames, ", ") + "}"
221245
}
246+
247+
func (q Query) TableIdentifierForMySQL() string {
248+
escapedNames := make([]string, 0, 3)
249+
for _, p := range []string{q.Table.Catalog, q.Table.Schema, q.Table.Name} {
250+
if p != "" {
251+
escapedNames = append(escapedNames, fmt.Sprintf("`%s`", p))
252+
}
253+
}
254+
return strings.Join(escapedNames, ".")
255+
}

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
{{define "copyfromCodeGoSqlDriver"}}
2+
{{range .GoQueries}}
3+
{{if eq .Cmd ":copyfrom" }}
4+
var readerHandlerSequenceFor{{.MethodName}} uint32 = 1
5+
6+
func convertRowsFor{{.MethodName}}(w *io.PipeWriter, {{.Arg.SlicePair}}) {
7+
e := mysqltsv.NewEncoder(w, {{ len .Arg.Fields }}, nil)
8+
for _, row := range {{.Arg.Name}} {
9+
{{- with $arg := .Arg }}
10+
{{- range $arg.Fields}}
11+
{{- if eq .Type "string"}}
12+
e.AppendString({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
13+
{{- else if eq .Type "[]byte"}}
14+
e.AppendBytes({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
15+
{{- else}}
16+
e.AppendValue({{if eq (len $arg.Fields) 1}}row{{else}}row.{{.Name}}{{end}})
17+
{{- end}}
18+
{{- end}}
19+
{{- end}}
20+
}
21+
w.CloseWithError(e.Close())
22+
}
23+
24+
{{range .Comments}}//{{.}}
25+
{{end -}}
26+
// {{.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.
27+
// Use this in a transaction and use SHOW WARNINGS to check for any problems and roll back if you want to.
28+
// 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
29+
func (q *Queries) {{.MethodName}}(ctx context.Context{{if $.EmitMethodsWithDBArgument}}, db DBTX{{end}}, {{.Arg.SlicePair}}) (int64, error) {
30+
pr, pw := io.Pipe()
31+
defer pr.Close()
32+
rh := fmt.Sprintf("{{.MethodName}}_%d", atomic.AddUint32(&readerHandlerSequenceFor{{.MethodName}}, 1))
33+
mysql.RegisterReaderHandler(rh, func() io.Reader { return pr })
34+
defer mysql.DeregisterReaderHandler(rh)
35+
go convertRowsFor{{.MethodName}}(pw, {{.Arg.Name}})
36+
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))
37+
if err != nil {
38+
return 0, err
39+
}
40+
return result.RowsAffected()
41+
}
42+
43+
{{end}}
44+
{{end}}
45+
{{end}}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,10 @@ func (r iteratorFor{{.MethodName}}) Err() error {
3939
{{end -}}
4040
{{- if $.EmitMethodsWithDBArgument -}}
4141
func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.SlicePair}}) (int64, error) {
42-
return db.CopyFrom(ctx, {{.TableIdentifier}}, {{.Arg.ColumnNames}}, &iteratorFor{{.MethodName}}{rows: {{.Arg.Name}}})
42+
return db.CopyFrom(ctx, {{.TableIdentifierAsGoSlice}}, {{.Arg.ColumnNamesAsGoSlice}}, &iteratorFor{{.MethodName}}{rows: {{.Arg.Name}}})
4343
{{- else -}}
4444
func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.SlicePair}}) (int64, error) {
45-
return q.db.CopyFrom(ctx, {{.TableIdentifier}}, {{.Arg.ColumnNames}}, &iteratorFor{{.MethodName}}{rows: {{.Arg.Name}}})
45+
return q.db.CopyFrom(ctx, {{.TableIdentifierAsGoSlice}}, {{.Arg.ColumnNamesAsGoSlice}}, &iteratorFor{{.MethodName}}{rows: {{.Arg.Name}}})
4646
{{- end}}
4747
}
4848

‎internal/codegen/golang/templates/template.tmpl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,8 @@ import (
186186
{{define "copyfromCode"}}
187187
{{if .SQLDriver.IsPGX }}
188188
{{- template "copyfromCodePgx" .}}
189+
{{else}}
190+
{{- template "copyfromCodeGoSqlDriver" .}}
189191
{{end}}
190192
{{end}}
191193

‎internal/endtoend/testdata/copyfrom/mysql/go/copyfrom.go

Lines changed: 74 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/endtoend/testdata/copyfrom/mysql/go/db.go

Lines changed: 31 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/endtoend/testdata/copyfrom/mysql/go/models.go

Lines changed: 16 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 25 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
CREATE TABLE foo (a text, b integer, c DATETIME, d DATE);
2+
3+
-- name: InsertValues :copyfrom
4+
INSERT INTO foo (a, b, c, d) VALUES (?, ?, ?, ?);
5+
6+
-- name: InsertSingleValue :copyfrom
7+
INSERT INTO foo (a) VALUES (?);

‎internal/endtoend/testdata/copyfrom/mysql/sqlc.json

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
{
2+
"version": "1",
3+
"packages": [
4+
{
5+
"path": "go",
6+
"sql_package": "database/sql",
7+
"sql_driver": "github.com/go-sql-driver/mysql",
8+
"engine": "mysql",
9+
"name": "querytest",
10+
"schema": "query.sql",
11+
"queries": "query.sql"
12+
}
13+
]
14+
}

‎internal/endtoend/testdata/go.mod

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@ module github.com/kyleconroy/sqlc/endtoend
33
go 1.18
44

55
require (
6+
github.com/go-sql-driver/mysql v1.7.0
67
github.com/gofrs/uuid v4.0.0+incompatible
78
github.com/google/uuid v1.3.0
9+
github.com/hexon/mysqltsv v0.1.0
810
github.com/jackc/pgconn v1.5.1-0.20200601181101-fa742c524853
911
github.com/jackc/pgtype v1.6.2
1012
github.com/jackc/pgx/v4 v4.6.1-0.20200606145419-4e5062306904

‎internal/endtoend/testdata/go.sum

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,17 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
99
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
1010
github.com/friendsofgo/errors v0.9.2 h1:X6NYxef4efCBdwI7BgS820zFaN7Cphrmb+Pljdzjtgk=
1111
github.com/friendsofgo/errors v0.9.2/go.mod h1:yCvFW5AkDIL9qn7suHVLiI/gH228n7PC4Pn44IGoTOI=
12+
github.com/go-sql-driver/mysql v1.7.0 h1:ueSltNNllEqE3qcWBTD0iQd3IpL/6U+mJxLkazJ7YPc=
13+
github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
1214
github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
1315
github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
1416
github.com/gofrs/uuid v4.0.0+incompatible h1:1SD/1F5pU8p29ybwgQSwpQk+mwdRrXCYuPhW6m+TnJw=
1517
github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
1618
github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
1719
github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
1820
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
21+
github.com/hexon/mysqltsv v0.1.0 h1:48wYQlsPH8ZEkKAVCdsOYzMYAlEoevw8ZWD8rqYPdlg=
22+
github.com/hexon/mysqltsv v0.1.0/go.mod h1:p3vPBkpxebjHWF1bWKYNcXx5pFu+yAG89QZQEKSvVrY=
1923
github.com/jackc/chunkreader v1.0.0/go.mod h1:RT6O25fNZIuasFJRyZ4R/Y2BbhasbmZXF9QQ7T3kePo=
2024
github.com/jackc/chunkreader/v2 v2.0.0/go.mod h1:odVSm741yZoC3dpHEUXIqA9tQRhFrgOHwnPIn9lDKlk=
2125
github.com/jackc/chunkreader/v2 v2.0.1 h1:i+RDz65UE+mmpjTfyz0MoVTnzeYxroil2G82ki7MGG8=

0 commit comments

Comments
 (0)