driver.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169
  1. // Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
  2. //
  3. // This Source Code Form is subject to the terms of the Mozilla Public
  4. // License, v. 2.0. If a copy of the MPL was not distributed with this file,
  5. // You can obtain one at http://mozilla.org/MPL/2.0/.
  6. // Package mysql provides a MySQL driver for Go's database/sql package.
  7. //
  8. // The driver should be used via the database/sql package:
  9. //
  10. // import "database/sql"
  11. // import _ "github.com/go-sql-driver/mysql"
  12. //
  13. // db, err := sql.Open("mysql", "user:password@/dbname")
  14. //
  15. // See https://github.com/go-sql-driver/mysql#usage for details
  16. package mysql
  17. import (
  18. "database/sql"
  19. "database/sql/driver"
  20. "net"
  21. "sync"
  22. )
  23. // MySQLDriver is exported to make the driver directly accessible.
  24. // In general the driver is used via the database/sql package.
  25. type MySQLDriver struct{}
  26. // DialFunc is a function which can be used to establish the network connection.
  27. // Custom dial functions must be registered with RegisterDial
  28. type DialFunc func(addr string) (net.Conn, error)
  29. var (
  30. dialsLock sync.RWMutex
  31. dials map[string]DialFunc
  32. )
  33. // RegisterDial registers a custom dial function. It can then be used by the
  34. // network address mynet(addr), where mynet is the registered new network.
  35. // addr is passed as a parameter to the dial function.
  36. func RegisterDial(net string, dial DialFunc) {
  37. dialsLock.Lock()
  38. defer dialsLock.Unlock()
  39. if dials == nil {
  40. dials = make(map[string]DialFunc)
  41. }
  42. dials[net] = dial
  43. }
  44. // Open new Connection.
  45. // See https://github.com/go-sql-driver/mysql#dsn-data-source-name for how
  46. // the DSN string is formatted
  47. func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
  48. var err error
  49. // New mysqlConn
  50. mc := &mysqlConn{
  51. maxAllowedPacket: maxPacketSize,
  52. maxWriteSize: maxPacketSize - 1,
  53. closech: make(chan struct{}),
  54. }
  55. mc.cfg, err = ParseDSN(dsn)
  56. if err != nil {
  57. return nil, err
  58. }
  59. mc.parseTime = mc.cfg.ParseTime
  60. // Connect to Server
  61. dialsLock.RLock()
  62. dial, ok := dials[mc.cfg.Net]
  63. dialsLock.RUnlock()
  64. if ok {
  65. mc.netConn, err = dial(mc.cfg.Addr)
  66. } else {
  67. nd := net.Dialer{Timeout: mc.cfg.Timeout}
  68. mc.netConn, err = nd.Dial(mc.cfg.Net, mc.cfg.Addr)
  69. }
  70. if err != nil {
  71. if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
  72. errLog.Print("net.Error from Dial()': ", nerr.Error())
  73. return nil, driver.ErrBadConn
  74. }
  75. return nil, err
  76. }
  77. // Enable TCP Keepalives on TCP connections
  78. if tc, ok := mc.netConn.(*net.TCPConn); ok {
  79. if err := tc.SetKeepAlive(true); err != nil {
  80. // Don't send COM_QUIT before handshake.
  81. mc.netConn.Close()
  82. mc.netConn = nil
  83. return nil, err
  84. }
  85. }
  86. // Call startWatcher for context support (From Go 1.8)
  87. mc.startWatcher()
  88. mc.buf = newBuffer(mc.netConn)
  89. // Set I/O timeouts
  90. mc.buf.timeout = mc.cfg.ReadTimeout
  91. mc.writeTimeout = mc.cfg.WriteTimeout
  92. // Reading Handshake Initialization Packet
  93. authData, plugin, err := mc.readHandshakePacket()
  94. if err != nil {
  95. mc.cleanup()
  96. return nil, err
  97. }
  98. if plugin == "" {
  99. plugin = defaultAuthPlugin
  100. }
  101. // Send Client Authentication Packet
  102. authResp, err := mc.auth(authData, plugin)
  103. if err != nil {
  104. // try the default auth plugin, if using the requested plugin failed
  105. errLog.Print("could not use requested auth plugin '"+plugin+"': ", err.Error())
  106. plugin = defaultAuthPlugin
  107. authResp, err = mc.auth(authData, plugin)
  108. if err != nil {
  109. mc.cleanup()
  110. return nil, err
  111. }
  112. }
  113. if err = mc.writeHandshakeResponsePacket(authResp, plugin); err != nil {
  114. mc.cleanup()
  115. return nil, err
  116. }
  117. // Handle response to auth packet, switch methods if possible
  118. if err = mc.handleAuthResult(authData, plugin); err != nil {
  119. // Authentication failed and MySQL has already closed the connection
  120. // (https://dev.mysql.com/doc/internals/en/authentication-fails.html).
  121. // Do not send COM_QUIT, just cleanup and return the error.
  122. mc.cleanup()
  123. return nil, err
  124. }
  125. if mc.cfg.MaxAllowedPacket > 0 {
  126. mc.maxAllowedPacket = mc.cfg.MaxAllowedPacket
  127. } else {
  128. // Get max allowed packet size
  129. maxap, err := mc.getSystemVar("max_allowed_packet")
  130. if err != nil {
  131. mc.Close()
  132. return nil, err
  133. }
  134. mc.maxAllowedPacket = stringToInt(maxap) - 1
  135. }
  136. if mc.maxAllowedPacket < maxPacketSize {
  137. mc.maxWriteSize = mc.maxAllowedPacket
  138. }
  139. // Handle DSN Params
  140. err = mc.handleParams()
  141. if err != nil {
  142. mc.Close()
  143. return nil, err
  144. }
  145. return mc, nil
  146. }
  147. func init() {
  148. sql.Register("mysql", &MySQLDriver{})
  149. }