/
cmd.go
74 lines (65 loc) · 1.74 KB
/
cmd.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
package cmd
import (
"bytes"
"fmt"
"io/ioutil"
"sort"
"strings"
"h12.io/csv"
)
type MySQLCmd struct {
Value interface{}
TagKey string
ExpandPath []string
Engine string
Replace bool
}
func (cmd MySQLCmd) CreateDB(name string) string {
return fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s DEFAULT CHARACTER SET utf8mb4;\n", name)
}
func (cmd MySQLCmd) CreateTable(fullTableName string) (string, error) {
w := new(bytes.Buffer)
fmt.Fprintf(w, "CREATE TABLE IF NOT EXISTS %s (\n", fullTableName)
fields, err := cmd.Fields()
if err != nil {
return "", err
}
var pks []string
sort.Sort(fields)
for _, field := range fields {
if field.Tag.Get("PK") == "yes" {
pks = append(pks, field.Name)
fmt.Fprintf(w, "\t%s %s,\n", field.Name, field.Tag.Get("TYPE"))
}
}
for _, field := range fields {
if field.Tag.Get("PK") != "yes" {
fmt.Fprintf(w, "\t%s %s,\n", field.Name, field.Tag.Get("TYPE"))
}
}
fmt.Fprintf(w, "\tPRIMARY KEY (%s)\n", strings.Join(pks, ","))
fmt.Fprintf(w, ") ENGINE=%s DEFAULT CHARSET=utf8mb4;\n", cmd.Engine)
return w.String(), nil
}
func (cmd MySQLCmd) LoadDataTemplate() (string, error) {
fields, err := cmd.Fields()
if err != nil {
return "", err
}
replaceOrIgnore := "IGNORE"
if cmd.Replace {
replaceOrIgnore = "REPLACE"
}
return fmt.Sprintf("LOAD DATA LOCAL INFILE 'Reader::%%s' "+
"%s INTO TABLE %%s "+
"CHARACTER SET utf8mb4 "+
`FIELDS OPTIONALLY ENCLOSED BY '"' `+
"(%s);\n", replaceOrIgnore, strings.Join(fields.Names(), ", ")), nil
}
func (cmd MySQLCmd) Fields() (csv.Fields, error) {
enc := csv.NewEncoder(ioutil.Discard).SetTagKey(cmd.TagKey).SetExpandPath(cmd.ExpandPath...)
if err := enc.Encode(cmd.Value); err != nil {
return nil, err
}
return enc.Fields(), nil
}