用Go向MySQL导入.csv文件（转载）

版权声明：转载请注明。 https://blog.csdn.net/puss0/article/details/90411618

今天来更新一个很少碰到，但碰到了又让人十分蛋疼的问题——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);

• 数据库驱动用的是"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)
	}
}

[注：以上代码不能直接食用，除非你装了MySQL]

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

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

• 什么鬼？文件未注册？文件还需要注册？明明在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) //<--根源所在
	}
	...
}

错误就来自这段代码，问题很明显了，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()
}

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)
	}
}

运行成功。需要注意的是，sql语句中的文件名和RegisterLocalFile带入的文件名必须完全一样，而且是绝对路径。

---

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

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

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

type Config struct {
	...
	AllowAllFiles           bool // Allow all files to be used with LOAD DATA LOCAL INFILE
	...
}

看到这里是不是很开心呢？不过不要开心的太早了，虽然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)
			}
		...
		}
	}
}

还记得sql.Open函数吗？他的第二个参数需要一个字符串来指定用户名，密码，表名等等信息。parseDSNParams函数就是在解析这个字符串，并且设置相应的参数。所以我们只需要在原来的dburl最后加上&allowAllFiles=true就可以了。

dburl = "root:******@/test?charset=utf8&parseTime=True&loc=Local&allowAllFiles=true"

• 其他都不用改，再次运行，成功。

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

终点站到了，请您携带好随身物品，下车请走好，欢迎大家再次乘坐。