correctly close rows objects if there are multiple result sets

Author: maddyblue

conn.go

@@ -1339,7 +1339,12 @@ func (rs *rows) Close() error {
 		switch err {
 		case nil:
 		case io.EOF:
-			return nil
+			// rs.Next can return io.EOF on both 'Z' (ready for query) and 'T' (row
+			// description, used with HasNextResultSet). We need to fetch messages until
+			// we hit a 'Z', which is done by waiting for done to be set.
+			if rs.done {
+				return nil
+			}
 		default:
 			return err
 		}

conn_test.go

@@ -1583,3 +1583,32 @@ func TestRowsResultTag(t *testing.T) {
 		}
 	}
 }
+
+// TestQuickClose tests that closing a query early allows a subsequent query to work.
+func TestQuickClose(t *testing.T) {
+	db := openTestConn(t)
+	defer db.Close()
+
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+	rows, err := tx.Query("SELECT 1; SELECT 2;")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := rows.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	var id int
+	if err := tx.QueryRow("SELECT 3").Scan(&id); err != nil {
+		t.Fatal(err)
+	}
+	if id != 3 {
+		t.Fatalf("unexpected %d", id)
+	}
+	if err := tx.Commit(); err != nil {
+		t.Fatal(err)
+	}
+}