Update go-sql-driver/mysql package (#7019)

2026-07-19 12:10:24 +03:00 · 2018-12-29 21:59:03 +05:30
parent f24c017e9a
commit 75cd4201b0
20 changed files with 2008 additions and 801 deletions
@@ -9,7 +9,10 @@
 package mysql

 import (
+	"context"
+	"database/sql"
 	"database/sql/driver"
+	"io"
 	"net"
 	"strconv"
 	"strings"
@@ -29,7 +32,14 @@ type mysqlConn struct {
 	status           statusFlag
 	sequence         uint8
 	parseTime        bool
-	strict           bool
+
+	// for context support (Go 1.8+)
+	watching bool
+	watcher  chan<- context.Context
+	closech  chan struct{}
+	finished chan<- struct{}
+	canceled atomicError // set non-nil if conn is canceled
+	closed   atomicBool  // set when conn is closed, before closech is closed
 }

 // Handles parameters set in DSN after the connection is established
@@ -62,22 +72,41 @@ func (mc *mysqlConn) handleParams() (err error) {
 	return
 }

+func (mc *mysqlConn) markBadConn(err error) error {
+	if mc == nil {
+		return err
+	}
+	if err != errBadConnNoWrite {
+		return err
+	}
+	return driver.ErrBadConn
+}
+
 func (mc *mysqlConn) Begin() (driver.Tx, error) {
-	if mc.netConn == nil {
+	return mc.begin(false)
+}
+
+func (mc *mysqlConn) begin(readOnly bool) (driver.Tx, error) {
+	if mc.closed.IsSet() {
 		errLog.Print(ErrInvalidConn)
 		return nil, driver.ErrBadConn
 	}
-	err := mc.exec("START TRANSACTION")
+	var q string
+	if readOnly {
+		q = "START TRANSACTION READ ONLY"
+	} else {
+		q = "START TRANSACTION"
+	}
+	err := mc.exec(q)
 	if err == nil {
 		return &mysqlTx{mc}, err
 	}
-
-	return nil, err
+	return nil, mc.markBadConn(err)
 }

 func (mc *mysqlConn) Close() (err error) {
 	// Makes Close idempotent
-	if mc.netConn != nil {
+	if !mc.closed.IsSet() {
 		err = mc.writeCommandPacket(comQuit)
 	}

@@ -91,26 +120,39 @@ func (mc *mysqlConn) Close() (err error) {
 // is called before auth or on auth failure because MySQL will have already
 // closed the network connection.
 func (mc *mysqlConn) cleanup() {
-	// Makes cleanup idempotent
-	if mc.netConn != nil {
-		if err := mc.netConn.Close(); err != nil {
-			errLog.Print(err)
-		}
-		mc.netConn = nil
+	if !mc.closed.TrySet(true) {
+		return
 	}
-	mc.cfg = nil
-	mc.buf.nc = nil
+
+	// Makes cleanup idempotent
+	close(mc.closech)
+	if mc.netConn == nil {
+		return
+	}
+	if err := mc.netConn.Close(); err != nil {
+		errLog.Print(err)
+	}
+}
+
+func (mc *mysqlConn) error() error {
+	if mc.closed.IsSet() {
+		if err := mc.canceled.Value(); err != nil {
+			return err
+		}
+		return ErrInvalidConn
+	}
+	return nil
 }

 func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) {
-	if mc.netConn == nil {
+	if mc.closed.IsSet() {
 		errLog.Print(ErrInvalidConn)
 		return nil, driver.ErrBadConn
 	}
 	// Send command
 	err := mc.writeCommandPacketStr(comStmtPrepare, query)
 	if err != nil {
-		return nil, err
+		return nil, mc.markBadConn(err)
 	}

 	stmt := &mysqlStmt{
@@ -140,11 +182,11 @@ func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (strin
 		return "", driver.ErrSkip
 	}

-	buf := mc.buf.takeCompleteBuffer()
-	if buf == nil {
+	buf, err := mc.buf.takeCompleteBuffer()
+	if err != nil {
 		// can not take the buffer. Something must be wrong with the connection
-		errLog.Print(ErrBusyBuffer)
-		return "", driver.ErrBadConn
+		errLog.Print(err)
+		return "", ErrInvalidConn
 	}
 	buf = buf[:0]
 	argPos := 0
@@ -257,7 +299,7 @@ func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (strin
 }

 func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, error) {
-	if mc.netConn == nil {
+	if mc.closed.IsSet() {
 		errLog.Print(ErrInvalidConn)
 		return nil, driver.ErrBadConn
 	}
@@ -271,7 +313,6 @@ func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, err
 			return nil, err
 		}
 		query = prepared
-		args = nil
 	}
 	mc.affectedRows = 0
 	mc.insertId = 0
@@ -283,32 +324,43 @@ func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, err
 			insertId:     int64(mc.insertId),
 		}, err
 	}
-	return nil, err
+	return nil, mc.markBadConn(err)
 }

 // Internal function to execute commands
 func (mc *mysqlConn) exec(query string) error {
 	// Send command
-	err := mc.writeCommandPacketStr(comQuery, query)
-	if err != nil {
-		return err
+	if err := mc.writeCommandPacketStr(comQuery, query); err != nil {
+		return mc.markBadConn(err)
 	}

 	// Read Result
 	resLen, err := mc.readResultSetHeaderPacket()
-	if err == nil && resLen > 0 {
-		if err = mc.readUntilEOF(); err != nil {
+	if err != nil {
+		return err
+	}
+
+	if resLen > 0 {
+		// columns
+		if err := mc.readUntilEOF(); err != nil {
 			return err
 		}

-		err = mc.readUntilEOF()
+		// rows
+		if err := mc.readUntilEOF(); err != nil {
+			return err
+		}
 	}

-	return err
+	return mc.discardResults()
 }

 func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, error) {
-	if mc.netConn == nil {
+	return mc.query(query, args)
+}
+
+func (mc *mysqlConn) query(query string, args []driver.Value) (*textRows, error) {
+	if mc.closed.IsSet() {
 		errLog.Print(ErrInvalidConn)
 		return nil, driver.ErrBadConn
 	}
@@ -322,7 +374,6 @@ func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, erro
 			return nil, err
 		}
 		query = prepared
-		args = nil
 	}
 	// Send command
 	err := mc.writeCommandPacketStr(comQuery, query)
@@ -335,15 +386,22 @@ func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, erro
 			rows.mc = mc

 			if resLen == 0 {
-				// no columns, no more data
-				return emptyRows{}, nil
+				rows.rs.done = true
+
+				switch err := rows.NextResultSet(); err {
+				case nil, io.EOF:
+					return rows, nil
+				default:
+					return nil, err
+				}
 			}
+
 			// Columns
-			rows.columns, err = mc.readColumns(resLen)
+			rows.rs.columns, err = mc.readColumns(resLen)
 			return rows, err
 		}
 	}
-	return nil, err
+	return nil, mc.markBadConn(err)
 }

 // Gets the value of the given MySQL System Variable
@@ -359,7 +417,7 @@ func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
 	if err == nil {
 		rows := new(textRows)
 		rows.mc = mc
-		rows.columns = []mysqlField{{fieldType: fieldTypeVarChar}}
+		rows.rs.columns = []mysqlField{{fieldType: fieldTypeVarChar}}

 		if resLen > 0 {
 			// Columns
@@ -375,3 +433,211 @@ func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
 	}
 	return nil, err
 }
+
+// finish is called when the query has canceled.
+func (mc *mysqlConn) cancel(err error) {
+	mc.canceled.Set(err)
+	mc.cleanup()
+}
+
+// finish is called when the query has succeeded.
+func (mc *mysqlConn) finish() {
+	if !mc.watching || mc.finished == nil {
+		return
+	}
+	select {
+	case mc.finished <- struct{}{}:
+		mc.watching = false
+	case <-mc.closech:
+	}
+}
+
+// Ping implements driver.Pinger interface
+func (mc *mysqlConn) Ping(ctx context.Context) (err error) {
+	if mc.closed.IsSet() {
+		errLog.Print(ErrInvalidConn)
+		return driver.ErrBadConn
+	}
+
+	if err = mc.watchCancel(ctx); err != nil {
+		return
+	}
+	defer mc.finish()
+
+	if err = mc.writeCommandPacket(comPing); err != nil {
+		return mc.markBadConn(err)
+	}
+
+	return mc.readResultOK()
+}
+
+// BeginTx implements driver.ConnBeginTx interface
+func (mc *mysqlConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	if err := mc.watchCancel(ctx); err != nil {
+		return nil, err
+	}
+	defer mc.finish()
+
+	if sql.IsolationLevel(opts.Isolation) != sql.LevelDefault {
+		level, err := mapIsolationLevel(opts.Isolation)
+		if err != nil {
+			return nil, err
+		}
+		err = mc.exec("SET TRANSACTION ISOLATION LEVEL " + level)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return mc.begin(opts.ReadOnly)
+}
+
+func (mc *mysqlConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	dargs, err := namedValueToValue(args)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := mc.watchCancel(ctx); err != nil {
+		return nil, err
+	}
+
+	rows, err := mc.query(query, dargs)
+	if err != nil {
+		mc.finish()
+		return nil, err
+	}
+	rows.finish = mc.finish
+	return rows, err
+}
+
+func (mc *mysqlConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+	dargs, err := namedValueToValue(args)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := mc.watchCancel(ctx); err != nil {
+		return nil, err
+	}
+	defer mc.finish()
+
+	return mc.Exec(query, dargs)
+}
+
+func (mc *mysqlConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+	if err := mc.watchCancel(ctx); err != nil {
+		return nil, err
+	}
+
+	stmt, err := mc.Prepare(query)
+	mc.finish()
+	if err != nil {
+		return nil, err
+	}
+
+	select {
+	default:
+	case <-ctx.Done():
+		stmt.Close()
+		return nil, ctx.Err()
+	}
+	return stmt, nil
+}
+
+func (stmt *mysqlStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
+	dargs, err := namedValueToValue(args)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := stmt.mc.watchCancel(ctx); err != nil {
+		return nil, err
+	}
+
+	rows, err := stmt.query(dargs)
+	if err != nil {
+		stmt.mc.finish()
+		return nil, err
+	}
+	rows.finish = stmt.mc.finish
+	return rows, err
+}
+
+func (stmt *mysqlStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
+	dargs, err := namedValueToValue(args)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := stmt.mc.watchCancel(ctx); err != nil {
+		return nil, err
+	}
+	defer stmt.mc.finish()
+
+	return stmt.Exec(dargs)
+}
+
+func (mc *mysqlConn) watchCancel(ctx context.Context) error {
+	if mc.watching {
+		// Reach here if canceled,
+		// so the connection is already invalid
+		mc.cleanup()
+		return nil
+	}
+	// When ctx is already cancelled, don't watch it.
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	// When ctx is not cancellable, don't watch it.
+	if ctx.Done() == nil {
+		return nil
+	}
+	// When watcher is not alive, can't watch it.
+	if mc.watcher == nil {
+		return nil
+	}
+
+	mc.watching = true
+	mc.watcher <- ctx
+	return nil
+}
+
+func (mc *mysqlConn) startWatcher() {
+	watcher := make(chan context.Context, 1)
+	mc.watcher = watcher
+	finished := make(chan struct{})
+	mc.finished = finished
+	go func() {
+		for {
+			var ctx context.Context
+			select {
+			case ctx = <-watcher:
+			case <-mc.closech:
+				return
+			}
+
+			select {
+			case <-ctx.Done():
+				mc.cancel(ctx.Err())
+			case <-finished:
+			case <-mc.closech:
+				return
+			}
+		}
+	}()
+}
+
+func (mc *mysqlConn) CheckNamedValue(nv *driver.NamedValue) (err error) {
+	nv.Value, err = converter{}.ConvertValue(nv.Value)
+	return
+}
+
+// ResetSession implements driver.SessionResetter.
+// (From Go 1.10)
+func (mc *mysqlConn) ResetSession(ctx context.Context) error {
+	if mc.closed.IsSet() {
+		return driver.ErrBadConn
+	}
+	return nil
+}