From 199d1051dd382340dfbe0d2719139cfc563ab3ca Mon Sep 17 00:00:00 2001 From: Peter Mattis Date: Tue, 2 Feb 2016 09:58:19 -0500 Subject: [PATCH] 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 #4016. --- GLOCKFILE | 3 +- acceptance/main_test.go | 3 +- cli/cli_test.go | 8 ++++ cli/flags.go | 4 +- cli/sql.go | 42 ++++++++++-------- cli/sql_util.go | 38 +++++++++------- cli/sql_util_test.go | 98 +++++++++++++++++++++++++---------------- client/doc.go | 2 +- sql/bench_test.go | 2 +- sql/driver/wire.go | 2 +- sql/pgwire/types.go | 6 +-- sql/pgwire/v3.go | 2 +- sql/pgwire_test.go | 12 ++--- storage/log_test.go | 2 +- 14 files changed, 132 insertions(+), 92 deletions(-) diff --git a/GLOCKFILE b/GLOCKFILE index c2ff61f3df90..805a368e5078 100644 --- a/GLOCKFILE +++ b/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 diff --git a/acceptance/main_test.go b/acceptance/main_test.go index 56b741867e9d..af6588a1d80d 100644 --- a/acceptance/main_test.go +++ b/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") diff --git a/cli/cli_test.go b/cli/cli_test.go index cc4dc78d8d27..9e8212ff5ed2 100644 --- a/cli/cli_test.go +++ b/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() { diff --git a/cli/flags.go b/cli/flags.go index 3e839813c11c..0284eb845ca6 100644 --- a/cli/flags.go +++ b/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 diff --git a/cli/sql.go b/cli/sql.go index b19ff0c1a3c8..51c1be7d0f78 100644 --- a/cli/sql.go +++ b/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 } } } diff --git a/cli/sql_util.go b/cli/sql_util.go index 69bc69bca88f..ed8b65fed1d1 100644 --- a/cli/sql_util.go +++ b/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,7 +66,7 @@ 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. @@ -71,30 +74,28 @@ func runQueryWithFormat(db *sql.DB, format fmtMap, query string, parameters ...i []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 } diff --git a/cli/sql_util_test.go b/cli/sql_util_test.go index 5b44e76f750b..e91839c1bb36 100644 --- a/cli/sql_util_test.go +++ b/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() diff --git a/client/doc.go b/client/doc.go index f80b6a4d8404..39fbf3213b45 100644 --- a/client/doc.go +++ b/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 diff --git a/sql/bench_test.go b/sql/bench_test.go index b3c3648eccfd..7dce83bbd317 100644 --- a/sql/bench_test.go +++ b/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)) { diff --git a/sql/driver/wire.go b/sql/driver/wire.go index a9278f71a112..ccfc061ebc83 100644 --- a/sql/driver/wire.go +++ b/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 diff --git a/sql/pgwire/types.go b/sql/pgwire/types.go index 4e16fc4f3e0e..2d67f0518761 100644 --- a/sql/pgwire/types.go +++ b/sql/pgwire/types.go @@ -24,7 +24,7 @@ import ( "strconv" "time" - "github.com/lib/pq/oid" + "github.com/cockroachdb/pq/oid" "github.com/cockroachdb/cockroach/sql/parser" "github.com/cockroachdb/cockroach/util" @@ -185,8 +185,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. diff --git a/sql/pgwire/v3.go b/sql/pgwire/v3.go index 3eff4a8540d0..8d81b06fd974 100644 --- a/sql/pgwire/v3.go +++ b/sql/pgwire/v3.go @@ -25,7 +25,7 @@ import ( "strconv" "strings" - "github.com/lib/pq/oid" + "github.com/cockroachdb/pq/oid" "github.com/cockroachdb/cockroach/roachpb" "github.com/cockroachdb/cockroach/sql" diff --git a/sql/pgwire_test.go b/sql/pgwire_test.go index dc1663c22998..176d6a7b7b33 100644 --- a/sql/pgwire_test.go +++ b/sql/pgwire_test.go @@ -25,7 +25,7 @@ import ( "reflect" "testing" - "github.com/lib/pq" + "github.com/cockroachdb/pq" "github.com/cockroachdb/cockroach/security" "github.com/cockroachdb/cockroach/security/securitytest" @@ -467,7 +467,7 @@ func TestCmdCompleteVsEmptyStatements(t *testing.T) { } defer db.Close() - // lib/pq handles the empty query response by returning a nil driver.Result. + // cockroachdb/pq handles the empty query response by returning a nil driver.Result. // Unfortunately sql.Exec wraps that, nil or not, in a sql.Result which doesn't // expose the underlying driver.Result. // sql.Result does however have methods which attempt to dereference the underlying @@ -479,7 +479,7 @@ func TestCmdCompleteVsEmptyStatements(t *testing.T) { if err != nil { t.Fatal(err) } - _, _ = nonempty.RowsAffected() // should not panic if lib/pq returned a non-nil result. + _, _ = nonempty.RowsAffected() // should not panic if cockroachdb/pq returned a non-nil result. empty, err := db.Exec(" ") if err != nil { @@ -488,12 +488,12 @@ func TestCmdCompleteVsEmptyStatements(t *testing.T) { defer func() { _ = recover() }() - _, _ = empty.RowsAffected() // should panic if lib/pq returned a nil result as expected. + _, _ = empty.RowsAffected() // should panic if cockroachdb/pq returned a nil result as expected. t.Fatal("should not get here -- empty result from empty query should panic first") // TODO(dt): clean this up with testify/assert and add tests for less trivial empty queries. } -// Unfortunately lib/pq doesn't expose returned command tags directly, but we can test +// Unfortunately cockroachdb/pq doesn't expose returned command tags directly, but we can test // the methods where it depends on their values (Begin, Commit, RowsAffected for INSERTs). func TestPGCommandTags(t *testing.T) { defer leaktest.AfterTest(t) @@ -526,7 +526,7 @@ func TestPGCommandTags(t *testing.T) { t.Fatal(err) } - // lib/pq has a special-case for INSERT (due to oids), so test insert and update statements. + // cockroachdb/pq has a special-case for INSERT (due to oids), so test insert and update statements. res, err := db.Exec("INSERT INTO testing.tags VALUES (1, 1), (2, 2)") if err != nil { t.Fatal(err) diff --git a/storage/log_test.go b/storage/log_test.go index 8df7538107eb..7c6769fd14c9 100644 --- a/storage/log_test.go +++ b/storage/log_test.go @@ -23,7 +23,7 @@ import ( "os" "testing" - _ "github.com/lib/pq" + _ "github.com/cockroachdb/pq" "github.com/cockroachdb/cockroach/client" "github.com/cockroachdb/cockroach/keys"