Add context support

For backward compatibility with Go versions < 1.8, the file ctx_backport.go
contains a redefinition of context.Context as well as a few other things.
The result is that the vast majority of the code can be shared between Go
versions.

Author: oscarzhao

README.md

@@ -418,6 +418,11 @@ Version 1.0 of the driver recommended adding `&charset=utf8` (alias for `SET NAM
 
 See http://dev.mysql.com/doc/refman/5.7/en/charset-unicode.html for more details on MySQL's Unicode support.
 
+## Context Support
+Since go1.8, context is introduced to `database/sql` for better control on timeout and cancellation.
+New interfaces such as `driver.QueryerContext`, `driver.ExecerContext` are introduced. See more details on [context support to database/sql package](https://golang.org/doc/go1.8#database_sql, "sql").
+
+In Go-MySQL-Driver, we implemented these interfaces for structs `mysqlConn`, `mysqlStmt` and `mysqlTx`.
 
 ## Testing / Development
 To run the driver tests you may need to adjust the configuration. See the [Testing Wiki-Page](https://github.com/go-sql-driver/mysql/wiki/Testing "Testing") for details.

connection.go

@@ -42,7 +42,7 @@ func (mc *mysqlConn) handleParams() (err error) {
 			charsets := strings.Split(val, ",")
 			for i := range charsets {
 				// ignore errors here - a charset may not exist
-				err = mc.exec("SET NAMES " + charsets[i])
+				err = mc.exec(backgroundCtx(), "SET NAMES "+charsets[i])
 				if err == nil {
 					break
 				}
@@ -53,7 +53,7 @@ func (mc *mysqlConn) handleParams() (err error) {
 
 		// System Vars
 		default:
-			err = mc.exec("SET " + param + "=" + val + "")
+			err = mc.exec(backgroundCtx(), "SET "+param+"="+val+"")
 			if err != nil {
 				return
 			}
@@ -63,12 +63,17 @@ func (mc *mysqlConn) handleParams() (err error) {
 	return
 }
 
+// Begin implements driver.Conn interface
 func (mc *mysqlConn) Begin() (driver.Tx, error) {
+	return mc.beginTx(backgroundCtx(), txOptions{})
+}
+
+func (mc *mysqlConn) beginTx(ctx mysqlContext, opts txOptions) (driver.Tx, error) {
 	if mc.netConn == nil {
 		errLog.Print(ErrInvalidConn)
 		return nil, driver.ErrBadConn
 	}
-	err := mc.exec("START TRANSACTION")
+	err := mc.exec(ctx, "START TRANSACTION")
 	if err == nil {
 		return &mysqlTx{mc}, err
 	}
@@ -79,7 +84,7 @@ func (mc *mysqlConn) Begin() (driver.Tx, error) {
 func (mc *mysqlConn) Close() (err error) {
 	// Makes Close idempotent
 	if mc.netConn != nil {
-		err = mc.writeCommandPacket(comQuit)
+		err = mc.writeCommandPacket(backgroundCtx(), comQuit)
 	}
 
 	mc.cleanup()
@@ -104,12 +109,16 @@ func (mc *mysqlConn) cleanup() {
 }
 
 func (mc *mysqlConn) Prepare(query string) (driver.Stmt, error) {
+	return mc.prepareContext(backgroundCtx(), query)
+}
+
+func (mc *mysqlConn) prepareContext(ctx mysqlContext, query string) (driver.Stmt, error) {
 	if mc.netConn == nil {
 		errLog.Print(ErrInvalidConn)
 		return nil, driver.ErrBadConn
 	}
 	// Send command
-	err := mc.writeCommandPacketStr(comStmtPrepare, query)
+	err := mc.writeCommandPacketStr(ctx, comStmtPrepare, query)
 	if err != nil {
 		return nil, err
 	}
@@ -258,6 +267,10 @@ func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (strin
 }
 
 func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, error) {
+	return mc.ExecContext(backgroundCtx(), query, args)
+}
+
+func (mc *mysqlConn) ExecContext(ctx mysqlContext, query string, args []driver.Value) (driver.Result, error) {
 	if mc.netConn == nil {
 		errLog.Print(ErrInvalidConn)
 		return nil, driver.ErrBadConn
@@ -276,7 +289,7 @@ func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, err
 	mc.affectedRows = 0
 	mc.insertId = 0
 
-	err := mc.exec(query)
+	err := mc.exec(ctx, query)
 	if err == nil {
 		return &mysqlResult{
 			affectedRows: int64(mc.affectedRows),
@@ -287,9 +300,9 @@ func (mc *mysqlConn) Exec(query string, args []driver.Value) (driver.Result, err
 }
 
 // Internal function to execute commands
-func (mc *mysqlConn) exec(query string) error {
+func (mc *mysqlConn) exec(ctx mysqlContext, query string) error {
 	// Send command
-	if err := mc.writeCommandPacketStr(comQuery, query); err != nil {
+	if err := mc.writeCommandPacketStr(ctx, comQuery, query); err != nil {
 		return err
 	}
 
@@ -314,7 +327,12 @@ func (mc *mysqlConn) exec(query string) error {
 	return mc.discardResults()
 }
 
+// Query implements driver.Queryer interface
 func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, error) {
+	return mc.queryContext(backgroundCtx(), query, args)
+}
+
+func (mc *mysqlConn) queryContext(ctx mysqlContext, query string, args []driver.Value) (driver.Rows, error) {
 	if mc.netConn == nil {
 		errLog.Print(ErrInvalidConn)
 		return nil, driver.ErrBadConn
@@ -331,7 +349,7 @@ func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, erro
 		query = prepared
 	}
 	// Send command
-	err := mc.writeCommandPacketStr(comQuery, query)
+	err := mc.writeCommandPacketStr(ctx, comQuery, query)
 	if err == nil {
 		// Read Result
 		var resLen int
@@ -362,7 +380,7 @@ func (mc *mysqlConn) Query(query string, args []driver.Value) (driver.Rows, erro
 // The returned byte slice is only valid until the next read
 func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
 	// Send command
-	if err := mc.writeCommandPacketStr(comQuery, "SELECT @@"+name); err != nil {
+	if err := mc.writeCommandPacketStr(backgroundCtx(), comQuery, "SELECT @@"+name); err != nil {
 		return nil, err
 	}
 

connection_ctx.go

@@ -0,0 +1,48 @@
+// +build go1.8
+
+// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
+//
+// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+)
+
+// Ping implements driver.Pinger interface
+func (mc *mysqlConn) Ping(ctx context.Context) error {
+	if mc.netConn == nil {
+		errLog.Print(ErrInvalidConn)
+		return driver.ErrBadConn
+	}
+	if err := mc.writeCommandPacket(ctx, comPing); err != nil {
+		errLog.Print(err)
+		return err
+	}
+
+	if _, err := mc.readResultOK(); err != nil {
+		errLog.Print(err)
+		return err
+	}
+	return nil
+}
+
+// BeginTx implements driver.ConnBeginTx interface
+func (mc *mysqlConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	return mc.beginTx(ctx, txOptions(opts))
+}
+
+func (mc *mysqlConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+	return mc.prepareContext(ctx, query)
+}
+
+// QueryContext implements driver.QueryerContext interface
+func (mc *mysqlConn) QueryContext(ctx context.Context, query string, args []driver.Value) (driver.Rows, error) {
+	return mc.queryContext(ctx, query, args)
+}

ctx_backport.go

@@ -0,0 +1,61 @@
+// +build !go1.8
+
+// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
+//
+// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package mysql
+
+import (
+	"time"
+)
+
+// txOptions is defined for compatibility with Go 1.8's driver.TxOptions struct.
+type txOptions struct {
+	Isolation int
+	ReadOnly  bool
+}
+
+// mysqlContext is a copy of context.Context from Go 1.7 and later.
+type mysqlContext interface {
+	Deadline() (deadline time.Time, ok bool)
+
+	Done() <-chan struct{}
+
+	Err() error
+
+	Value(key interface{}) interface{}
+}
+
+// emptyCtx is copied from Go 1.7's context package.
+type emptyCtx int
+
+func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
+	return
+}
+
+func (*emptyCtx) Done() <-chan struct{} {
+	return nil
+}
+
+func (*emptyCtx) Err() error {
+	return nil
+}
+
+func (*emptyCtx) Value(key interface{}) interface{} {
+	return nil
+}
+
+func (e *emptyCtx) String() string {
+	return "context.Background"
+}
+
+var background = new(emptyCtx)
+
+func backgroundCtx() mysqlContext {
+	return background
+}

ctx_go18.go

@@ -0,0 +1,27 @@
+// +build go1.8
+
+// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
+//
+// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+)
+
+// The definitions below are for compatibility with older Go versions.
+// See ctx_backport.go for the definitions used in older Go versions.
+
+type txOptions driver.TxOptions
+
+type mysqlContext context.Context
+
+func backgroundCtx() mysqlContext {
+	return context.Background()
+}

driver.go

@@ -14,6 +14,7 @@
 //  db, err := sql.Open("mysql", "user:password@/dbname")
 //
 // See https://github.com/go-sql-driver/mysql#usage for details
+
 package mysql
 
 import (
@@ -95,7 +96,7 @@ func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
 	}
 
 	// Send Client Authentication Packet
-	if err = mc.writeAuthPacket(cipher); err != nil {
+	if err = mc.writeAuthPacket(backgroundCtx(), cipher); err != nil {
 		mc.cleanup()
 		return nil, err
 	}
@@ -157,20 +158,20 @@ func handleAuthResult(mc *mysqlConn, oldCipher []byte) error {
 			cipher = oldCipher
 		}
 
-		if err = mc.writeOldAuthPacket(cipher); err != nil {
+		if err = mc.writeOldAuthPacket(backgroundCtx(), cipher); err != nil {
 			return err
 		}
 		_, err = mc.readResultOK()
 	} else if mc.cfg.AllowCleartextPasswords && err == ErrCleartextPassword {
 		// Retry with clear text password for
 		// http://dev.mysql.com/doc/refman/5.7/en/cleartext-authentication-plugin.html
 		// http://dev.mysql.com/doc/refman/5.7/en/pam-authentication-plugin.html
-		if err = mc.writeClearAuthPacket(); err != nil {
+		if err = mc.writeClearAuthPacket(backgroundCtx()); err != nil {
 			return err
 		}
 		_, err = mc.readResultOK()
 	} else if mc.cfg.AllowNativePasswords && err == ErrNativePassword {
-		if err = mc.writeNativeAuthPacket(cipher); err != nil {
+		if err = mc.writeNativeAuthPacket(backgroundCtx(), cipher); err != nil {
 			return err
 		}
 		_, err = mc.readResultOK()

driver_test.go

@@ -182,6 +182,14 @@ func TestEmptyQuery(t *testing.T) {
 	})
 }
 
+func (dbt *DBTest) TestPing(t *testing.T) {
+	runTests(t, dsn, func(dbt *DBTest) {
+		if err := dbt.db.Ping(); err != nil {
+			dbt.fail("Ping", "Ping", err)
+		}
+	})
+}
+
 func TestCRUD(t *testing.T) {
 	runTests(t, dsn, func(dbt *DBTest) {
 		// Create Table

infile.go

@@ -93,7 +93,7 @@ func deferredClose(err *error, closer io.Closer) {
 	}
 }
 
-func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
+func (mc *mysqlConn) handleInFileRequest(ctx mysqlContext, name string) (err error) {
 	var rdr io.Reader
 	var data []byte
 	packetSize := 16 * 1024 // 16KB is small enough for disk readahead and large enough for TCP
@@ -153,7 +153,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
 		for err == nil {
 			n, err = rdr.Read(data[4:])
 			if n > 0 {
-				if ioErr := mc.writePacket(data[:4+n]); ioErr != nil {
+				if ioErr := mc.writePacket(ctx, data[:4+n]); ioErr != nil {
 					return ioErr
 				}
 			}
@@ -167,7 +167,7 @@ func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
 	if data == nil {
 		data = make([]byte, 4)
 	}
-	if ioErr := mc.writePacket(data[:4]); ioErr != nil {
+	if ioErr := mc.writePacket(ctx, data[:4]); ioErr != nil {
 		return ioErr
 	}