From df18c7fc0268ffdb7e8ff580b446e7e31a10def7 Mon Sep 17 00:00:00 2001 From: Raphael 'kena' Poss Date: Tue, 21 Jun 2022 12:47:35 +0200 Subject: [PATCH] cli/sql: support for client-side `\o` 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. --- pkg/cli/auth.go | 4 +- pkg/cli/clisqlexec/run_query.go | 5 +- pkg/cli/clisqlexec/run_query_test.go | 2 +- pkg/cli/clisqlshell/context.go | 10 ++ pkg/cli/clisqlshell/sql.go | 103 +++++++++++++++++- pkg/cli/interactive_tests/test_local_cmds.tcl | 30 +++++ pkg/cli/zip.go | 3 +- 7 files changed, 148 insertions(+), 9 deletions(-) diff --git a/pkg/cli/auth.go b/pkg/cli/auth.go index 4a499513959c..215627afd34d 100644 --- a/pkg/cli/auth.go +++ b/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{ diff --git a/pkg/cli/clisqlexec/run_query.go b/pkg/cli/clisqlexec/run_query.go index d62daa7f0c93..0db250bdd535 100644 --- a/pkg/cli/clisqlexec/run_query.go +++ b/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 } } diff --git a/pkg/cli/clisqlexec/run_query_test.go b/pkg/cli/clisqlexec/run_query_test.go index 4f831d05d4d3..ceff3ee41f54 100644 --- a/pkg/cli/clisqlexec/run_query_test.go +++ b/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) { diff --git a/pkg/cli/clisqlshell/context.go b/pkg/cli/clisqlshell/context.go index f94fa7287c69..0bee77ab93a8 100644 --- a/pkg/cli/clisqlshell/context.go +++ b/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 diff --git a/pkg/cli/clisqlshell/sql.go b/pkg/cli/clisqlshell/sql.go index d3cce0fe7e2e..4adaf79f3325 100644 --- a/pkg/cli/clisqlshell/sql.go +++ b/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 +} diff --git a/pkg/cli/interactive_tests/test_local_cmds.tcl b/pkg/cli/interactive_tests/test_local_cmds.tcl index 84f55f346130..28464aa8d0b0 100755 --- a/pkg/cli/interactive_tests/test_local_cmds.tcl +++ b/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 diff --git a/pkg/cli/zip.go b/pkg/cli/zip.go index b587f82af2ee..393803fb6b91 100644 --- a/pkg/cli/zip.go +++ b/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 {