woodpecker/vendor/github.com/go-sql-driver/mysql/infile.go
2015-09-29 17:34:44 -07:00

152 lines
3.6 KiB
Go

// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.
package mysql
import (
"database/sql/driver"
"fmt"
"io"
"os"
"strings"
)
var (
fileRegister map[string]bool
readerRegister map[string]func() io.Reader
)
func init() {
fileRegister = make(map[string]bool)
readerRegister = make(map[string]func() io.Reader)
}
// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
// filePath := "/home/gopher/data.csv"
// mysql.RegisterLocalFile(filePath)
// err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
// if err != nil {
// ...
//
func RegisterLocalFile(filePath string) {
fileRegister[strings.Trim(filePath, `"`)] = true
}
// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
delete(fileRegister, strings.Trim(filePath, `"`))
}
// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
// mysql.RegisterReaderHandler("data", func() io.Reader {
// var csvReader io.Reader // Some Reader that returns CSV data
// ... // Open Reader here
// return csvReader
// })
// err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
// if err != nil {
// ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
readerRegister[name] = handler
}
// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
delete(readerRegister, name)
}
func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
var rdr io.Reader
data := make([]byte, 4+mc.maxWriteSize)
if strings.HasPrefix(name, "Reader::") { // io.Reader
name = name[8:]
handler, inMap := readerRegister[name]
if handler != nil {
rdr = handler()
}
if rdr == nil {
if !inMap {
err = fmt.Errorf("Reader '%s' is not registered", name)
} else {
err = fmt.Errorf("Reader '%s' is <nil>", name)
}
}
} else { // File
name = strings.Trim(name, `"`)
if mc.cfg.allowAllFiles || fileRegister[name] {
rdr, err = os.Open(name)
} else {
err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
}
}
if rdc, ok := rdr.(io.ReadCloser); ok {
defer func() {
if err == nil {
err = rdc.Close()
} else {
rdc.Close()
}
}()
}
// send content packets
var ioErr error
if err == nil {
var n int
for err == nil && ioErr == nil {
n, err = rdr.Read(data[4:])
if n > 0 {
data[0] = byte(n)
data[1] = byte(n >> 8)
data[2] = byte(n >> 16)
data[3] = mc.sequence
ioErr = mc.writePacket(data[:4+n])
}
}
if err == io.EOF {
err = nil
}
if ioErr != nil {
errLog.Print(ioErr.Error())
return driver.ErrBadConn
}
}
// send empty packet (termination)
ioErr = mc.writePacket([]byte{
0x00,
0x00,
0x00,
mc.sequence,
})
if ioErr != nil {
errLog.Print(ioErr.Error())
return driver.ErrBadConn
}
// read OK packet
if err == nil {
return mc.readResultOK()
} else {
mc.readPacket()
}
return err
}