@@ -14,15 +14,27 @@ import (
14
14
"os"
15
15
"strings"
16
16
"sync"
17
+ "time"
17
18
)
18
19
19
20
var (
20
21
fileRegister map [string ]bool
21
22
fileRegisterLock sync.RWMutex
22
- readerRegister map [string ]func () io.Reader
23
+ readerRegister map [string ]func (EncoderInfo ) io.Reader
23
24
readerRegisterLock sync.RWMutex
24
25
)
25
26
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
+
26
38
// RegisterLocalFile adds the given file to the file allowlist,
27
39
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
28
40
// Alternatively you can allow the use of all local files with
@@ -66,10 +78,18 @@ func DeregisterLocalFile(filePath string) {
66
78
// if err != nil {
67
79
// ...
68
80
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 ) {
69
89
readerRegisterLock .Lock ()
70
90
// lazy map init
71
91
if readerRegister == nil {
72
- readerRegister = make (map [string ]func () io.Reader )
92
+ readerRegister = make (map [string ]func (EncoderInfo ) io.Reader )
73
93
}
74
94
75
95
readerRegister [name ] = handler
@@ -110,7 +130,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
110
130
readerRegisterLock .RUnlock ()
111
131
112
132
if inMap {
113
- rdr = handler ()
133
+ rdr = handler (mc . cfg . encoderInfo () )
114
134
if rdr != nil {
115
135
if cl , ok := rdr .(io.Closer ); ok {
116
136
defer deferredClose (& err , cl )
0 commit comments