Merge #83118

83118: cli/sql: support for client-side `\o` r=rafiss a=knz

Release note (cli change): `cockroach sql` (and thus `cockroach demo`
too) now support the client-side commands `\o` and `\qecho` like
`psql`. The command `\o` can redirect the output of SQL queries to a
file. `\qecho` adds an arbitrary text to the current query output file.

Co-authored-by: Raphael 'kena' Poss <[email protected]>

pkg/cli/auth.go

@@ -182,7 +182,7 @@ func runLogout(cmd *cobra.Command, args []string) (resErr error) {
 		username)
 	return sqlExecCtx.RunQueryAndFormatResults(
 		context.Background(),
-		sqlConn, os.Stdout, stderr, logoutQuery)
+		sqlConn, os.Stdout, os.Stdout, stderr, logoutQuery)
 }
 
 var authListCmd = &cobra.Command{
@@ -214,7 +214,7 @@ SELECT username,
   FROM system.web_sessions`)
 	return sqlExecCtx.RunQueryAndFormatResults(
 		context.Background(),
-		sqlConn, os.Stdout, stderr, logoutQuery)
+		sqlConn, os.Stdout, os.Stdout, stderr, logoutQuery)
 }
 
 var authCmds = []*cobra.Command{

pkg/cli/clisqlexec/run_query.go

@@ -39,9 +39,10 @@ func (sqlExecCtx *Context) RunQuery(
 
 // RunQueryAndFormatResults takes a 'query' with optional 'parameters'.
 // It runs the sql query and writes output to 'w'.
+// Timings is enabled are written to 'tw'.
 // Errors and warnings, if any, are printed to 'ew'.
 func (sqlExecCtx *Context) RunQueryAndFormatResults(
-	ctx context.Context, conn clisqlclient.Conn, w, ew io.Writer, fn clisqlclient.QueryFn,
+	ctx context.Context, conn clisqlclient.Conn, w, tw, ew io.Writer, fn clisqlclient.QueryFn,
 ) (err error) {
 	startTime := timeutil.Now()
 	rows, isMultiStatementQuery, err := fn(ctx, conn)
@@ -113,7 +114,7 @@ func (sqlExecCtx *Context) RunQueryAndFormatResults(
 		} else if !more {
 			// We must call maybeShowTimes after rows has been closed, which is after
 			// NextResultSet returns false.
-			sqlExecCtx.maybeShowTimes(ctx, conn, w, ew, isMultiStatementQuery, startTime, queryCompleteTime)
+			sqlExecCtx.maybeShowTimes(ctx, conn, tw, ew, isMultiStatementQuery, startTime, queryCompleteTime)
 			return nil
 		}
 	}

pkg/cli/clisqlexec/run_query_test.go

@@ -41,7 +41,7 @@ func runQueryAndFormatResults(
 ) (err error) {
 	return testExecCtx.RunQueryAndFormatResults(
 		context.Background(),
-		conn, w, ioutil.Discard, fn)
+		conn, w, ioutil.Discard, ioutil.Discard, fn)
 }
 
 func TestRunQuery(t *testing.T) {

pkg/cli/clisqlshell/context.go

@@ -11,7 +11,9 @@
 package clisqlshell
 
 import (
+	"bufio"
 	"context"
+	"io"
 	"os"
 	"time"
 
@@ -54,6 +56,14 @@ type internalContext struct {
 	stdout *os.File
 	stderr *os.File
 
+	// queryOutputFile is the output file configured via \o.
+	// This can be the same as stdout (\o without argument).
+	// Note: we use .queryOutput for query execution, which
+	// is buffered (via queryOutputBuf).
+	queryOutputFile *os.File
+	queryOutputBuf  *bufio.Writer
+	queryOutput     io.Writer
+
 	// quitAfterExecStmts tells the shell whether to quit
 	// after processing the execStmts.
 	quitAfterExecStmts bool

pkg/cli/clisqlshell/sql.go

@@ -67,13 +67,15 @@ Connection
   \c, \connect {[DB] [USER] [HOST] [PORT] | [URL]}
                     connect to a server or print the current connection URL.
                     (Omitted values reuse previous parameters. Use '-' to skip a field.)
-  \password [USERNAME]   
+  \password [USERNAME]
                     securely change the password for a user
 
 Input/Output
   \echo [STRING]    write the provided string to standard output.
   \i                execute commands from the specified file.
   \ir               as \i, but relative to the location of the current script.
+  \o [FILE]         send all query results to the specified file.
+  \qecho [STRING]   write the provided string to the query output stream (see \o).
 
 Informational
   \l                list all databases in the CockroachDB cluster.
@@ -1174,6 +1176,10 @@ func (c *cliState) doHandleCliCmd(loopState, nextState cliStateEnum) cliStateEnu
 	case `\echo`:
 		fmt.Fprintln(c.iCtx.stdout, strings.Join(cmd[1:], " "))
 
+	case `\qecho`:
+		fmt.Fprintln(c.iCtx.queryOutput, strings.Join(cmd[1:], " "))
+		c.maybeFlushOutput()
+
 	case `\set`:
 		return c.handleSet(cmd[1:], loopState, errState)
 
@@ -1189,6 +1195,9 @@ func (c *cliState) doHandleCliCmd(loopState, nextState cliStateEnum) cliStateEnu
 	case `\ir`:
 		return c.runInclude(cmd[1:], loopState, errState, true /* relative */)
 
+	case `\o`:
+		return c.runOpen(cmd[1:], loopState, errState)
+
 	case `\p`:
 		// This is analogous to \show but does not need a special case.
 		// Implemented for compatibility with psql.
@@ -1793,10 +1802,12 @@ func (c *cliState) doRunStatements(nextState cliStateEnum) cliStateEnum {
 				c.concatLines,
 			)
 		}
+		defer c.maybeFlushOutput()
 		return c.sqlExecCtx.RunQueryAndFormatResults(
 			ctx,
 			c.conn,
-			c.iCtx.stdout,
+			c.iCtx.queryOutput, // query output.
+			c.iCtx.stdout,      // timings.
 			c.iCtx.stderr,
 			q,
 		)
@@ -1831,8 +1842,12 @@ func (c *cliState) doRunStatements(nextState cliStateEnum) cliStateEnum {
 				traceType = "kv"
 			}
 			if err := c.runWithInterruptableCtx(func(ctx context.Context) error {
+				defer c.maybeFlushOutput()
 				return c.sqlExecCtx.RunQueryAndFormatResults(ctx,
-					c.conn, c.iCtx.stdout, c.iCtx.stderr,
+					c.conn,
+					c.iCtx.queryOutput, // query output
+					c.iCtx.stdout,      // timings
+					c.iCtx.stderr,      // errors
 					clisqlclient.MakeQuery(fmt.Sprintf("SHOW %s TRACE FOR SESSION", traceType)))
 			}); err != nil {
 				clierror.OutputError(c.iCtx.stderr, err, true /*showSeverity*/, false /*verbose*/)
@@ -1922,6 +1937,11 @@ func (c *cliState) doRunShell(state cliStateEnum, cmdIn, cmdOut, cmdErr *os.File
 		}
 		switch state {
 		case cliStart:
+			defer func() {
+				if err := c.closeOutputFile(); err != nil {
+					fmt.Fprintf(cmdErr, "warning: closing output file: %v\n", err)
+				}
+			}()
 			cleanupFn, err := c.configurePreShellDefaults(cmdIn, cmdOut, cmdErr)
 			defer cleanupFn()
 			if err != nil {
@@ -1990,6 +2010,8 @@ func (c *cliState) configurePreShellDefaults(
 	cmdIn, cmdOut, cmdErr *os.File,
 ) (cleanupFn func(), err error) {
 	c.iCtx.stdout = cmdOut
+	c.iCtx.queryOutputFile = cmdOut
+	c.iCtx.queryOutput = cmdOut
 	c.iCtx.stderr = cmdErr
 
 	if c.sqlExecCtx.TerminalOutput {
@@ -2036,6 +2058,7 @@ func (c *cliState) configurePreShellDefaults(
 		// The readline library may have a custom file descriptor for stdout.
 		// Use that for further output.
 		c.iCtx.stdout = c.ins.Stdout()
+		c.iCtx.queryOutputFile = c.ins.Stdout()
 
 		// If the user has used bind -v or bind -l in their ~/.editrc,
 		// this will reset the standard bindings. However we really
@@ -2317,3 +2340,77 @@ func (c *cliState) getSessionVarValue(sessionVar string) (string, error) {
 	}
 	return "", nil
 }
+
+func (c *cliState) runOpen(cmd []string, contState, errState cliStateEnum) (resState cliStateEnum) {
+	if len(cmd) > 1 {
+		return c.invalidSyntax(errState)
+	}
+
+	outputFile := "" // no file: reset to stdout
+	if len(cmd) == 1 {
+		outputFile = cmd[0]
+	}
+	if err := c.openOutputFile(outputFile); err != nil {
+		c.exitErr = err
+		fmt.Fprintf(c.iCtx.stderr, "%v\n", c.exitErr)
+		return errState
+	}
+	return contState
+}
+
+func (c *cliState) openOutputFile(file string) error {
+	var f *os.File
+	if file != "" {
+		// First check whether the new file can be opened.
+		// (We keep the previous one otherwise.)
+		// NB: permission 0666 mimics what psql does: fopen(file, "w").
+		var err error
+		f, err = os.OpenFile(file, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
+		if err != nil {
+			return err
+		}
+	}
+	// Close the previous file.
+	if err := c.closeOutputFile(); err != nil {
+		return err
+	}
+	if f != nil {
+		c.iCtx.queryOutputFile = f
+		c.iCtx.queryOutputBuf = bufio.NewWriter(f)
+		c.iCtx.queryOutput = c.iCtx.queryOutputBuf
+	}
+	return nil
+}
+
+func (c *cliState) maybeFlushOutput() {
+	if err := c.maybeFlushOutputInternal(); err != nil {
+		fmt.Fprintf(c.iCtx.stderr, "warning: flushing output file: %v", err)
+	}
+}
+
+func (c *cliState) maybeFlushOutputInternal() error {
+	if c.iCtx.queryOutputBuf == nil {
+		return nil
+	}
+	return c.iCtx.queryOutputBuf.Flush()
+}
+
+func (c *cliState) closeOutputFile() error {
+	if c.iCtx.queryOutputBuf == nil {
+		return nil
+	}
+	if err := c.maybeFlushOutputInternal(); err != nil {
+		return err
+	}
+	if c.iCtx.queryOutputFile == c.iCtx.stdout {
+		// Nothing to do.
+		return nil
+	}
+
+	// Close file and reset.
+	err := c.iCtx.queryOutputFile.Close()
+	c.iCtx.queryOutputFile = c.iCtx.stdout
+	c.iCtx.queryOutput = c.iCtx.stdout
+	c.iCtx.queryOutputBuf = nil
+	return err
+}

pkg/cli/interactive_tests/test_local_cmds.tcl

@@ -266,6 +266,36 @@ eexpect Description
 eexpect root@
 end_test
 
+start_test "Check that the output of queries can be redirected to a file."
+send "\\o logs/query-output.txt\r"
+eexpect root@
+send "select 'hello world';\r"
+eexpect root@
+# Check the output is flushed immediately.
+system "grep -q world logs/query-output.txt"
+send "\\qecho universe\r"
+eexpect root@
+# Check qecho works.
+system "grep -q universe logs/query-output.txt"
+# Check the previous output was not erased by subsequent command.
+system "grep -q world logs/query-output.txt"
+end_test
+
+start_test "Check that errors are not redirected."
+send "select planet;\r"
+# Check the output is not redirected.
+eexpect "column \"planet\" does not exist"
+system "grep -vq planet logs/query-output.txt"
+end_test
+
+start_test "Check that the query output can be reset to stdout."
+send "\\o\r"
+eexpect root@
+send "select 'hel'||'lo';\r"
+eexpect "hello"
+eexpect root@
+end_test
+
 # Finally terminate with Ctrl+D.
 send_eof
 eexpect eof

pkg/cli/zip.go

@@ -13,6 +13,7 @@ package cli
 import (
 	"context"
 	"fmt"
+	"io/ioutil"
 	"net"
 	"os"
 	"strings"
@@ -356,7 +357,7 @@ func (zc *debugZipContext) dumpTableDataForZip(
 			}
 			// Pump the SQL rows directly into the zip writer, to avoid
 			// in-RAM buffering.
-			return sqlExecCtx.RunQueryAndFormatResults(ctx, conn, w, stderr, clisqlclient.MakeQuery(query))
+			return sqlExecCtx.RunQueryAndFormatResults(ctx, conn, w, ioutil.Discard, stderr, clisqlclient.MakeQuery(query))
 		}()
 		if sqlErr != nil {
 			if cErr := zc.z.createError(s, name, sqlErr); cErr != nil {