Skip to content

Commit

Permalink
Add support for executing SQL statements from the command line.
Browse files Browse the repository at this point in the history
With this patch the user can use "sql -e ..." to execute statements,
print their results and exit, without an interactive prompt.

Fixes #3817
  • Loading branch information
knz committed Jan 26, 2016
1 parent fedc14b commit 4288962
Show file tree
Hide file tree
Showing 5 changed files with 159 additions and 18 deletions.
31 changes: 31 additions & 0 deletions cli/cli_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -467,6 +467,37 @@ range_max_bytes: 67108864
// zone ls
}

func Example_sql() {
c := newCLITest()
defer c.Stop()

c.RunWithArgs([]string{"sql", "-e", "create database t; create table t.f (x int, y int); insert into t.f values (42, 69)"})
c.RunWithArgs([]string{"sql", "-e", "select 3", "select * from t.f"})
c.RunWithArgs([]string{"sql", "-e", "begin", "select 3", "commit"})
c.RunWithArgs([]string{"sql", "-e", "select 3; select * from t.f"})

// Output:
// sql -e create database t; create table t.f (x int, y int); insert into t.f values (42, 69)
// OK
// sql -e select 3 select * from t.f
// 1 row
// 3
// 3
// 1 row
// x y
// 42 69
// sql -e begin select 3 commit
// OK
// 1 row
// 3
// 3
// OK
// sql -e select 3; select * from t.f
// 1 row
// x y
// 42 69
}

func Example_user() {
c := newCLITest()
defer c.Stop()
Expand Down
45 changes: 45 additions & 0 deletions cli/context.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Raphael 'kena' Poss (knz@cockroachlabs.com)

package cli

// Context holds parameters needed to setup a server.
import "github.com/cockroachdb/cockroach/server"

// Calling "cli".initFlags(ctx *Context) will initialize Context using
// command flags. Keep in sync with "cli/flags.go".
type Context struct {
// Embed the server context.
server.Context

// OneShotSQL indicates the SQL client should run the command-line
// statement(s) and terminate directly, without presenting a REPL to
// the user.
OneShotSQL bool
}

// NewContext returns a Context with default values.
func NewContext() *Context {
ctx := &Context{}
ctx.InitDefaults()
return ctx
}

// InitDefaults sets up the default values for a Context.
func (ctx *Context) InitDefaults() {
ctx.Context.InitDefaults()
ctx.OneShotSQL = false
}
20 changes: 16 additions & 4 deletions cli/flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ import (
"strings"

"github.com/spf13/cobra"

"github.com/cockroachdb/cockroach/server"
)

var maxResults int64
Expand Down Expand Up @@ -150,17 +148,26 @@ var flagUsage = map[string]string{
"password": `
The created user's password. If provided, disables prompting. Pass '-' to provide
the password on standard input.
`,
"execute": `
Execute the SQL statement(s) on the command line, then exit. Each
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. The results of the last SQL statement in each
positional argument are printed on the standard output.
`,
}

func normalizeStdFlagName(s string) string {
return strings.Replace(s, "_", "-", -1)
}

// initFlags sets the server.Context values to flag values.
// initFlags sets the cli.Context values to flag values.
// Keep in sync with "server/context.go". Values in Context should be
// settable here.
func initFlags(ctx *server.Context) {
func initFlags(ctx *Context) {
// Map any flags registered in the standard "flag" package into the
// top-level cockroach command.
pf := cockroachCmd.PersistentFlags()
Expand Down Expand Up @@ -250,6 +257,11 @@ func initFlags(ctx *server.Context) {
f.StringVar(&ctx.Certs, "certs", ctx.Certs, flagUsage["certs"])
}

{
f := sqlShellCmd.Flags()
f.BoolVarP(&ctx.OneShotSQL, "execute", "e", ctx.OneShotSQL, flagUsage["execute"])
}

// Max results flag for scan, reverse scan, and range list.
for _, cmd := range []*cobra.Command{scanCmd, reverseScanCmd, lsRangesCmd} {
f := cmd.Flags()
Expand Down
73 changes: 63 additions & 10 deletions cli/sql.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package cli

import (
"database/sql"
"fmt"
"io"
"os"
Expand All @@ -40,17 +41,12 @@ var sqlShellCmd = &cobra.Command{
Long: `
Open a sql shell running against the cockroach database at --addr.
`,
Run: runTerm, // TODO(tschottdorf): should be able to return err code when reading from stdin
Run: runTerm,
}

func runTerm(cmd *cobra.Command, args []string) {
if len(args) != 0 {
mustUsage(cmd)
return
}

db := makeSQLClient()
defer func() { _ = db.Close() }()
// runInteractive runs the SQL client interactively, presenting
// a prompt to the user for each statement.
func runInteractive(db *sql.DB) {

liner := liner.NewLiner()
defer func() {
Expand Down Expand Up @@ -79,6 +75,9 @@ func runTerm(cmd *cobra.Command, args []string) {
var stmt []string
var l string
var err error

exitCode := 0

for {
if len(stmt) == 0 {
l, err = liner.Prompt(fullPrompt)
Expand All @@ -88,6 +87,7 @@ func runTerm(cmd *cobra.Command, args []string) {
if err != nil {
if err != io.EOF {
fmt.Fprintf(osStderr, "Input error: %s\n", err)
exitCode = 1
}
break
}
Expand All @@ -112,11 +112,64 @@ func runTerm(cmd *cobra.Command, args []string) {
fullStmt := strings.Join(stmt, "\n")
liner.AppendHistory(fullStmt)

exitCode = 0
if err := runPrettyQuery(db, os.Stdout, fullStmt); err != nil {
fmt.Println(err)
fmt.Fprintln(osStderr, err)
exitCode = 1
}

// Clear the saved statement.
stmt = stmt[:0]
}

if exitCode != 0 {
os.Exit(exitCode)
}
}

// runOneStatement executes one statement and terminates
// 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)
}

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"))
}
}

}
}

func runTerm(cmd *cobra.Command, args []string) {
if !(len(args) >= 1 && context.OneShotSQL) && len(args) != 0 {
mustUsage(cmd)
return
}

db := makeSQLClient()
defer func() { _ = db.Close() }()

if context.OneShotSQL {
// Single-line sql; run as simple as possible, without noise on stdout.
runStatements(db, args)
} else {
runInteractive(db)
}

}
8 changes: 4 additions & 4 deletions cli/start.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,8 @@ import (
"github.com/spf13/cobra"
)

// Context is the CLI Context used for the server.
var context = server.NewContext()
// Context is the CLI Context used for the command-line client.
var context = NewContext()

var errMissingParams = errors.New("missing or invalid parameters")

Expand Down Expand Up @@ -183,7 +183,7 @@ func runStart(_ *cobra.Command, _ []string) error {
}

log.Info("starting cockroach node")
s, err := server.NewServer(context, stopper)
s, err := server.NewServer(&context.Context, stopper)
if err != nil {
return util.Errorf("failed to start Cockroach server: %s", err)
}
Expand Down Expand Up @@ -285,7 +285,7 @@ completed, the server exits.

// runQuit accesses the quit shutdown path.
func runQuit(_ *cobra.Command, _ []string) {
admin := client.NewAdminClient(&context.Context, context.Addr, client.Quit)
admin := client.NewAdminClient(&context.Context.Context, context.Addr, client.Quit)
body, err := admin.Get()
// TODO(tschottdorf): needs cleanup. An error here can happen if the shutdown
// happened faster than the HTTP request made it back.
Expand Down

0 comments on commit 4288962

Please sign in to comment.