cli: 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.

Fixes cockroachdb#4016.

GLOCKFILE

@@ -25,6 +25,7 @@ github.com/cockroachdb/c-lz4 c40aaae2fc50293eb8750b34632bc3efe813e23f
 github.com/cockroachdb/c-protobuf 4feb192131ea08dfbd7253a00868ad69cbb61b81
 github.com/cockroachdb/c-rocksdb c0124c907c74b579d9d3d48eb96471bef270bc25
 github.com/cockroachdb/c-snappy 5c6d0932e0adaffce4bfca7bdf2ac37f79952ccf
+github.com/cockroachdb/pq 77893094b774b29f293681e6ac0a9322fbf3ce25
 github.com/cockroachdb/stress aa7690c22fd0abd6168ed0e6c361e4f4c5f7ab25
 github.com/codahale/hdrhistogram e88be87d51429689cef99043a54150d733265cd7
 github.com/coreos/etcd c3824b10daa09fe99fba66a329db881adeb03aae

acceptance/main_test.go

@@ -36,7 +36,6 @@ import (
 	"github.com/docker/engine-api/types"
 	"github.com/docker/engine-api/types/container"
 	"github.com/docker/engine-api/types/strslice"
-	_ "github.com/lib/pq"
 
 	"github.com/cockroachdb/cockroach/acceptance/cluster"
 	"github.com/cockroachdb/cockroach/acceptance/terrafarm"
@@ -46,6 +45,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 flagDuration = flag.Duration("d", cluster.DefaultDuration, "duration to run the test")

cli/cli_test.go

@@ -543,6 +543,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)
@@ -573,6 +574,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_sql_escape() {

cli/flags.go

@@ -69,9 +69,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

@@ -29,6 +29,8 @@ import (
 	"github.com/cockroachdb/cockroach/util/log"
 	"github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
+
+	"github.com/cockroachdb/pq"
 )
 
 const (
@@ -225,30 +227,35 @@ func runInteractive(conn *sqlConn) (exitErr error) {
 // on error.
 func runStatements(conn *sqlConn, stmts []string) error {
 	for _, stmt := range stmts {
-		fullStmt := stmt + "\n"
-		cols, allRows, err := runQuery(conn, fullStmt)
-		if err != nil {
-			fmt.Fprintln(osStderr, err)
-			return err
-		}
-
-		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.
-			noun := "rows"
-			if len(allRows) == 1 {
-				noun = "row"
+		for {
+			cols, allRows, err := runQuery(conn, stmt)
+			if err != nil {
+				if err == pq.ErrNoMoreResults {
+					break
+				}
+				fmt.Fprintln(osStderr, err)
+				os.Exit(1)
 			}
 
-			fmt.Fprintf(os.Stdout, "%d %s\n", len(allRows), noun)
-
-			// 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.
+				noun := "rows"
+				if len(allRows) == 1 {
+					noun = "row"
+				}
+
+				fmt.Fprintf(os.Stdout, "%d %s\n", len(allRows), noun)
+
+				// 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 = nextResult
 		}
 	}
 	return nil

cli/sql_util.go

@@ -23,16 +23,21 @@ import (
 	"io"
 	"net"
 
-	"github.com/lib/pq"
-
 	"github.com/olekukonko/tablewriter"
 
 	"github.com/cockroachdb/cockroach/util/log"
+	"github.com/cockroachdb/pq"
 )
 
+// nextResult is a special SQL query which indicates we want the next result
+// for a multi-statement query. Instead of sending a query, we invoke
+// sqlConn.Next().
+const nextResult = `\next`
+
 type sqlConnI interface {
 	driver.Conn
 	driver.Queryer
+	Next() (driver.Rows, error)
 }
 
 type sqlConn struct {
@@ -65,6 +70,20 @@ func (c *sqlConn) Query(query string, args []driver.Value) (*sqlRows, error) {
 	return &sqlRows{Rows: rows, conn: c}, nil
 }
 
+func (c *sqlConn) Next() (*sqlRows, error) {
+	if c.conn == nil {
+		return nil, driver.ErrBadConn
+	}
+	rows, err := c.conn.Next()
+	if err == driver.ErrBadConn {
+		c.Close()
+	}
+	if err != nil {
+		return nil, err
+	}
+	return &sqlRows{Rows: rows, conn: c}, nil
+}
+
 func (c *sqlConn) Close() {
 	if c.conn != nil {
 		err := c.conn.Close()
@@ -118,16 +137,16 @@ type fmtMap map[string]func(driver.Value) string
 
 // runQuery takes a 'query' with optional 'parameters'.
 // It runs the sql query and returns a list of columns names and a list of rows.
-func runQuery(db *sqlConn, query string, parameters ...driver.Value) (
+func runQuery(conn *sqlConn, query string, parameters ...driver.Value) (
 	[]string, [][]string, error) {
-	return runQueryWithFormat(db, nil, query, parameters...)
+	return runQueryWithFormat(conn, 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 *sqlConn, format fmtMap, query string, parameters ...driver.Value) (
+func runQueryWithFormat(conn *sqlConn, format fmtMap, query string, parameters ...driver.Value) (
 	[]string, [][]string, error) {
 	// driver.Value is an alias for interface{}, but must adhere to a restricted
 	// set of types when being passed to driver.Queryer.Query (see
@@ -143,24 +162,36 @@ func runQueryWithFormat(db *sqlConn, format fmtMap, query string, parameters ...
 		}
 	}
 
-	rows, err := db.Query(query, parameters)
+	var rows *sqlRows
+	var err error
+	if query == nextResult {
+		rows, err = conn.Next()
+	} else {
+		rows, err = conn.Query(query, parameters)
+	}
 	if err != nil {
-		return nil, nil, fmt.Errorf("query error: %s", err)
+		return nil, nil, err
 	}
 
 	defer func() { _ = 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 *sqlConn, w io.Writer, query string, parameters ...driver.Value) error {
-	cols, allRows, err := runQuery(db, query, parameters...)
-	if err != nil {
-		return err
+func runPrettyQuery(conn *sqlConn, w io.Writer, query string, parameters ...driver.Value) error {
+	for {
+		cols, allRows, err := runQuery(conn, query, parameters...)
+		if err != nil {
+			if err == pq.ErrNoMoreResults {
+				return nil
+			}
+			return err
+		}
+		printQueryOutput(w, cols, allRows)
+		query = nextResult
+		parameters = nil
 	}
-	printQueryOutput(w, cols, allRows)
-	return nil
 }
 
 // sqlRowsToStrings turns 'rows' into a list of rows, each of which

cli/sql_util_test.go

@@ -129,38 +129,31 @@ OK
 	}
 	b.Reset()
 
-	// TODO(pmattis): This test case fails now as lib/pq doesn't handle multiple
-	// results correctly. We were previously incorrectly ignoring the error from
-	// sql.Rows.Err() which is what allowed the test to pass.
-
-	/**
-		// Test multiple results.
-		if err := runPrettyQuery(conn, &b, `SELECT 1; SELECT 2, 3; SELECT 'hello'`); err != nil {
-			t.Fatal(err)
-		}
-
-		expected = `
-	+---+
-	| 1 |
-	+---+
-	| 1 |
-	+---+
-	`
-		// TODO(pmattis): When #4016 is fixed, we should see:
-		// +---+---+
-		// | 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)
-		}
-		b.Reset()
-	**/
+	// Test multiple results.
+	if err := runPrettyQuery(conn, &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)
+	}
+	b.Reset()
 }

sql/bench_test.go

@@ -24,12 +24,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

sql/pgwire/types.go

@@ -25,11 +25,10 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/lib/pq/oid"
-
 	"github.com/cockroachdb/cockroach/sql/parser"
 	"github.com/cockroachdb/cockroach/util"
 	"github.com/cockroachdb/cockroach/util/log"
+	"github.com/cockroachdb/pq/oid"
 )
 
 //go:generate stringer -type=formatCode
@@ -201,8 +200,8 @@ func (b *writeBuffer) writeBinaryDatum(d parser.Datum) error {
 
 const pgTimeStampFormat = "2006-01-02 15:04:05.999999999-07:00"
 
-// formatTs formats t into a format lib/pq understands.
-// Mostly cribbed from github.com/lib/pq.
+// formatTs formats t into a format cockroachdb/pq understands.
+// Mostly cribbed from github.com/cockroachdb/pq.
 func formatTs(t time.Time) (b []byte) {
 	// Need to send dates before 0001 A.D. with " BC" suffix, instead of the
 	// minus sign preferred by Go.

sql/pgwire/v3.go

@@ -25,14 +25,13 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/lib/pq/oid"
-
 	"github.com/cockroachdb/cockroach/roachpb"
 	"github.com/cockroachdb/cockroach/sql"
 	"github.com/cockroachdb/cockroach/sql/parser"
 	"github.com/cockroachdb/cockroach/util"
 	"github.com/cockroachdb/cockroach/util/log"
 	"github.com/cockroachdb/cockroach/util/tracing"
+	"github.com/cockroachdb/pq/oid"
 )
 
 //go:generate stringer -type=clientMessageType