Add support for retrieving multiple results.

If the special `NEXT` query is executed, the driver looks for another
set of results on the connection. If no such results are found,
ErrNoMoreResults is returned.

Author: petermattis

conn.go

@@ -32,8 +32,11 @@ var (
 	ErrSSLNotSupported           = errors.New("pq: SSL is not enabled on the server")
 	ErrSSLKeyHasWorldPermissions = errors.New("pq: Private key file has group or world access. Permissions should be u=rw (0600) or less.")
 	ErrCouldNotDetectUsername    = errors.New("pq: Could not detect default username. Please provide one explicitly.")
+	ErrNoMoreResults             = errors.New("pq: no more results")
 )
 
+const NextResults = "NEXT"
+
 type drv struct{}
 
 func (d *drv) Open(name string) (driver.Conn, error) {
@@ -619,56 +622,73 @@ func (cn *conn) simpleExec(q string) (res driver.Result, commandTag string, err
 func (cn *conn) simpleQuery(q string) (res *rows, err error) {
 	defer cn.errRecover(&err)
 
-	cn.waitReadyForQuery()
-
-	// Mark the connection has having sent a query.
-	cn.readyForQuery = false
-	b := cn.writeBuf('Q')
-	b.string(q)
-	cn.send(b)
+	querySent := false
+	nextResult := q == NextResults
 
 	for {
+		if cn.readyForQuery && !querySent {
+			if nextResult {
+				return nil, ErrNoMoreResults
+			}
+
+			// Mark the connection has having sent a query.
+			cn.readyForQuery = false
+			b := cn.writeBuf('Q')
+			b.string(q)
+			cn.send(b)
+			querySent = true
+		}
 
 		t, r := cn.recv1()
 		switch t {
 		case 'C', 'I':
-			// We allow queries which don't return any results through Query as
-			// well as Exec.  We still have to give database/sql a rows object
-			// the user can close, though, to avoid connections from being
-			// leaked.  A "rows" with done=true works fine for that purpose.
-			if err != nil {
-				cn.bad = true
-				errorf("unexpected message %q in simple query execution", t)
-			}
-			if res == nil {
-				res = &rows{
-					cn: cn,
+			if nextResult || querySent {
+				// We allow queries which don't return any results through Query as
+				// well as Exec.  We still have to give database/sql a rows object
+				// the user can close, though, to avoid connections from being
+				// leaked.  A "rows" with done=true works fine for that purpose.
+				if err != nil {
+					cn.bad = true
+					errorf("unexpected message %q in simple query execution", t)
 				}
+				if res == nil {
+					res = &rows{
+						cn: cn,
+					}
+				}
+				res.done = true
 			}
-			res.done = true
 		case 'Z':
 			cn.processReadyForQuery(r)
-			// done
-			return
+			if querySent {
+				// done
+				return
+			}
 		case 'E':
-			res = nil
-			err = parseError(r)
+			if nextResult || querySent {
+				res = nil
+				err = parseError(r)
+			}
 		case 'D':
-			if res == nil {
-				cn.bad = true
-				errorf("unexpected DataRow in simple query execution")
+			if nextResult || querySent {
+				if res == nil {
+					cn.bad = true
+					errorf("unexpected DataRow in simple query execution")
+				}
+				// the query didn't fail; kick off to Next
+				cn.saveMessage(t, r)
+				return
 			}
-			// the query didn't fail; kick off to Next
-			cn.saveMessage(t, r)
-			return
 		case 'T':
-			// res might be non-nil here if we received a previous
-			// CommandComplete, but that's fine; just overwrite it
-			res = &rows{cn: cn}
-			res.colNames, res.colFmts, res.colTyps = parsePortalRowDescribe(r)
-
-			// To work around a bug in QueryRow in Go 1.2 and earlier, wait
-			// until the first DataRow has been received.
+			if nextResult || querySent {
+				// res might be non-nil here if we received a previous
+				// CommandComplete, but that's fine; just overwrite it
+				res = &rows{cn: cn}
+				res.colNames, res.colFmts, res.colTyps = parsePortalRowDescribe(r)
+
+				// To work around a bug in QueryRow in Go 1.2 and earlier, wait
+				// until the first DataRow has been received.
+			}
 		default:
 			cn.bad = true
 			errorf("unknown response for simple query: %q", t)

conn_test.go

@@ -136,6 +136,7 @@ func TestOpenURL(t *testing.T) {
 }
 
 const pgpass_file = "/tmp/pqgotest_pgpass"
+
 func TestPgpass(t *testing.T) {
 	testAssert := func(conninfo string, expected string, reason string) {
 		conn, err := openTestConnConninfo(conninfo)
@@ -339,6 +340,70 @@ func TestRowsCloseBeforeDone(t *testing.T) {
 	}
 }
 
+func TestMultipleResults(t *testing.T) {
+	db := openTestConn(t)
+	defer db.Close()
+
+	var val int
+	if err := db.QueryRow("SELECT 1; SELECT 2").Scan(&val); err != nil {
+		t.Fatal(err)
+	}
+	if val != 1 {
+		t.Fatalf("expected 1, but found %d", val)
+	}
+	if err := db.QueryRow(NextResults).Scan(&val); err != nil {
+		t.Fatal(err)
+	}
+	if val != 2 {
+		t.Fatalf("expected 2, but found %d", val)
+	}
+	if err := db.QueryRow(NextResults).Scan(&val); err != ErrNoMoreResults {
+		t.Fatalf("expected %s, but found %v", ErrNoMoreResults, err)
+	}
+
+	// Now test discarding the second result.
+	if err := db.QueryRow("SELECT 3; SELECT 4").Scan(&val); err != nil {
+		t.Fatal(err)
+	}
+	if val != 3 {
+		t.Fatalf("expected 3, but found %d", val)
+	}
+	if err := db.QueryRow("SELECT 5").Scan(&val); err != nil {
+		t.Fatal(err)
+	}
+	if val != 5 {
+		t.Fatalf("expected 5, but found %d", val)
+	}
+}
+
+func TestTxnMultipleResults(t *testing.T) {
+	db := openTestConn(t)
+	defer db.Close()
+
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tx.Rollback()
+
+	var val int
+	if err := tx.QueryRow("SELECT 1; SELECT 2").Scan(&val); err != nil {
+		t.Fatal(err)
+	}
+	if val != 1 {
+		t.Fatalf("expected 1, but found %d", val)
+	}
+	if err := tx.QueryRow(NextResults).Scan(&val); err != nil {
+		t.Fatal(err)
+	}
+	if val != 2 {
+		t.Fatalf("expected 2, but found %d", val)
+	}
+	if err := tx.QueryRow(NextResults).Scan(&val); err != ErrNoMoreResults {
+		t.Fatalf("expected %s, but found %v", ErrNoMoreResults, err)
+	}
+}
+
 func TestParameterCountMismatch(t *testing.T) {
 	db := openTestConn(t)
 	defer db.Close()