Skip to content

Commit 6f7eca0

Browse files
committed
Introduce RegisterReaderHandlerWithEncoderInfo
It is the same as RegisterReaderHandler, but receives some configuration settings which are needed to correctly encode the TSV. (For now only the Location). issue #1416
1 parent 8503110 commit 6f7eca0

File tree

2 files changed

+24
-3
lines changed

2 files changed

+24
-3
lines changed

AUTHORS

+1
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ Jeffrey Charles <jeffreycharles at gmail.com>
5252
Jerome Meyer <jxmeyer at gmail.com>
5353
Jiajia Zhong <zhong2plus at gmail.com>
5454
Jian Zhen <zhenjl at gmail.com>
55+
Jille Timmermans <jille at quis.cx>
5556
Joshua Prunier <joshua.prunier at gmail.com>
5657
Julien Lefevre <julien.lefevr at gmail.com>
5758
Julien Schmidt <go-sql-driver at julienschmidt.com>

infile.go

+23-3
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,27 @@ import (
1414
"os"
1515
"strings"
1616
"sync"
17+
"time"
1718
)
1819

1920
var (
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
// ...
6880
func 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

Comments
 (0)