database/sql: add context methods

Add context methods to sql and sql/driver methods. If
the driver doesn't implement context methods the connection
pool will still handle timeouts when a query fails to return
in time or when a connection is not available from the pool
in time.

There will be a follow-up CL that will add support for
context values that specify transaction levels and modes
that a driver can use.

Fixes #15123

Change-Id: Ia99f3957aa3f177b23044dd99d4ec217491a30a7
Reviewed-on: https://go-review.googlesource.com/29381Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

src/database/sql/ctxutil.go

+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package sql
+
+import (
+	"context"
+	"database/sql/driver"
+	"errors"
+)
+
+func ctxDriverPrepare(ctx context.Context, ci driver.Conn, query string) (driver.Stmt, error) {
+	if ciCtx, is := ci.(driver.ConnPrepareContext); is {
+		return ciCtx.PrepareContext(ctx, query)
+	}
+	type R struct {
+		err   error
+		panic interface{}
+		si    driver.Stmt
+	}
+
+	rc := make(chan R, 1)
+	go func() {
+		r := R{}
+		defer func() {
+			if v := recover(); v != nil {
+				r.panic = v
+			}
+			rc <- r
+		}()
+		r.si, r.err = ci.Prepare(query)
+	}()
+	select {
+	case <-ctx.Done():
+		go func() {
+			<-rc
+			close(rc)
+		}()
+		return nil, ctx.Err()
+	case r := <-rc:
+		if r.panic != nil {
+			panic(r.panic)
+		}
+		return r.si, r.err
+	}
+}
+
+func ctxDriverExec(ctx context.Context, execer driver.Execer, query string, dargs []driver.Value) (driver.Result, error) {
+	if execerCtx, is := execer.(driver.ExecerContext); is {
+		return execerCtx.ExecContext(ctx, query, dargs)
+	}
+	type R struct {
+		err   error
+		panic interface{}
+		resi  driver.Result
+	}
+
+	rc := make(chan R, 1)
+	go func() {
+		r := R{}
+		defer func() {
+			if v := recover(); v != nil {
+				r.panic = v
+			}
+			rc <- r
+		}()
+		r.resi, r.err = execer.Exec(query, dargs)
+	}()
+	select {
+	case <-ctx.Done():
+		go func() {
+			<-rc
+			close(rc)
+		}()
+		return nil, ctx.Err()
+	case r := <-rc:
+		if r.panic != nil {
+			panic(r.panic)
+		}
+		return r.resi, r.err
+	}
+}
+
+func ctxDriverQuery(ctx context.Context, queryer driver.Queryer, query string, dargs []driver.Value) (driver.Rows, error) {
+	if queryerCtx, is := queryer.(driver.QueryerContext); is {
+		return queryerCtx.QueryContext(ctx, query, dargs)
+	}
+	type R struct {
+		err   error
+		panic interface{}
+		rowsi driver.Rows
+	}
+
+	rc := make(chan R, 1)
+	go func() {
+		r := R{}
+		defer func() {
+			if v := recover(); v != nil {
+				r.panic = v
+			}
+			rc <- r
+		}()
+		r.rowsi, r.err = queryer.Query(query, dargs)
+	}()
+	select {
+	case <-ctx.Done():
+		go func() {
+			<-rc
+			close(rc)
+		}()
+		return nil, ctx.Err()
+	case r := <-rc:
+		if r.panic != nil {
+			panic(r.panic)
+		}
+		return r.rowsi, r.err
+	}
+}
+
+func ctxDriverStmtExec(ctx context.Context, si driver.Stmt, dargs []driver.Value) (driver.Result, error) {
+	if siCtx, is := si.(driver.StmtExecContext); is {
+		return siCtx.ExecContext(ctx, dargs)
+	}
+	type R struct {
+		err   error
+		panic interface{}
+		resi  driver.Result
+	}
+
+	rc := make(chan R, 1)
+	go func() {
+		r := R{}
+		defer func() {
+			if v := recover(); v != nil {
+				r.panic = v
+			}
+			rc <- r
+		}()
+		r.resi, r.err = si.Exec(dargs)
+	}()
+	select {
+	case <-ctx.Done():
+		go func() {
+			<-rc
+			close(rc)
+		}()
+		return nil, ctx.Err()
+	case r := <-rc:
+		if r.panic != nil {
+			panic(r.panic)
+		}
+		return r.resi, r.err
+	}
+}
+
+func ctxDriverStmtQuery(ctx context.Context, si driver.Stmt, dargs []driver.Value) (driver.Rows, error) {
+	if siCtx, is := si.(driver.StmtQueryContext); is {
+		return siCtx.QueryContext(ctx, dargs)
+	}
+	type R struct {
+		err   error
+		panic interface{}
+		rowsi driver.Rows
+	}
+
+	rc := make(chan R, 1)
+	go func() {
+		r := R{}
+		defer func() {
+			if v := recover(); v != nil {
+				r.panic = v
+			}
+			rc <- r
+		}()
+		r.rowsi, r.err = si.Query(dargs)
+	}()
+	select {
+	case <-ctx.Done():
+		go func() {
+			<-rc
+			close(rc)
+		}()
+		return nil, ctx.Err()
+	case r := <-rc:
+		if r.panic != nil {
+			panic(r.panic)
+		}
+		return r.rowsi, r.err
+	}
+}
+
+var errLevelNotSupported = errors.New("sql: selected isolation level is not supported")
+
+func ctxDriverBegin(ctx context.Context, ci driver.Conn) (driver.Tx, error) {
+	if ciCtx, is := ci.(driver.ConnBeginContext); is {
+		return ciCtx.BeginContext(ctx)
+	}
+	// TODO(kardianos): check the transaction level in ctx. If set and non-default
+	// then return an error here as the BeginContext driver value is not supported.
+
+	type R struct {
+		err   error
+		panic interface{}
+		txi   driver.Tx
+	}
+	rc := make(chan R, 1)
+	go func() {
+		r := R{}
+		defer func() {
+			if v := recover(); v != nil {
+				r.panic = v
+			}
+			rc <- r
+		}()
+		r.txi, r.err = ci.Begin()
+	}()
+	select {
+	case <-ctx.Done():
+		go func() {
+			<-rc
+			close(rc)
+		}()
+		return nil, ctx.Err()
+	case r := <-rc:
+		if r.panic != nil {
+			panic(r.panic)
+		}
+		return r.txi, r.err
+	}
+}

src/database/sql/driver/driver.go

@@ -8,7 +8,10 @@
 // Most code should use package sql.
 package driver
 
-import "errors"
+import (
+	"context"
+	"errors"
+)
 
 // Value is a value that drivers must be able to handle.
 // It is either nil or an instance of one of these types:
@@ -65,6 +68,12 @@ type Execer interface {
 	Exec(query string, args []Value) (Result, error)
 }
 
+// ExecerContext is like execer, but must honor the context timeout and return
+// when the context is cancelled.
+type ExecerContext interface {
+	ExecContext(ctx context.Context, query string, args []Value) (Result, error)
+}
+
 // Queryer is an optional interface that may be implemented by a Conn.
 //
 // If a Conn does not implement Queryer, the sql package's DB.Query will
@@ -76,6 +85,12 @@ type Queryer interface {
 	Query(query string, args []Value) (Rows, error)
 }
 
+// QueryerContext is like Queryer, but most honor the context timeout and return
+// when the context is cancelled.
+type QueryerContext interface {
+	QueryContext(ctx context.Context, query string, args []Value) (Rows, error)
+}
+
 // Conn is a connection to a database. It is not used concurrently
 // by multiple goroutines.
 //
@@ -98,6 +113,23 @@ type Conn interface {
 	Begin() (Tx, error)
 }
 
+// ConnPrepareContext enhances the Conn interface with context.
+type ConnPrepareContext interface {
+	// PrepareContext returns a prepared statement, bound to this connection.
+	// context is for the preparation of the statement,
+	// it must not store the context within the statement itself.
+	PrepareContext(ctx context.Context, query string) (Stmt, error)
+}
+
+// ConnBeginContext enhances the Conn interface with context.
+type ConnBeginContext interface {
+	// BeginContext starts and returns a new transaction.
+	// the provided context should be used to roll the transaction back
+	// if it is cancelled. If there is an isolation level in context
+	// that is not supported by the driver an error must be returned.
+	BeginContext(ctx context.Context) (Tx, error)
+}
+
 // Result is the result of a query execution.
 type Result interface {
 	// LastInsertId returns the database's auto-generated ID
@@ -139,6 +171,18 @@ type Stmt interface {
 	Query(args []Value) (Rows, error)
 }
 
+// StmtExecContext enhances the Stmt interface by providing Exec with context.
+type StmtExecContext interface {
+	// ExecContext must honor the context timeout and return when it is cancelled.
+	ExecContext(ctx context.Context, args []Value) (Result, error)
+}
+
+// StmtQueryContext enhances the Stmt interface by providing Query with context.
+type StmtQueryContext interface {
+	// QueryContext must honor the context timeout and return when it is cancelled.
+	QueryContext(ctx context.Context, args []Value) (Rows, error)
+}
+
 // ColumnConverter may be optionally implemented by Stmt if the
 // statement is aware of its own columns' types and can convert from
 // any type to a driver Value.

src/database/sql/sql_test.go

@@ -5,6 +5,7 @@
 package sql
 
 import (
+	"context"
 	"database/sql/driver"
 	"errors"
 	"fmt"
@@ -1159,17 +1160,19 @@ func TestMaxOpenConnsOnBusy(t *testing.T) {
 
 	db.SetMaxOpenConns(3)
 
-	conn0, err := db.conn(cachedOrNewConn)
+	ctx := context.Background()
+
+	conn0, err := db.conn(ctx, cachedOrNewConn)
 	if err != nil {
 		t.Fatalf("db open conn fail: %v", err)
 	}
 
-	conn1, err := db.conn(cachedOrNewConn)
+	conn1, err := db.conn(ctx, cachedOrNewConn)
 	if err != nil {
 		t.Fatalf("db open conn fail: %v", err)
 	}
 
-	conn2, err := db.conn(cachedOrNewConn)
+	conn2, err := db.conn(ctx, cachedOrNewConn)
 	if err != nil {
 		t.Fatalf("db open conn fail: %v", err)
 	}

src/go/build/deps_test.go

@@ -228,8 +228,8 @@ var pkgDeps = map[string][]string{
 	"compress/lzw":             {"L4"},
 	"compress/zlib":            {"L4", "compress/flate"},
 	"context":                  {"errors", "fmt", "reflect", "sync", "time"},
-	"database/sql":             {"L4", "container/list", "database/sql/driver"},
-	"database/sql/driver":      {"L4", "time"},
+	"database/sql":             {"L4", "container/list", "context", "database/sql/driver"},
+	"database/sql/driver":      {"L4", "context", "time"},
 	"debug/dwarf":              {"L4"},
 	"debug/elf":                {"L4", "OS", "debug/dwarf", "compress/zlib"},
 	"debug/gosym":              {"L4"},