From 4288962803979d272724fe2f81c1c1c8c7809857 Mon Sep 17 00:00:00 2001
From: Raphael 'kena' Poss <knz@cockroachlabs.com>
Date: Tue, 26 Jan 2016 21:42:57 +0000
Subject: [PATCH] Add support for executing SQL statements from the command
 line.

With this patch the user can use "sql -e ..." to execute statements,
print their results and exit, without an interactive prompt.

Fixes #3817
---
 cli/cli_test.go | 31 +++++++++++++++++++++
 cli/context.go  | 45 ++++++++++++++++++++++++++++++
 cli/flags.go    | 20 +++++++++++---
 cli/sql.go      | 73 ++++++++++++++++++++++++++++++++++++++++++-------
 cli/start.go    |  8 +++---
 5 files changed, 159 insertions(+), 18 deletions(-)
 create mode 100644 cli/context.go

diff --git a/cli/cli_test.go b/cli/cli_test.go
index 443d757dfd8a..821c3f65f9cb 100644
--- a/cli/cli_test.go
+++ b/cli/cli_test.go
@@ -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()
diff --git a/cli/context.go b/cli/context.go
new file mode 100644
index 000000000000..f241dab0b60d
--- /dev/null
+++ b/cli/context.go
@@ -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
+}
diff --git a/cli/flags.go b/cli/flags.go
index 598f92db62f1..9abab25f625a 100644
--- a/cli/flags.go
+++ b/cli/flags.go
@@ -22,8 +22,6 @@ import (
 	"strings"
 
 	"github.com/spf13/cobra"
-
-	"github.com/cockroachdb/cockroach/server"
 )
 
 var maxResults int64
@@ -150,6 +148,15 @@ 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.
 `,
 }
 
@@ -157,10 +164,10 @@ 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()
@@ -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()
diff --git a/cli/sql.go b/cli/sql.go
index 2bebac0e5faa..77f76ee6a810 100644
--- a/cli/sql.go
+++ b/cli/sql.go
@@ -17,6 +17,7 @@
 package cli
 
 import (
+	"database/sql"
 	"fmt"
 	"io"
 	"os"
@@ -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() {
@@ -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)
@@ -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
 		}
@@ -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)
+	}
+
 }
diff --git a/cli/start.go b/cli/start.go
index c7f4dac71474..462aa95d2df1 100644
--- a/cli/start.go
+++ b/cli/start.go
@@ -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")
 
@@ -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)
 	}
@@ -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.