database/sql: Use Tx.ctx in Tx non-context methods

The Tx methods Query and Exec uses context.Background()
even Tx was created by context.
This patch enables using Tx.ctx in all Tx methods
which do not has context arg.
Backward compatibility:
- If Tx has created without context, nothing changes.
- If Tx has created with context and non-context method is called:
  - If context is expired, the execution fails,
    but it can fail on Commit or Rollback as well,
    so in terms of whole transaction - nothing changes.
  - If context is not expired, nothing changes too.

Fixes #20098
Change-Id: I9570a2deaace5875bb4c5dcf7b3a084a6bcd0d00
Reviewed-on: https://go-review.googlesource.com/44956
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Reviewed-by: Daniel Theophanes <kardianos@gmail.com>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

1 files changed, 29 insertions, 0 deletions

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 8a477edf1a..06877a6081 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -439,6 +439,35 @@ func TestTxContextWait(t *testing.T) {
 	waitForFree(t, db, 5*time.Second, 0)
 }
 
+// TestTxUsesContext tests the transaction behavior when the tx was created by context,
+// but for query execution used methods without context
+func TestTxUsesContext(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Millisecond)
+	defer cancel()
+
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		// Guard against the context being canceled before BeginTx completes.
+		if err == context.DeadlineExceeded {
+			t.Skip("tx context canceled prior to first use")
+		}
+		t.Fatal(err)
+	}
+
+	// This will trigger the *fakeConn.Prepare method which will take time
+	// performing the query. The ctxDriverPrepare func will check the context
+	// after this and close the rows and return an error.
+	_, err = tx.Query("WAIT|1s|SELECT|people|age,name|")
+	if err != context.DeadlineExceeded {
+		t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
+	}
+
+	waitForFree(t, db, 5*time.Second, 0)
+}
+
 func TestMultiResultSetQuery(t *testing.T) {
 	db := newTestDB(t, "people")
 	defer closeDB(t, db)