@@ -14,15 +14,27 @@ import (
1414 "os"
1515 "strings"
1616 "sync"
17+ "time"
1718)
1819
1920var (
2021 fileRegister map [string ]bool
2122 fileRegisterLock sync.RWMutex
22- readerRegister map [string ]func () io.Reader
23+ readerRegister map [string ]func (EncoderInfo ) io.Reader
2324 readerRegisterLock sync.RWMutex
2425)
2526
27+ // EncoderInfo contains the settings needed to encode a TSV file for LOAD DATA INFILE.
28+ type EncoderInfo struct {
29+ Location * time.Location
30+ }
31+
32+ func (c * Config ) encoderInfo () EncoderInfo {
33+ return EncoderInfo {
34+ Location : c .Loc ,
35+ }
36+ }
37+
2638// RegisterLocalFile adds the given file to the file allowlist,
2739// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
2840// Alternatively you can allow the use of all local files with
@@ -66,10 +78,18 @@ func DeregisterLocalFile(filePath string) {
6678// if err != nil {
6779// ...
6880func RegisterReaderHandler (name string , handler func () io.Reader ) {
81+ RegisterReaderHandlerWithEncoderInfo (name , func (EncoderInfo ) io.Reader {
82+ return handler ()
83+ })
84+ }
85+
86+ // RegisterReaderHandlerWithEncoderInfo is like RegisterReaderHandler but the
87+ // callback receives the information needed to correctly encode a TSV file.
88+ func RegisterReaderHandlerWithEncoderInfo (name string , handler func (EncoderInfo ) io.Reader ) {
6989 readerRegisterLock .Lock ()
7090 // lazy map init
7191 if readerRegister == nil {
72- readerRegister = make (map [string ]func () io.Reader )
92+ readerRegister = make (map [string ]func (EncoderInfo ) io.Reader )
7393 }
7494
7595 readerRegister [name ] = handler
@@ -110,7 +130,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
110130 readerRegisterLock .RUnlock ()
111131
112132 if inMap {
113- rdr = handler ()
133+ rdr = handler (mc . cfg . encoderInfo () )
114134 if rdr != nil {
115135 if cl , ok := rdr .(io.Closer ); ok {
116136 defer deferredClose (& err , cl )
0 commit comments