Merge pull request #535 from mjibson/ctx

maddyblue · web-flow · commit a6657b2386e9 · 2017-01-31T20:15:56.000-05:00
Add context methods
diff --git a/.travis.yml b/.travis.yml
@@ -1,10 +1,11 @@
 language: go
 
 go:
-  - 1.5
-  - 1.6
-  - 1.7
-  - tip
+  - 1.5.x
+  - 1.6.x
+  - 1.7.x
+  - 1.8.x
+  - master
 
 sudo: true
 
diff --git a/conn.go b/conn.go
@@ -98,6 +98,15 @@ type conn struct {
 	namei     int
 	scratch   [512]byte
 	txnStatus transactionStatus
+	txnClosed chan<- struct{}
+
+	// Save connection arguments to use during CancelRequest.
+	dialer Dialer
+	opts   values
+
+	// Cancellation key data for use with CancelRequest messages.
+	processID int
+	secretKey int
 
 	parameterStatus parameterStatus
 
@@ -307,7 +316,10 @@ func DialOpen(d Dialer, name string) (_ driver.Conn, err error) {
 		}
 	}
 
-	cn := &conn{}
+	cn := &conn{
+		opts:   o,
+		dialer: d,
+	}
 	err = cn.handleDriverSettings(o)
 	if err != nil {
 		return nil, err
@@ -529,7 +541,15 @@ func (cn *conn) Begin() (_ driver.Tx, err error) {
 	return cn, nil
 }
 
+func (cn *conn) closeTxn() {
+	if cn.txnClosed != nil {
+		close(cn.txnClosed)
+		cn.txnClosed = nil
+	}
+}
+
 func (cn *conn) Commit() (err error) {
+	defer cn.closeTxn()
 	if cn.bad {
 		return driver.ErrBadConn
 	}
@@ -565,6 +585,7 @@ func (cn *conn) Commit() (err error) {
 }
 
 func (cn *conn) Rollback() (err error) {
+	defer cn.closeTxn()
 	if cn.bad {
 		return driver.ErrBadConn
 	}
@@ -796,7 +817,11 @@ func (cn *conn) Close() (err error) {
 }
 
 // Implement the "Queryer" interface
-func (cn *conn) Query(query string, args []driver.Value) (_ driver.Rows, err error) {
+func (cn *conn) Query(query string, args []driver.Value) (driver.Rows, error) {
+	return cn.query(query, args)
+}
+
+func (cn *conn) query(query string, args []driver.Value) (_ *rows, err error) {
 	if cn.bad {
 		return nil, driver.ErrBadConn
 	}
@@ -1074,6 +1099,7 @@ func (cn *conn) startup(o values) {
 		t, r := cn.recv()
 		switch t {
 		case 'K':
+			cn.processBackendKeyData(r)
 		case 'S':
 			cn.processParameterStatus(r)
 		case 'R':
@@ -1301,6 +1327,7 @@ func (cn *conn) parseComplete(commandTag string) (driver.Result, string) {
 
 type rows struct {
 	cn       *conn
+	closed   chan<- struct{}
 	colNames []string
 	colTyps  []oid.Oid
 	colFmts  []format
@@ -1309,6 +1336,9 @@ type rows struct {
 }
 
 func (rs *rows) Close() error {
+	if rs.closed != nil {
+		defer close(rs.closed)
+	}
 	// no need to look at cn.bad as Next() will
 	for {
 		err := rs.Next(nil)
@@ -1513,6 +1543,11 @@ func (cn *conn) readReadyForQuery() {
 	}
 }
 
+func (c *conn) processBackendKeyData(r *readBuf) {
+	c.processID = r.int32()
+	c.secretKey = r.int32()
+}
+
 func (cn *conn) readParseResponse() {
 	t, r := cn.recv1()
 	switch t {
diff --git a/conn_go18.go b/conn_go18.go
@@ -0,0 +1,92 @@
+// +build go1.8
+
+package pq
+
+import (
+	"context"
+	"database/sql/driver"
+	"errors"
+)
+
+// 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))
+	for i, nv := range args {
+		list[i] = nv.Value
+	}
+	var closed chan<- struct{}
+	if ctx.Done() != nil {
+		closed = watchCancel(ctx, cn.cancel)
+	}
+	r, err := cn.query(query, list)
+	if err != nil {
+		return nil, err
+	}
+	r.closed = closed
+	return r, nil
+}
+
+// Implement the "ExecerContext" interface
+func (cn *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+	list := make([]driver.Value, len(args))
+	for i, nv := range args {
+		list[i] = nv.Value
+	}
+
+	if ctx.Done() != nil {
+		closed := watchCancel(ctx, cn.cancel)
+		defer close(closed)
+	}
+
+	return cn.Exec(query, list)
+}
+
+// Implement the "ConnBeginTx" interface
+func (cn *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	if opts.Isolation != 0 {
+		return nil, errors.New("isolation levels not supported")
+	}
+	if opts.ReadOnly {
+		return nil, errors.New("read-only transactions not supported")
+	}
+	tx, err := cn.Begin()
+	if err != nil {
+		return nil, err
+	}
+	if ctx.Done() != nil {
+		cn.txnClosed = watchCancel(ctx, cn.cancel)
+	}
+	return tx, nil
+}
+
+func watchCancel(ctx context.Context, cancel func()) chan<- struct{} {
+	closed := make(chan struct{})
+	go func() {
+		select {
+		case <-ctx.Done():
+			cancel()
+		case <-closed:
+		}
+	}()
+	return closed
+}
+
+func (cn *conn) cancel() {
+	var err error
+	can := &conn{}
+	can.c, err = dial(cn.dialer, cn.opts)
+	if err != nil {
+		return
+	}
+	can.ssl(cn.opts)
+
+	defer can.errRecover(&err)
+
+	w := can.writeBuf(0)
+	w.int32(80877102) // cancel request code
+	w.int32(cn.processID)
+	w.int32(cn.secretKey)
+
+	can.sendStartupPacket(w)
+	_ = can.c.Close()
+}
diff --git a/go18_test.go b/go18_test.go
@@ -2,7 +2,12 @@
 
 package pq
 
-import "testing"
+import (
+	"context"
+	"database/sql"
+	"testing"
+	"time"
+)
 
 func TestMultipleSimpleQuery(t *testing.T) {
 	db := openTestConn(t)
@@ -66,3 +71,95 @@ func TestMultipleSimpleQuery(t *testing.T) {
 		t.Fatal("unexpected result set")
 	}
 }
+
+func TestContextCancelExec(t *testing.T) {
+	db := openTestConn(t)
+	defer db.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// Delay execution for just a bit until db.ExecContext has begun.
+	go func() {
+		time.Sleep(time.Millisecond * 10)
+		cancel()
+	}()
+
+	// Not canceled until after the exec has started.
+	if _, err := db.ExecContext(ctx, "select pg_sleep(1)"); err == nil {
+		t.Fatal("expected error")
+	} else if err.Error() != "pq: canceling statement due to user request" {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	// Context is already canceled, so error should come before execution.
+	if _, err := db.ExecContext(ctx, "select pg_sleep(1)"); err == nil {
+		t.Fatal("expected error")
+	} else if err.Error() != "context canceled" {
+		t.Fatalf("unexpected error: %s", err)
+	}
+}
+
+func TestContextCancelQuery(t *testing.T) {
+	db := openTestConn(t)
+	defer db.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// Delay execution for just a bit until db.QueryContext has begun.
+	go func() {
+		time.Sleep(time.Millisecond * 10)
+		cancel()
+	}()
+
+	// Not canceled until after the exec has started.
+	if _, err := db.QueryContext(ctx, "select pg_sleep(1)"); err == nil {
+		t.Fatal("expected error")
+	} else if err.Error() != "pq: canceling statement due to user request" {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	// Context is already canceled, so error should come before execution.
+	if _, err := db.QueryContext(ctx, "select pg_sleep(1)"); err == nil {
+		t.Fatal("expected error")
+	} else if err.Error() != "context canceled" {
+		t.Fatalf("unexpected error: %s", err)
+	}
+}
+
+func TestContextCancelBegin(t *testing.T) {
+	db := openTestConn(t)
+	defer db.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Delay execution for just a bit until tx.Exec has begun.
+	go func() {
+		time.Sleep(time.Millisecond * 10)
+		cancel()
+	}()
+
+	// Not canceled until after the exec has started.
+	if _, err := tx.Exec("select pg_sleep(1)"); err == nil {
+		t.Fatal("expected error")
+	} else if err.Error() != "pq: canceling statement due to user request" {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	// Transaction is canceled, so expect an error.
+	if _, err := tx.Query("select pg_sleep(1)"); err == nil {
+		t.Fatal("expected error")
+	} else if err != sql.ErrTxDone {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	// Context is canceled, so cannot begin a transaction.
+	if _, err := db.BeginTx(ctx, nil); err == nil {
+		t.Fatal("expected error")
+	} else if err.Error() != "context canceled" {
+		t.Fatalf("unexpected error: %s", err)
+	}
+}

-Original file line number
+Diff line change
@@ @@ -1,10 +1,11 @@ @@
 language: go
 go:
 -  - 1.5
 -  - 1.6
 -  - 1.7
 -  - tip
 +  - 1.5.x
 +  - 1.6.x
 +  - 1.7.x
 +  - 1.8.x
 +  - master
 sudo: true