Add support for multiple results to SQL statements.

Use the new functionality in lib/pq to select the next set of results
when multiple statements were executed.

Switch to using the "postgres" SQL driver in the cli tests.

Remove runPrettyQueryWithFormat. It was only used in tests.

Fixes cockroachdb#4016.

GLOCKFILE

@@ -20,7 +20,8 @@ github.com/cockroachdb/c-lz4 c40aaae2fc50293eb8750b34632bc3efe813e23f
 github.com/cockroachdb/c-protobuf 4feb192131ea08dfbd7253a00868ad69cbb61b81
 github.com/cockroachdb/c-rocksdb b7fb7bddcb55be35eacdf67e9e2c931083ce02c4
 github.com/cockroachdb/c-snappy 5c6d0932e0adaffce4bfca7bdf2ac37f79952ccf
-github.com/cockroachdb/stress 574c7f17016a4db745b88f6643700995110bdd07 
+github.com/cockroachdb/pq 77bd85500f4521560720328957bae47a78a90ed1
+github.com/cockroachdb/stress 574c7f17016a4db745b88f6643700995110bdd07
 github.com/cockroachdb/yacc 443154b1852a8702b07d675da6cd97cd9177a316
 github.com/codahale/hdrhistogram 954f16e8b9ef0e5d5189456aa4c1202758e04f17
 github.com/coreos/etcd 8199147cf859882f625523446c28ae2e53b2432f

acceptance/main_test.go

@@ -33,8 +33,6 @@ import (
 	"testing"
 	"time"
 
-	_ "github.com/lib/pq"
-
 	"github.com/samalba/dockerclient"
 
 	"github.com/cockroachdb/cockroach/acceptance/cluster"
@@ -46,6 +44,7 @@ import (
 	"github.com/cockroachdb/cockroach/util/log"
 	"github.com/cockroachdb/cockroach/util/randutil"
 	"github.com/cockroachdb/cockroach/util/stop"
+	_ "github.com/cockroachdb/pq"
 )
 
 var duration = flag.Duration("d", 5*time.Second, "duration to run the test")

cli/cli_test.go

@@ -541,6 +541,7 @@ func Example_sql() {
 	c.RunWithArgs([]string{"sql", "-e", "select * from t.f"})
 	c.RunWithArgs([]string{"sql", "-e", "show databases"})
 	c.RunWithArgs([]string{"sql", "-e", "explain select 3"})
+	c.RunWithArgs([]string{"sql", "-e", "select 1; select 2"})
 
 	// Output:
 	// sql -e create database t; create table t.f (x int, y int); insert into t.f values (42, 69)
@@ -571,6 +572,13 @@ func Example_sql() {
 	// 1 row
 	// Level	Type	Description
 	// 0	"empty"	"-"
+	// sql -e select 1; select 2
+	// 1 row
+	// 1
+	// 1
+	// 1 row
+	// 2
+	// 2
 }
 
 func Example_user() {

cli/flags.go

@@ -82,8 +82,8 @@ subsequent positional argument on the command line may contain
 one or more SQL statements, separated by semicolons. If an
 error occurs in any statement, the command exits with a
 non-zero status code and further statements are not
-executed. Only the results of the first SQL statement in each
-positional argument are printed on the standard output.`),
+executed. The results of each SQL statement are printed on
+the standard output.`),
 	"join": wrapText(`
 A comma-separated list of addresses to use when a new node is joining
 an existing cluster. For the first node in a cluster, --join should

cli/sql.go

@@ -24,6 +24,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/cockroachdb/pq"
 	"github.com/peterh/liner"
 	"github.com/spf13/cobra"
 )
@@ -137,26 +138,31 @@ func runInteractive(db *sql.DB, dbURL string) {
 // on error.
 func runStatements(db *sql.DB, stmts []string) {
 	for _, stmt := range stmts {
-		fullStmt := stmt + "\n"
-		cols, allRows, err := runQuery(db, fullStmt)
-		if err != nil {
-			fmt.Fprintln(osStderr, err)
-			os.Exit(1)
-		}
+		for {
+			cols, allRows, err := runQuery(db, stmt)
+			if err != nil {
+				if err == pq.ErrNoMoreResults {
+					break
+				}
+				fmt.Fprintln(osStderr, err)
+				os.Exit(1)
+			}
 
-		if len(cols) == 0 {
-			// No result selected, inform the user.
-			fmt.Fprintln(os.Stdout, "OK")
-		} else {
-			// Some results selected, inform the user about how much data to expect.
-			fmt.Fprintf(os.Stdout, "%d row%s\n", len(allRows),
-				map[bool]string{true: "", false: "s"}[len(allRows) == 1])
-
-			// Then print the results themselves.
-			fmt.Fprintln(os.Stdout, strings.Join(cols, "\t"))
-			for _, row := range allRows {
-				fmt.Fprintln(os.Stdout, strings.Join(row, "\t"))
+			if len(cols) == 0 {
+				// No result selected, inform the user.
+				fmt.Fprintln(os.Stdout, "OK")
+			} else {
+				// Some results selected, inform the user about how much data to expect.
+				fmt.Fprintf(os.Stdout, "%d row%s\n", len(allRows),
+					map[bool]string{true: "", false: "s"}[len(allRows) == 1])
+
+				// Then print the results themselves.
+				fmt.Fprintln(os.Stdout, strings.Join(cols, "\t"))
+				for _, row := range allRows {
+					fmt.Fprintln(os.Stdout, strings.Join(row, "\t"))
+				}
 			}
+			stmt = pq.NextResults
 		}
 	}
 }

cli/sql_util.go

@@ -22,7 +22,7 @@ import (
 	"io"
 
 	// Import postgres driver.
-	_ "github.com/lib/pq"
+	"github.com/cockroachdb/pq"
 
 	"github.com/olekukonko/tablewriter"
 
@@ -49,6 +49,9 @@ func makeSQLClient() (*sql.DB, string) {
 	if err != nil {
 		panicf("failed to initialize SQL client: %s", err)
 	}
+	// Ensure we use only one connection so that retrieval of multiple results
+	// can work without assumptions about connection reuse.
+	db.SetMaxOpenConns(1)
 	return db, sqlURL
 }
 
@@ -63,38 +66,36 @@ func runQuery(db *sql.DB, query string, parameters ...interface{}) (
 	return runQueryWithFormat(db, nil, query, parameters...)
 }
 
-// runQuery takes a 'query' with optional 'parameters'.
+// runQueryWithFormat takes a 'query' with optional 'parameters'.
 // It runs the sql query and returns a list of columns names and a list of rows.
 // If 'format' is not nil, the values with column name
 // found in the map are run through the corresponding callback.
 func runQueryWithFormat(db *sql.DB, format fmtMap, query string, parameters ...interface{}) (
 	[]string, [][]string, error) {
 	rows, err := db.Query(query, parameters...)
 	if err != nil {
-		return nil, nil, fmt.Errorf("query error: %s", err)
+		return nil, nil, err
 	}
 
 	defer rows.Close()
 	return sqlRowsToStrings(rows, format)
 }
 
-// runPrettyQueryWithFormat takes a 'query' with optional 'parameters'.
+// runPrettyQuery takes a 'query' with optional 'parameters'.
 // It runs the sql query and writes pretty output to 'w'.
 func runPrettyQuery(db *sql.DB, w io.Writer, query string, parameters ...interface{}) error {
-	return runPrettyQueryWithFormat(db, w, nil, query, parameters...)
-}
-
-// runPrettyQueryWithFormat takes a 'query' with optional 'parameters'.
-// It runs the sql query and writes pretty output to 'w'.
-// If 'format' is not nil, the values with column name
-// found in the map are run through the corresponding callback.
-func runPrettyQueryWithFormat(db *sql.DB, w io.Writer, format fmtMap, query string, parameters ...interface{}) error {
-	cols, allRows, err := runQueryWithFormat(db, format, query, parameters...)
-	if err != nil {
-		return err
+	for {
+		cols, allRows, err := runQueryWithFormat(db, nil, query, parameters...)
+		if err != nil {
+			if err == pq.ErrNoMoreResults {
+				return nil
+			}
+			return err
+		}
+		printQueryOutput(w, cols, allRows)
+		query = pq.NextResults
+		parameters = nil
 	}
-	printQueryOutput(w, cols, allRows)
-	return nil
 }
 
 // sqlRowsToStrings turns 'rows' into a list of rows, each of which
@@ -135,6 +136,9 @@ func sqlRowsToStrings(rows *sql.Rows, format fmtMap) ([]string, [][]string, erro
 		}
 		allRows = append(allRows, rowStrings)
 	}
+	if err := rows.Err(); err != nil {
+		return nil, nil, err
+	}
 
 	return cols, allRows, nil
 }

cli/sql_util_test.go

@@ -20,33 +20,33 @@ import (
 	"bytes"
 	"database/sql"
 	"fmt"
+	"os"
 	"reflect"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/security"
 	"github.com/cockroachdb/cockroach/server"
+	"github.com/cockroachdb/cockroach/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/util/leaktest"
 )
 
-func makeTestDBClient(t *testing.T, s *server.TestServer) *sql.DB {
-	db, err := sql.Open("cockroach", fmt.Sprintf("https://%s@%s?certs=%s",
-		security.RootUser,
-		s.ServingAddr(),
-		security.EmbeddedCertsDir))
-	if err != nil {
-		t.Fatal(err)
-	}
-	return db
-}
-
 func TestRunQuery(t *testing.T) {
 	defer leaktest.AfterTest(t)
 	s := server.StartTestServer(nil)
-	db := makeTestDBClient(t, s)
 
+	url, cleanup := sqlutils.PGUrl(t, s, security.RootUser, os.TempDir(), "TestRunQuery")
+	db, err := sql.Open("postgres", url.String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cleanup()
 	defer db.Close()
 	defer s.Stop()
 
+	// Ensure we use only one connection so that retrieval of multiple results
+	// can work without assumptions about connection reuse.
+	db.SetMaxOpenConns(1)
+
 	// Use a buffer as the io.Writer.
 	var b bytes.Buffer
 
@@ -75,9 +75,9 @@ OK
 	}
 
 	expectedRows := [][]string{
-		{`parentID`, `INT`, `false`, `NULL`},
-		{`name`, `STRING`, `false`, `NULL`},
-		{`id`, `INT`, `true`, `NULL`},
+		{`"parentID"`, `"INT"`, `false`, `NULL`},
+		{`"name"`, `"STRING"`, `false`, `NULL`},
+		{`"id"`, `"INT"`, `true`, `NULL`},
 	}
 	if !reflect.DeepEqual(expectedRows, rows) {
 		t.Fatalf("expected:\n%v\ngot:\n%v", expectedRows, rows)
@@ -88,13 +88,13 @@ OK
 	}
 
 	expected = `
-+----------+--------+-------+---------+
-|  Field   |  Type  | Null  | Default |
-+----------+--------+-------+---------+
-| parentID | INT    | false | NULL    |
-| name     | STRING | false | NULL    |
-| id       | INT    | true  | NULL    |
-+----------+--------+-------+---------+
++------------+----------+-------+---------+
+|   Field    |   Type   | Null  | Default |
++------------+----------+-------+---------+
+| "parentID" | "INT"    | false | NULL    |
+| "name"     | "STRING" | false | NULL    |
+| "id"       | "INT"    | true  | NULL    |
++------------+----------+-------+---------+
 `
 
 	if a, e := b.String(), expected[1:]; a != e {
@@ -108,11 +108,38 @@ OK
 	}
 
 	expected = `
-+----------+------------+----+
-| parentID |    name    | id |
-+----------+------------+----+
-|        1 | descriptor |  3 |
-+----------+------------+----+
++----------+--------------+----+
+| parentID |     name     | id |
++----------+--------------+----+
+|        1 | "descriptor" |  3 |
++----------+--------------+----+
+`
+	if a, e := b.String(), expected[1:]; a != e {
+		t.Fatalf("expected output:\n%s\ngot:\n%s", e, a)
+	}
+	b.Reset()
+
+	// Test multiple results.
+	if err := runPrettyQuery(db, &b, `SELECT 1; SELECT 2, 3; SELECT 'hello'`); err != nil {
+		t.Fatal(err)
+	}
+
+	expected = `
++---+
+| 1 |
++---+
+| 1 |
++---+
++---+---+
+| 2 | 3 |
++---+---+
+| 2 | 3 |
++---+---+
++---------+
+| 'hello' |
++---------+
+| "hello" |
++---------+
 `
 	if a, e := b.String(), expected[1:]; a != e {
 		t.Fatalf("expected output:\n%s\ngot:\n%s", e, a)
@@ -121,22 +148,17 @@ OK
 
 	// Test custom formatting.
 	newFormat := func(val interface{}) string {
-		return fmt.Sprintf("--> %#v <--", val)
+		return fmt.Sprintf("--> %s <--", val)
 	}
 
-	if err := runPrettyQueryWithFormat(db, &b, fmtMap{"name": newFormat},
-		`SELECT * FROM system.namespace WHERE name=$1`, "descriptor"); err != nil {
+	_, rows, err = runQueryWithFormat(db, fmtMap{"name": newFormat},
+		`SELECT * FROM system.namespace WHERE name=$1`, "descriptor")
+	if err != nil {
 		t.Fatal(err)
 	}
 
-	expected = `
-+----------+----------------------+----+
-| parentID |         name         | id |
-+----------+----------------------+----+
-|        1 | --> "descriptor" <-- |  3 |
-+----------+----------------------+----+
-`
-	if a, e := b.String(), expected[1:]; a != e {
+	expected = `[1 --> descriptor <-- 3]`
+	if a, e := fmt.Sprint(rows[0]), expected; a != e {
 		t.Fatalf("expected output:\n%s\ngot:\n%s", e, a)
 	}
 	b.Reset()

client/doc.go

@@ -16,7 +16,7 @@
 
 /*
 Package client and its KV API has been deprecated for external usage. Please use
-a postgres-compatible SQL driver (e.g. github.com/lib/pq). For more details, see
+a postgres-compatible SQL driver (e.g. github.com/cockroachdb/pq). For more details, see
 http://www.cockroachlabs.com/blog/sql-in-cockroachdb-mapping-table-data-to-key-value-storage/.
 
 Package client provides clients for accessing the various externally-facing

sql/bench_test.go

@@ -25,12 +25,12 @@ import (
 	"testing"
 
 	_ "github.com/go-sql-driver/mysql"
-	_ "github.com/lib/pq"
 
 	"github.com/cockroachdb/cockroach/security"
 	"github.com/cockroachdb/cockroach/server"
 	"github.com/cockroachdb/cockroach/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/util/tracing"
+	_ "github.com/cockroachdb/pq"
 )
 
 func benchmarkCockroach(b *testing.B, f func(b *testing.B, db *sql.DB)) {

sql/driver/wire.go

@@ -96,7 +96,7 @@ func (d Datum) Value() (driver.Value, error) {
 		val = t.FloatVal
 	case *Datum_DecimalVal:
 		// For now, we just return the decimal string, to be consistent
-		// with lib/pq's driver.
+		// with cockroachdb/pq's driver.
 		val = t.DecimalVal
 	case *Datum_BytesVal:
 		val = t.BytesVal