用Go向MySQL导入.csv文件

今天来更新一个很少碰到，但碰到了又让人十分蛋疼的问题——Go语言中执行MySQL的load data local infile语句报local file 'xxx' is not registered错误该如何解决。

上车请刷卡，没卡的乘客请投币，上车的乘客请往车厢中部走，汽车起步，请坐稳扶好。

情景在现：我要在Go语言中执行sql语句，往MySQL中导入一个.csv文件，sql语句如下：

load data LOCAL infile './user.csv' into table users
CHARACTER SET utf8 fields terminated by ',' lines terminated by '\r\n' ignore 1 lines(age,name);
1
2
数据库驱动用的是"github.com/go-sql-driver/mysql"，go语言代码如下：

package main

import (
"database/sql"
"fmt"
_ "github.com/go-sql-driver/mysql"
)

var (
db *sql.DB
dburl = "root:******@/test?charset=utf8&parseTime=True&loc=Local"
sql = `load data LOCAL infile 'F://user.csv' into table users CHARACTER SET utf8 fields terminated by ',' lines terminated by '\r\n' ignore 1 lines(age,name);`
)

func init() {
var err error
db, err = sql.Open("mysql", dburl)
if err != nil {
panic(err)
}
}

func main() {
_, err := db.Exec(sql)
if err != nil {
fmt.Println(err)
}
}
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
[注：以上代码不能直接食用，除非你装了MySQL]

一切看起来都很美好，然而运行的时候。。。会报错下面的错误。

local file 'F://user.csv' is not registered
1

[龙叔对不住啦：手动捂脸]

什么鬼？文件未注册？文件还需要注册？明明在cmd里面执行这句sql就没有问题，偏偏在go里面就不行，而且怎么弄都不行，简直想砸电。。“电脑太贵了”。。简直想砸键盘，然而。。。。砸完键盘之后，问题还是要解决的。

经过一番挖地三尺的搜索，终于找到了问题的根源所在。

在MySQL驱动的infile.go文件中，有一个叫handleInFileRequest的函数，其中有这样一段代码：

package mysql

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
...
fr := fileRegister[name] // 这里的 name 就是文件名：F://user.csv
if mc.cfg.AllowAllFiles || fr {
var file *os.File
var fi os.FileInfo

if file, err = os.Open(name); err == nil {
defer deferredClose(&err, file)

// get file size
if fi, err = file.Stat(); err == nil {
rdr = file
if fileSize := int(fi.Size()); fileSize < packetSize {
packetSize = fileSize
}
}
}
} else {
err = fmt.Errorf("local file '%s' is not registered", name) //<--根源所在
}
...
}
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
错误就来自这段代码，问题很明显了，mc.cfg.AllowAllFiles = false并且fr为空。因此我们可以从两个方面来解决这个问题。

方案一、解决fr为空的问题。fileRegister实际上是一个map，在infile.go中定义，并且有一个函数作为接口：

package mysql

var (
fileRegister map[string]bool
)

func RegisterLocalFile(filePath string) {
fileRegisterLock.Lock()
// lazy map init
if fileRegister == nil {
fileRegister = make(map[string]bool)
}

fileRegister[strings.Trim(filePath, `"`)] = true
fileRegisterLock.Unlock()
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
RegisterLocalFile是一个导出函数，可以在我们的main函数中调用，但前提是不能以_的方式导入驱动包github.com/go-sql-driver/mysql。将代码修改如下：

package main

import (
"database/sql"
"fmt"
"github.com/go-sql-driver/mysql"
)

var (
db *sql.DB
dburl = "root:******@/test?charset=utf8&parseTime=True&loc=Local"
sql = `load data LOCAL infile 'F://user.csv' into table users CHARACTER SET utf8 fields terminated by ',' lines terminated by '\r\n' ignore 1 lines(age,name);`
)

func init() {
var err error
db, err = sql.Open("mysql", dburl)
if err != nil {
panic(err)
}
}

func main() {
mysql.RegisterLocalFile("F://user.csv")
_, err := db.Exec(sql)
if err != nil {
fmt.Println(err)
}
}
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
运行成功。需要注意的是，sql语句中的文件名和RegisterLocalFile带入的文件名必须完全一样，而且是绝对路径。

方案二、我们还可以通过修改mc.cfg.AllowAllFiles的值来解决这个问题。handleInFileRequest是mysqlConn的函数，我们需要看看mysqlConn的定义。

type mysqlConn struct {
...
cfg *Config
...
}

1
2
3
4
5
6
去掉那些无关紧要的字段，我们看到cfg字段是*Config类型，再看看Config结构体的定义。

type Config struct {
...
AllowAllFiles bool // Allow all files to be used with LOAD DATA LOCAL INFILE
...
}
1
2
3
4
5
看到这里是不是很开心呢？不过不要开心的太早了，虽然AllowAllFiles是导出的，但是cfg是未导出字段，并且mysqlConn并没有提供有关cfg的接口，玩个蛋。。。

就在我以为此路不通而心灰意冷自怨自艾濒临放弃怨天尤人的时候，天无绝人之路。

在dsn.go中，我发现了这样一个函数：

package mysql

func parseDSNParams(cfg *Config, params string) (err error) {
for _, v := range strings.Split(params, "&") {
param := strings.SplitN(v, "=", 2)
if len(param) != 2 {
continue
}

// cfg params
switch value := param[1]; param[0] {
// Disable INFILE whitelist / enable all files
case "allowAllFiles":
var isBool bool
cfg.AllowAllFiles, isBool = readBool(value)
if !isBool {
return errors.New("invalid bool value: " + value)
}
...
}
}
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
还记得sql.Open函数吗？他的第二个参数需要一个字符串来指定用户名，密码，表名等等信息。parseDSNParams函数就是在解析这个字符串，并且设置相应的参数。所以我们只需要在原来的dburl最后加上&allowAllFiles=true就可以了。

dburl = "root:******@/test?charset=utf8&parseTime=True&loc=Local&allowAllFiles=true"
1
其他都不用改，再次运行，成功。

本方案只针对MySQL使用github.com/go-sql-driver/mysql作为数据库驱动的情况，至于用的是database/sql还是gorm还是sqlx都没关系。如果数据库驱动不是go-sql-driver/mysql，遇到类似的问题也可参考以上的思路。
---------------------