infile.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181
  1. // Go MySQL Driver - A MySQL-Driver for Go's database/sql package
  2. //
  3. // Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
  4. //
  5. // This Source Code Form is subject to the terms of the Mozilla Public
  6. // License, v. 2.0. If a copy of the MPL was not distributed with this file,
  7. // You can obtain one at http://mozilla.org/MPL/2.0/.
  8. package mysql
  9. import (
  10. "fmt"
  11. "io"
  12. "os"
  13. "strings"
  14. "sync"
  15. )
  16. var (
  17. fileRegister map[string]bool
  18. fileRegisterLock sync.RWMutex
  19. readerRegister map[string]func() io.Reader
  20. readerRegisterLock sync.RWMutex
  21. )
  22. // RegisterLocalFile adds the given file to the file whitelist,
  23. // so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
  24. // Alternatively you can allow the use of all local files with
  25. // the DSN parameter 'allowAllFiles=true'
  26. //
  27. // filePath := "/home/gopher/data.csv"
  28. // mysql.RegisterLocalFile(filePath)
  29. // err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
  30. // if err != nil {
  31. // ...
  32. //
  33. func RegisterLocalFile(filePath string) {
  34. fileRegisterLock.Lock()
  35. // lazy map init
  36. if fileRegister == nil {
  37. fileRegister = make(map[string]bool)
  38. }
  39. fileRegister[strings.Trim(filePath, `"`)] = true
  40. fileRegisterLock.Unlock()
  41. }
  42. // DeregisterLocalFile removes the given filepath from the whitelist.
  43. func DeregisterLocalFile(filePath string) {
  44. fileRegisterLock.Lock()
  45. delete(fileRegister, strings.Trim(filePath, `"`))
  46. fileRegisterLock.Unlock()
  47. }
  48. // RegisterReaderHandler registers a handler function which is used
  49. // to receive a io.Reader.
  50. // The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
  51. // If the handler returns a io.ReadCloser Close() is called when the
  52. // request is finished.
  53. //
  54. // mysql.RegisterReaderHandler("data", func() io.Reader {
  55. // var csvReader io.Reader // Some Reader that returns CSV data
  56. // ... // Open Reader here
  57. // return csvReader
  58. // })
  59. // err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
  60. // if err != nil {
  61. // ...
  62. //
  63. func RegisterReaderHandler(name string, handler func() io.Reader) {
  64. readerRegisterLock.Lock()
  65. // lazy map init
  66. if readerRegister == nil {
  67. readerRegister = make(map[string]func() io.Reader)
  68. }
  69. readerRegister[name] = handler
  70. readerRegisterLock.Unlock()
  71. }
  72. // DeregisterReaderHandler removes the ReaderHandler function with
  73. // the given name from the registry.
  74. func DeregisterReaderHandler(name string) {
  75. readerRegisterLock.Lock()
  76. delete(readerRegister, name)
  77. readerRegisterLock.Unlock()
  78. }
  79. func deferredClose(err *error, closer io.Closer) {
  80. closeErr := closer.Close()
  81. if *err == nil {
  82. *err = closeErr
  83. }
  84. }
  85. func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
  86. var rdr io.Reader
  87. var data []byte
  88. packetSize := 16 * 1024 // 16KB is small enough for disk readahead and large enough for TCP
  89. if mc.maxWriteSize < packetSize {
  90. packetSize = mc.maxWriteSize
  91. }
  92. if idx := strings.Index(name, "Reader::"); idx == 0 || (idx > 0 && name[idx-1] == '/') { // io.Reader
  93. // The server might return an an absolute path. See issue #355.
  94. name = name[idx+8:]
  95. readerRegisterLock.RLock()
  96. handler, inMap := readerRegister[name]
  97. readerRegisterLock.RUnlock()
  98. if inMap {
  99. rdr = handler()
  100. if rdr != nil {
  101. if cl, ok := rdr.(io.Closer); ok {
  102. defer deferredClose(&err, cl)
  103. }
  104. } else {
  105. err = fmt.Errorf("Reader '%s' is <nil>", name)
  106. }
  107. } else {
  108. err = fmt.Errorf("Reader '%s' is not registered", name)
  109. }
  110. } else { // File
  111. name = strings.Trim(name, `"`)
  112. fileRegisterLock.RLock()
  113. fr := fileRegister[name]
  114. fileRegisterLock.RUnlock()
  115. if mc.cfg.AllowAllFiles || fr {
  116. var file *os.File
  117. var fi os.FileInfo
  118. if file, err = os.Open(name); err == nil {
  119. defer deferredClose(&err, file)
  120. // get file size
  121. if fi, err = file.Stat(); err == nil {
  122. rdr = file
  123. if fileSize := int(fi.Size()); fileSize < packetSize {
  124. packetSize = fileSize
  125. }
  126. }
  127. }
  128. } else {
  129. err = fmt.Errorf("local file '%s' is not registered", name)
  130. }
  131. }
  132. // send content packets
  133. if err == nil {
  134. data := make([]byte, 4+packetSize)
  135. var n int
  136. for err == nil {
  137. n, err = rdr.Read(data[4:])
  138. if n > 0 {
  139. if ioErr := mc.writePacket(data[:4+n]); ioErr != nil {
  140. return ioErr
  141. }
  142. }
  143. }
  144. if err == io.EOF {
  145. err = nil
  146. }
  147. }
  148. // send empty packet (termination)
  149. if data == nil {
  150. data = make([]byte, 4)
  151. }
  152. if ioErr := mc.writePacket(data[:4]); ioErr != nil {
  153. return ioErr
  154. }
  155. // read OK packet
  156. if err == nil {
  157. return mc.readResultOK()
  158. }
  159. mc.readPacket()
  160. return err
  161. }