Merge branch 'context' into v3-experimental

query.go

@@ -4,6 +4,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"golang.org/x/net/context"
 	"time"
 )
 
@@ -55,7 +56,9 @@ func (rows *Rows) FieldDescriptions() []FieldDescription {
 	return rows.fields
 }
 
-func (rows *Rows) close() {
+// Close closes the rows, making the connection ready for use again. It is safe
+// to call Close after rows is already closed.
+func (rows *Rows) Close() {
 	if rows.closed {
 		return
 	}
@@ -67,6 +70,8 @@ func (rows *Rows) close() {
 
 	rows.closed = true
 
+	rows.err = rows.conn.termContext(rows.err)
+
 	if rows.err == nil {
 		if rows.conn.shouldLog(LogLevelInfo) {
 			endTime := time.Now()
@@ -81,63 +86,10 @@ func (rows *Rows) close() {
 	}
 }
 
-func (rows *Rows) readUntilReadyForQuery() {
-	for {
-		t, r, err := rows.conn.rxMsg()
-		if err != nil {
-			rows.close()
-			return
-		}
-
-		switch t {
-		case readyForQuery:
-			rows.conn.rxReadyForQuery(r)
-			rows.close()
-			return
-		case rowDescription:
-		case dataRow:
-		case commandComplete:
-		case bindComplete:
-		case errorResponse:
-			err = rows.conn.rxErrorResponse(r)
-			if rows.err == nil {
-				rows.err = err
-			}
-		default:
-			err = rows.conn.processContextFreeMsg(t, r)
-			if err != nil {
-				rows.close()
-				return
-			}
-		}
-	}
-}
-
-// Close closes the rows, making the connection ready for use again. It is safe
-// to call Close after rows is already closed.
-func (rows *Rows) Close() {
-	if rows.closed {
-		return
-	}
-	rows.readUntilReadyForQuery()
-	rows.close()
-}
-
 func (rows *Rows) Err() error {
 	return rows.err
 }
 
-// abort signals that the query was not successfully sent to the server.
-// This differs from Fatal in that it is not necessary to readUntilReadyForQuery
-func (rows *Rows) abort(err error) {
-	if rows.err != nil {
-		return
-	}
-
-	rows.err = err
-	rows.close()
-}
-
 // Fatal signals an error occurred after the query was sent to the server. It
 // closes the rows automatically.
 func (rows *Rows) Fatal(err error) {
@@ -169,10 +121,6 @@ func (rows *Rows) Next() bool {
 		}
 
 		switch t {
-		case readyForQuery:
-			rows.conn.rxReadyForQuery(r)
-			rows.close()
-			return false
 		case dataRow:
 			fieldCount := r.readInt16()
 			if int(fieldCount) != len(rows.fields) {
@@ -183,7 +131,9 @@ func (rows *Rows) Next() bool {
 			rows.mr = r
 			return true
 		case commandComplete:
-		case bindComplete:
+			rows.Close()
+			return false
+
 		default:
 			err = rows.conn.processContextFreeMsg(t, r)
 			if err != nil {
@@ -441,32 +391,7 @@ func (rows *Rows) AfterClose(f func(*Rows)) {
 // be returned in an error state. So it is allowed to ignore the error returned
 // from Query and handle it in *Rows.
 func (c *Conn) Query(sql string, args ...interface{}) (*Rows, error) {
-	c.lastActivityTime = time.Now()
-
-	rows := c.getRows(sql, args)
-
-	if err := c.lock(); err != nil {
-		rows.abort(err)
-		return rows, err
-	}
-	rows.unlockConn = true
-
-	ps, ok := c.preparedStatements[sql]
-	if !ok {
-		var err error
-		ps, err = c.Prepare("", sql)
-		if err != nil {
-			rows.abort(err)
-			return rows, rows.err
-		}
-	}
-	rows.sql = ps.SQL
-	rows.fields = ps.FieldDescriptions
-	err := c.sendPreparedQuery(ps, args...)
-	if err != nil {
-		rows.abort(err)
-	}
-	return rows, rows.err
+	return c.QueryContext(context.Background(), sql, args...)
 }
 
 func (c *Conn) getRows(sql string, args []interface{}) *Rows {
@@ -492,3 +417,51 @@ func (c *Conn) QueryRow(sql string, args ...interface{}) *Row {
 	rows, _ := c.Query(sql, args...)
 	return (*Row)(rows)
 }
+
+func (c *Conn) QueryContext(ctx context.Context, sql string, args ...interface{}) (rows *Rows, err error) {
+	err = c.waitForPreviousCancelQuery(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	c.lastActivityTime = time.Now()
+
+	rows = c.getRows(sql, args)
+
+	if err := c.lock(); err != nil {
+		rows.Fatal(err)
+		return rows, err
+	}
+	rows.unlockConn = true
+
+	ps, ok := c.preparedStatements[sql]
+	if !ok {
+		var err error
+		ps, err = c.PrepareExContext(ctx, "", sql, nil)
+		if err != nil {
+			rows.Fatal(err)
+			return rows, rows.err
+		}
+	}
+	rows.sql = ps.SQL
+	rows.fields = ps.FieldDescriptions
+
+	err = c.initContext(ctx)
+	if err != nil {
+		rows.Fatal(err)
+		return rows, err
+	}
+
+	err = c.sendPreparedQuery(ps, args...)
+	if err != nil {
+		rows.Fatal(err)
+		err = c.termContext(err)
+	}
+
+	return rows, err
+}
+
+func (c *Conn) QueryRowContext(ctx context.Context, sql string, args ...interface{}) *Row {
+	rows, _ := c.QueryContext(ctx, sql, args...)
+	return (*Row)(rows)
+}