Bump version to 5.11.0

vendor/github.com/lib/pq/conn_go18.go

@@ -7,10 +7,13 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
-	"sync/atomic"
 	"time"
 )
 
+const (
+	watchCancelDialContextTimeout = time.Second * 10
+)
+
 // Implement the "QueryerContext" interface
 func (cn *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
 	list := make([]driver.Value, len(args))
@@ -43,6 +46,14 @@ func (cn *conn) ExecContext(ctx context.Context, query string, args []driver.Nam
 	return cn.Exec(query, list)
 }
 
+// Implement the "ConnPrepareContext" interface
+func (cn *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+	if finish := cn.watchCancel(ctx); finish != nil {
+		defer finish()
+	}
+	return cn.Prepare(query)
+}
+
 // Implement the "ConnBeginTx" interface
 func (cn *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
 	var mode string
@@ -103,13 +114,13 @@ func (cn *conn) watchCancel(ctx context.Context) func() {
 				}
 
 				// Set the connection state to bad so it does not get reused.
-				cn.setBad()
+				cn.err.set(ctx.Err())
 
 				// At this point the function level context is canceled,
 				// so it must not be used for the additional network
 				// request to cancel the query.
 				// Create a new context to pass into the dial.
-				ctxCancel, cancel := context.WithTimeout(context.Background(), time.Second*10)
+				ctxCancel, cancel := context.WithTimeout(context.Background(), watchCancelDialContextTimeout)
 				defer cancel()
 
 				_ = cn.cancel(ctxCancel)
@@ -119,7 +130,7 @@ func (cn *conn) watchCancel(ctx context.Context) func() {
 		return func() {
 			select {
 			case <-finished:
-				cn.setBad()
+				cn.err.set(ctx.Err())
 				cn.Close()
 			case finished <- struct{}{}:
 			}
@@ -145,11 +156,8 @@ func (cn *conn) cancel(ctx context.Context) error {
 	defer c.Close()
 
 	{
-		bad := &atomic.Value{}
-		bad.Store(false)
 		can := conn{
-			c:   c,
-			bad: bad,
+			c: c,
 		}
 		err = can.ssl(o)
 		if err != nil {
@@ -172,3 +180,68 @@ func (cn *conn) cancel(ctx context.Context) error {
 		return err
 	}
 }
+
+// Implement the "StmtQueryContext" interface
+func (st *stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
+	list := make([]driver.Value, len(args))
+	for i, nv := range args {
+		list[i] = nv.Value
+	}
+	finish := st.watchCancel(ctx)
+	r, err := st.query(list)
+	if err != nil {
+		if finish != nil {
+			finish()
+		}
+		return nil, err
+	}
+	r.finish = finish
+	return r, nil
+}
+
+// Implement the "StmtExecContext" interface
+func (st *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
+	list := make([]driver.Value, len(args))
+	for i, nv := range args {
+		list[i] = nv.Value
+	}
+
+	if finish := st.watchCancel(ctx); finish != nil {
+		defer finish()
+	}
+
+	return st.Exec(list)
+}
+
+// watchCancel is implemented on stmt in order to not mark the parent conn as bad
+func (st *stmt) watchCancel(ctx context.Context) func() {
+	if done := ctx.Done(); done != nil {
+		finished := make(chan struct{})
+		go func() {
+			select {
+			case <-done:
+				// At this point the function level context is canceled,
+				// so it must not be used for the additional network
+				// request to cancel the query.
+				// Create a new context to pass into the dial.
+				ctxCancel, cancel := context.WithTimeout(context.Background(), watchCancelDialContextTimeout)
+				defer cancel()
+
+				_ = st.cancel(ctxCancel)
+				finished <- struct{}{}
+			case <-finished:
+			}
+		}()
+		return func() {
+			select {
+			case <-finished:
+			case finished <- struct{}{}:
+			}
+		}
+	}
+	return nil
+}
+
+func (st *stmt) cancel(ctx context.Context) error {
+	return st.cn.cancel(ctx)
+}