connection.go 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461
  1. // Go MySQL Driver - A MySQL-Driver for Go's database/sql package
  2. //
  3. // Copyright 2012 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. "database/sql/driver"
  11. "io"
  12. "net"
  13. "strconv"
  14. "strings"
  15. "time"
  16. )
  17. // a copy of context.Context for Go 1.7 and earlier
  18. type mysqlContext interface {
  19. Done() <-chan struct{}
  20. Err() error
  21. // defined in context.Context, but not used in this driver:
  22. // Deadline() (deadline time.Time, ok bool)
  23. // Value(key interface{}) interface{}
  24. }
  25. type mysqlConn struct {
  26. buf buffer
  27. netConn net.Conn
  28. affectedRows uint64
  29. insertId uint64
  30. cfg *Config
  31. maxAllowedPacket int
  32. maxWriteSize int
  33. writeTimeout time.Duration
  34. flags clientFlag
  35. status statusFlag
  36. sequence uint8
  37. parseTime bool
  38. // for context support (Go 1.8+)
  39. watching bool
  40. watcher chan<- mysqlContext
  41. closech chan struct{}
  42. finished chan<- struct{}
  43. canceled atomicError // set non-nil if conn is canceled
  44. closed atomicBool // set when conn is closed, before closech is closed
  45. }
  46. // Handles parameters set in DSN after the connection is established
  47. func (mc *mysqlConn) handleParams() (err error) {
  48. for param, val := range mc.cfg.Params {
  49. switch param {
  50. // Charset
  51. case "charset":
  52. charsets := strings.Split(val, ",")
  53. for i := range charsets {
  54. // ignore errors here - a charset may not exist
  55. err = mc.exec("SET NAMES " + charsets[i])
  56. if err == nil {
  57. break
  58. }
  59. }
  60. if err != nil {
  61. return
  62. }
  63. // System Vars
  64. default:
  65. err = mc.exec("SET " + param + "=" + val + "")
  66. if err != nil {
  67. return
  68. }
  69. }
  70. }
  71. return
  72. }
  73. func (mc *mysqlConn) markBadConn(err error) error {
  74. if mc == nil {
  75. return err
  76. }
  77. if err != errBadConnNoWrite {
  78. return err
  79. }
  80. return driver.ErrBadConn
  81. }
  82. func (mc *mysqlConn) Begin() (driver.Tx, error) {
  83. return mc.begin(false)
  84. }
  85. func (mc *mysqlConn) begin(readOnly bool) (driver.Tx, error) {
  86. if mc.closed.IsSet() {
  87. errLog.Print(ErrInvalidConn)
  88. return nil, driver.ErrBadConn
  89. }
  90. var q string
  91. if readOnly {
  92. q = "START TRANSACTION READ ONLY"
  93. } else {
  94. q = "START TRANSACTION"
  95. }
  96. err := mc.exec(q)
  97. if err == nil {
  98. return &mysqlTx{mc}, err
  99. }
  100. return nil, mc.markBadConn(err)
  101. }
  102. func (mc *mysqlConn) Close() (err error) {
  103. // Makes Close idempotent
  104. if !mc.closed.IsSet() {
  105. err = mc.writeCommandPacket(comQuit)
  106. }
  107. mc.cleanup()
  108. return
  109. }
  110. // Closes the network connection and unsets internal variables. Do not call this
  111. // function after successfully authentication, call Close instead. This function
  112. // is called before auth or on auth failure because MySQL will have already
  113. // closed the network connection.
  114. func (mc *mysqlConn) cleanup() {
  115. if !mc.closed.TrySet(true) {
  116. return
  117. }
  118. // Makes cleanup idempotent
  119. close(mc.closech)
  120. if mc.netConn == nil {
  121. return
  122. }
  123. if err := mc.netConn.Close(); err != nil {
  124. errLog.Print(err)
  125. }
  126. }
  127. func (mc *mysqlConn) error() error {
  128. if mc.closed.IsSet() {
  129. if err := mc.canceled.Value(); err != nil {
  130. return err
  131. }
  132. return ErrInvalidConn
  133. }
  134. return nil
  135. }
  136. func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) {
  137. if mc.closed.IsSet() {
  138. errLog.Print(ErrInvalidConn)
  139. return nil, driver.ErrBadConn
  140. }
  141. // Send command
  142. err := mc.writeCommandPacketStr(comStmtPrepare, query)
  143. if err != nil {
  144. return nil, mc.markBadConn(err)
  145. }
  146. stmt := &mysqlStmt{
  147. mc: mc,
  148. }
  149. // Read Result
  150. columnCount, err := stmt.readPrepareResultPacket()
  151. if err == nil {
  152. if stmt.paramCount > 0 {
  153. if err = mc.readUntilEOF(); err != nil {
  154. return nil, err
  155. }
  156. }
  157. if columnCount > 0 {
  158. err = mc.readUntilEOF()
  159. }
  160. }
  161. return stmt, err
  162. }
  163. func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (string, error) {
  164. // Number of ? should be same to len(args)
  165. if strings.Count(query, "?") != len(args) {
  166. return "", driver.ErrSkip
  167. }
  168. buf := mc.buf.takeCompleteBuffer()
  169. if buf == nil {
  170. // can not take the buffer. Something must be wrong with the connection
  171. errLog.Print(ErrBusyBuffer)
  172. return "", ErrInvalidConn
  173. }
  174. buf = buf[:0]
  175. argPos := 0
  176. for i := 0; i < len(query); i++ {
  177. q := strings.IndexByte(query[i:], '?')
  178. if q == -1 {
  179. buf = append(buf, query[i:]...)
  180. break
  181. }
  182. buf = append(buf, query[i:i+q]...)
  183. i += q
  184. arg := args[argPos]
  185. argPos++
  186. if arg == nil {
  187. buf = append(buf, "NULL"...)
  188. continue
  189. }
  190. switch v := arg.(type) {
  191. case int64:
  192. buf = strconv.AppendInt(buf, v, 10)
  193. case float64:
  194. buf = strconv.AppendFloat(buf, v, 'g', -1, 64)
  195. case bool:
  196. if v {
  197. buf = append(buf, '1')
  198. } else {
  199. buf = append(buf, '0')
  200. }
  201. case time.Time:
  202. if v.IsZero() {
  203. buf = append(buf, "'0000-00-00'"...)
  204. } else {
  205. v := v.In(mc.cfg.Loc)
  206. v = v.Add(time.Nanosecond * 500) // To round under microsecond
  207. year := v.Year()
  208. year100 := year / 100
  209. year1 := year % 100
  210. month := v.Month()
  211. day := v.Day()
  212. hour := v.Hour()
  213. minute := v.Minute()
  214. second := v.Second()
  215. micro := v.Nanosecond() / 1000
  216. buf = append(buf, []byte{
  217. '\'',
  218. digits10[year100], digits01[year100],
  219. digits10[year1], digits01[year1],
  220. '-',
  221. digits10[month], digits01[month],
  222. '-',
  223. digits10[day], digits01[day],
  224. ' ',
  225. digits10[hour], digits01[hour],
  226. ':',
  227. digits10[minute], digits01[minute],
  228. ':',
  229. digits10[second], digits01[second],
  230. }...)
  231. if micro != 0 {
  232. micro10000 := micro / 10000
  233. micro100 := micro / 100 % 100
  234. micro1 := micro % 100
  235. buf = append(buf, []byte{
  236. '.',
  237. digits10[micro10000], digits01[micro10000],
  238. digits10[micro100], digits01[micro100],
  239. digits10[micro1], digits01[micro1],
  240. }...)
  241. }
  242. buf = append(buf, '\'')
  243. }
  244. case []byte:
  245. if v == nil {
  246. buf = append(buf, "NULL"...)
  247. } else {
  248. buf = append(buf, "_binary'"...)
  249. if mc.status&statusNoBackslashEscapes == 0 {
  250. buf = escapeBytesBackslash(buf, v)
  251. } else {
  252. buf = escapeBytesQuotes(buf, v)
  253. }
  254. buf = append(buf, '\'')
  255. }
  256. case string:
  257. buf = append(buf, '\'')
  258. if mc.status&statusNoBackslashEscapes == 0 {
  259. buf = escapeStringBackslash(buf, v)
  260. } else {
  261. buf = escapeStringQuotes(buf, v)
  262. }
  263. buf = append(buf, '\'')
  264. default:
  265. return "", driver.ErrSkip
  266. }
  267. if len(buf)+4 > mc.maxAllowedPacket {
  268. return "", driver.ErrSkip
  269. }
  270. }
  271. if argPos != len(args) {
  272. return "", driver.ErrSkip
  273. }
  274. return string(buf), nil
  275. }
  276. func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, error) {
  277. if mc.closed.IsSet() {
  278. errLog.Print(ErrInvalidConn)
  279. return nil, driver.ErrBadConn
  280. }
  281. if len(args) != 0 {
  282. if !mc.cfg.InterpolateParams {
  283. return nil, driver.ErrSkip
  284. }
  285. // try to interpolate the parameters to save extra roundtrips for preparing and closing a statement
  286. prepared, err := mc.interpolateParams(query, args)
  287. if err != nil {
  288. return nil, err
  289. }
  290. query = prepared
  291. }
  292. mc.affectedRows = 0
  293. mc.insertId = 0
  294. err := mc.exec(query)
  295. if err == nil {
  296. return &mysqlResult{
  297. affectedRows: int64(mc.affectedRows),
  298. insertId: int64(mc.insertId),
  299. }, err
  300. }
  301. return nil, mc.markBadConn(err)
  302. }
  303. // Internal function to execute commands
  304. func (mc *mysqlConn) exec(query string) error {
  305. // Send command
  306. if err := mc.writeCommandPacketStr(comQuery, query); err != nil {
  307. return mc.markBadConn(err)
  308. }
  309. // Read Result
  310. resLen, err := mc.readResultSetHeaderPacket()
  311. if err != nil {
  312. return err
  313. }
  314. if resLen > 0 {
  315. // columns
  316. if err := mc.readUntilEOF(); err != nil {
  317. return err
  318. }
  319. // rows
  320. if err := mc.readUntilEOF(); err != nil {
  321. return err
  322. }
  323. }
  324. return mc.discardResults()
  325. }
  326. func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, error) {
  327. return mc.query(query, args)
  328. }
  329. func (mc *mysqlConn) query(query string, args []driver.Value) (*textRows, error) {
  330. if mc.closed.IsSet() {
  331. errLog.Print(ErrInvalidConn)
  332. return nil, driver.ErrBadConn
  333. }
  334. if len(args) != 0 {
  335. if !mc.cfg.InterpolateParams {
  336. return nil, driver.ErrSkip
  337. }
  338. // try client-side prepare to reduce roundtrip
  339. prepared, err := mc.interpolateParams(query, args)
  340. if err != nil {
  341. return nil, err
  342. }
  343. query = prepared
  344. }
  345. // Send command
  346. err := mc.writeCommandPacketStr(comQuery, query)
  347. if err == nil {
  348. // Read Result
  349. var resLen int
  350. resLen, err = mc.readResultSetHeaderPacket()
  351. if err == nil {
  352. rows := new(textRows)
  353. rows.mc = mc
  354. if resLen == 0 {
  355. rows.rs.done = true
  356. switch err := rows.NextResultSet(); err {
  357. case nil, io.EOF:
  358. return rows, nil
  359. default:
  360. return nil, err
  361. }
  362. }
  363. // Columns
  364. rows.rs.columns, err = mc.readColumns(resLen)
  365. return rows, err
  366. }
  367. }
  368. return nil, mc.markBadConn(err)
  369. }
  370. // Gets the value of the given MySQL System Variable
  371. // The returned byte slice is only valid until the next read
  372. func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
  373. // Send command
  374. if err := mc.writeCommandPacketStr(comQuery, "SELECT @@"+name); err != nil {
  375. return nil, err
  376. }
  377. // Read Result
  378. resLen, err := mc.readResultSetHeaderPacket()
  379. if err == nil {
  380. rows := new(textRows)
  381. rows.mc = mc
  382. rows.rs.columns = []mysqlField{{fieldType: fieldTypeVarChar}}
  383. if resLen > 0 {
  384. // Columns
  385. if err := mc.readUntilEOF(); err != nil {
  386. return nil, err
  387. }
  388. }
  389. dest := make([]driver.Value, resLen)
  390. if err = rows.readRow(dest); err == nil {
  391. return dest[0].([]byte), mc.readUntilEOF()
  392. }
  393. }
  394. return nil, err
  395. }
  396. // finish is called when the query has canceled.
  397. func (mc *mysqlConn) cancel(err error) {
  398. mc.canceled.Set(err)
  399. mc.cleanup()
  400. }
  401. // finish is called when the query has succeeded.
  402. func (mc *mysqlConn) finish() {
  403. if !mc.watching || mc.finished == nil {
  404. return
  405. }
  406. select {
  407. case mc.finished <- struct{}{}:
  408. mc.watching = false
  409. case <-mc.closech:
  410. }
  411. }