database/sql: fix Tx.Query

Fixes #2784

R=golang-dev, bradfitz
CC=golang-dev
https://golang.org/cl/5574073

2 files changed, 39 insertions, 2 deletions

diff --git a/src/pkg/database/sql/sql.go b/src/pkg/database/sql/sql.go
index 70499b9a95..7e226b17dc 100644
--- a/src/pkg/database/sql/sql.go
+++ b/src/pkg/database/sql/sql.go
@@ -556,8 +556,11 @@ func (tx *Tx) Query(query string, args ...interface{}) (*Rows, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer stmt.Close()
-	return stmt.Query(args...)
+	rows, err := stmt.Query(args...)
+	if err == nil {
+		rows.closeStmt = stmt
+	}
+	return rows, err
 }
 
 // QueryRow executes a query that is expected to return at most one row.
diff --git a/src/pkg/database/sql/sql_test.go b/src/pkg/database/sql/sql_test.go
index 3fb137eb24..08db6d38ff 100644
--- a/src/pkg/database/sql/sql_test.go
+++ b/src/pkg/database/sql/sql_test.go
@@ -311,6 +311,40 @@ func TestTxStmt(t *testing.T) {
 	}
 }
 
+// Issue: http://golang.org/issue/2784
+// This test didn't fail before because we got luckly with the fakedb driver.
+// It was failing, and now not, in github.com/bradfitz/go-sql-test
+func TestTxQuery(t *testing.T) {
+	db := newTestDB(t, "")
+	defer closeDB(t, db)
+	exec(t, db, "CREATE|t1|name=string,age=int32,dead=bool")
+	exec(t, db, "INSERT|t1|name=Alice")
+
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tx.Rollback()
+
+	r, err := tx.Query("SELECT|t1|name|")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !r.Next() {
+		if r.Err() != nil {
+			t.Fatal(r.Err())
+		}
+		t.Fatal("expected one row")
+	}
+
+	var x string
+	err = r.Scan(&x)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
 // Tests fix for issue 2542, that we release a lock when querying on
 // a closed connection.
 func TestIssue2542Deadlock(t *testing.T) {