poc: context support

Author: serprex

client/conn.go

@@ -35,6 +35,8 @@ type Conn struct {
 	// Connection read and write timeouts to set on the connection
 	ReadTimeout  time.Duration
 	WriteTimeout time.Duration
+	contexts     chan context.Context
+	closed       bool
 
 	// The buffer size to use in the packet connection
 	BufferSize int
@@ -136,6 +138,7 @@ func ConnectWithDialer(ctx context.Context, network, addr, user, password, dbNam
 	c.password = password
 	c.db = dbName
 	c.proto = network
+	c.contexts = make(chan context.Context)
 
 	// use default charset here, utf-8
 	c.charset = mysql.DEFAULT_CHARSET
@@ -184,6 +187,21 @@ func ConnectWithDialer(ctx context.Context, network, addr, user, password, dbNam
 		}
 	}
 
+	go func() {
+		ctx := context.Background()
+		for {
+			var ok bool
+			select {
+			case <-ctx.Done():
+				_ = c.Conn.SetDeadline(time.Unix(0, 0))
+			case ctx, ok = <-c.contexts:
+				if !ok {
+					return
+				}
+			}
+		}
+	}()
+
 	return c, nil
 }
 
@@ -208,8 +226,19 @@ func (c *Conn) handshake() error {
 	return nil
 }
 
+func (c *Conn) watchCtx(ctx context.Context) func() {
+	c.contexts <- ctx
+	return func() {
+		c.contexts <- context.Background()
+	}
+}
+
 // Close directly closes the connection. Use Quit() to first send COM_QUIT to the server and then close the connection.
 func (c *Conn) Close() error {
+	if !c.closed {
+		close(c.contexts)
+		c.closed = true
+	}
 	return c.Conn.Close()
 }
 
@@ -309,6 +338,11 @@ func (c *Conn) Execute(command string, args ...interface{}) (*mysql.Result, erro
 	}
 }
 
+func (c *Conn) ExecuteContext(ctx context.Context, command string, args ...interface{}) (*mysql.Result, error) {
+	defer c.watchCtx(ctx)
+	return c.Execute(command, args)
+}
+
 // ExecuteMultiple will call perResultCallback for every result of the multiple queries
 // that are executed.
 //
@@ -363,6 +397,11 @@ func (c *Conn) ExecuteMultiple(query string, perResultCallback ExecPerResultCall
 	return mysql.NewResult(rs), nil
 }
 
+func (c *Conn) ExecuteMultipleContext(ctx context.Context, query string, perResultCallback ExecPerResultCallback) (*mysql.Result, error) {
+	defer c.watchCtx(ctx)
+	return c.ExecuteMultiple(query, perResultCallback)
+}
+
 // ExecuteSelectStreaming will call perRowCallback for every row in resultset
 // WITHOUT saving any row data to Result.{Values/RawPkg/RowDatas} fields.
 // When given, perResultCallback will be called once per result
@@ -376,11 +415,33 @@ func (c *Conn) ExecuteSelectStreaming(command string, result *mysql.Result, perR
 	return c.readResultStreaming(false, result, perRowCallback, perResultCallback)
 }
 
+func (c *Conn) ExecuteSelectStreamingContext(ctx context.Context, command string, result *mysql.Result, perRowCallback SelectPerRowCallback, perResultCallback SelectPerResultCallback) error {
+	defer c.watchCtx(ctx)
+	return c.ExecuteSelectStreaming(command, result, perRowCallback, perResultCallback)
+}
+
 func (c *Conn) Begin() error {
 	_, err := c.exec("BEGIN")
 	return errors.Trace(err)
 }
 
+func (c *Conn) BeginTx(ctx context.Context, readOnly bool, txIsolation string) error {
+	defer c.watchCtx(ctx)()
+
+	if txIsolation != "" {
+		if _, err := c.exec("SET TRANSACTION ISOLATION LEVEL " + txIsolation); err != nil {
+			return errors.Trace(err)
+		}
+	}
+	var err error
+	if readOnly {
+		_, err = c.exec("START TRANSACTION READ ONLY")
+	} else {
+		_, err = c.exec("START TRANSACTION")
+	}
+	return errors.Trace(err)
+}
+
 func (c *Conn) Commit() error {
 	_, err := c.exec("COMMIT")
 	return errors.Trace(err)

driver/driver.go

@@ -3,6 +3,7 @@
 package driver
 
 import (
+	"context"
 	"crypto/tls"
 	"database/sql"
 	sqldriver "database/sql/driver"
@@ -184,8 +185,13 @@ func (d driver) Open(dsn string) (sqldriver.Conn, error) {
 
 type CheckNamedValueFunc func(*sqldriver.NamedValue) error
 
-var _ sqldriver.NamedValueChecker = &conn{}
-var _ sqldriver.Validator = &conn{}
+var (
+	_ sqldriver.NamedValueChecker = &conn{}
+	_ sqldriver.Validator         = &conn{}
+	_ sqldriver.Conn              = &conn{}
+	_ sqldriver.ConnBeginTx       = &conn{}
+	_ sqldriver.QueryerContext    = &conn{}
+)
 
 type state struct {
 	valid bool
@@ -242,6 +248,27 @@ func (c *conn) Begin() (sqldriver.Tx, error) {
 	return &tx{c.Conn}, nil
 }
 
+var isolationLevelTransactionIsolation = map[sql.IsolationLevel]string{
+	sql.LevelDefault:         "",
+	sql.LevelRepeatableRead:  "REPEATABLE READ",
+	sql.LevelReadCommitted:   "READ COMMITTED",
+	sql.LevelReadUncommitted: "READ UNCOMMITTED",
+	sql.LevelSerializable:    "SERIALIZABLE",
+}
+
+func (c *conn) BeginTx(ctx context.Context, opts sqldriver.TxOptions) (sqldriver.Tx, error) {
+	isolation := sql.IsolationLevel(opts.Isolation)
+	txIsolation, ok := isolationLevelTransactionIsolation[isolation]
+	if !ok {
+		return nil, fmt.Errorf("invalid mysql transaction isolation level %s", isolation)
+	}
+	err := c.Conn.BeginTx(ctx, opts.ReadOnly, txIsolation)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	return &tx{c.Conn}, nil
+}
+
 func buildArgs(args []sqldriver.Value) []interface{} {
 	a := make([]interface{}, len(args))
 
@@ -252,6 +279,17 @@ func buildArgs(args []sqldriver.Value) []interface{} {
 	return a
 }
 
+func buildNamedArgs(args []sqldriver.NamedValue) []interface{} {
+	a := make([]interface{}, len(args))
+
+	for i, arg := range args {
+		// TODO named parameter support
+		a[i] = arg.Value
+	}
+
+	return a
+}
+
 func (st *state) replyError(err error) error {
 	isBadConnection := mysql.ErrorEqual(err, mysql.ErrBadConn)
 
@@ -284,6 +322,15 @@ func (c *conn) Query(query string, args []sqldriver.Value) (sqldriver.Rows, erro
 	return newRows(r.Resultset)
 }
 
+func (c *conn) QueryContext(ctx context.Context, query string, args []sqldriver.NamedValue) (sqldriver.Rows, error) {
+	a := buildNamedArgs(args)
+	r, err := c.Conn.ExecuteContext(ctx, query, a...)
+	if err != nil {
+		return nil, c.state.replyError(err)
+	}
+	return newRows(r.Resultset)
+}
+
 type stmt struct {
 	*client.Stmt
 	connectionState *state

packet/conn.go

@@ -29,7 +29,6 @@ type Conn struct {
 
 	readTimeout  time.Duration
 	writeTimeout time.Duration
-	ctx          context.Context
 
 	// Buffered reader for net.Conn in Non-TLS connection only to address replication performance issue.
 	// See https://github.com/go-mysql-org/go-mysql/pull/422 for more details.