cli/sql: implement the \i and \ir client-side command

Release note (cli change): The SQL shell (`cockroach sql`, `cockroach
demo`) now support the `\i` and `\ir` client-side command which reads
SQL file and evaluates its content in-place.

`\ir` differs from `\i` in that the file name is resolved relative
to the location of the script containing the `\ir` command. This makes
`\ir` likely more desirable in the general case.

Occurences of `\q` inside a file included via `\i`/`\ir stop
evaluation of the file and resume evaluation of the file that included
it.

This feature is compatible with the identically named `psql` commands.
It is meant to help compose complex initialization scripts from a
library of standard components.

For example, one could be defining each table and its initial contents
in separate SQL files, and then use different super-files to include
different tables depending on the desired final schema.

pkg/cli/cli_test.go

@@ -2089,3 +2089,46 @@ func Example_read_from_file() {
 	// ERROR: column "undefined" does not exist
 	// SQLSTATE: 42703
 }
+
+// Example_includes tests the \i command.
+func Example_includes() {
+	c := newCLITest(cliTestParams{})
+	defer c.cleanup()
+
+	c.RunWithArgs([]string{"sql", "-f", "testdata/i_twolevels1.sql"})
+	c.RunWithArgs([]string{"sql", "-f", "testdata/i_multiline.sql"})
+	c.RunWithArgs([]string{"sql", "-f", "testdata/i_stopmiddle.sql"})
+	c.RunWithArgs([]string{"sql", "-f", "testdata/i_maxrecursion.sql"})
+
+	// Output:
+	// sql -f testdata/i_twolevels1.sql
+	// > SELECT 123;
+	// ?column?
+	// 123
+	// > SELECT 789;
+	// ?column?
+	// 789
+	// ?column?
+	// 456
+	// sql -f testdata/i_multiline.sql
+	// ERROR: at or near "\": syntax error
+	// SQLSTATE: 42601
+	// DETAIL: source SQL:
+	// SELECT -- incomplete statement, \i invalid
+	// \i testdata/i_twolevels2.sql
+	// ^
+	// HINT: try \h SELECT
+	// ERROR: at or near "\": syntax error
+	// SQLSTATE: 42601
+	// DETAIL: source SQL:
+	// SELECT -- incomplete statement, \i invalid
+	// \i testdata/i_twolevels2.sql
+	// ^
+	// HINT: try \h SELECT
+	// sql -f testdata/i_stopmiddle.sql
+	// ?column?
+	// 123
+	// sql -f testdata/i_maxrecursion.sql
+	// \i: too many recursion levels (max 10)
+	// ERROR: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: testdata/i_maxrecursion.sql: \i: too many recursion levels (max 10)
+}

pkg/cli/sql.go

@@ -66,6 +66,8 @@ Query Buffer
 
 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.
 
 Informational
   \l                list all databases in the CockroachDB cluster.
@@ -126,6 +128,12 @@ type cliState struct {
 	// buf is used to read lines if isInteractive is false.
 	buf *bufio.Reader
 
+	// levels is the number of inclusion recursion levels.
+	levels int
+	// includeDir is the directory relative to which relative
+	// includes (\ir) resolve the file name.
+	includeDir string
+
 	// The prompt at the beginning of a multi-line entry.
 	fullPrompt string
 	// The prompt on a continuation line in a multi-line entry.
@@ -1044,6 +1052,12 @@ func (c *cliState) doHandleCliCmd(loopState, nextState cliStateEnum) cliStateEnu
 	case `\!`:
 		return c.runSyscmd(c.lastInputLine, loopState, errState)
 
+	case `\i`:
+		return c.runInclude(cmd[1:], loopState, errState, false /* relative */)
+
+	case `\ir`:
+		return c.runInclude(cmd[1:], loopState, errState, true /* relative */)
+
 	case `\p`:
 		// This is analogous to \show but does not need a special case.
 		// Implemented for compatibility with psql.
@@ -1113,6 +1127,65 @@ func (c *cliState) doHandleCliCmd(loopState, nextState cliStateEnum) cliStateEnu
 	return loopState
 }
 
+const maxRecursionLevels = 10
+
+func (c *cliState) runInclude(
+	cmd []string, contState, errState cliStateEnum, relative bool,
+) (resState cliStateEnum) {
+	if len(cmd) != 1 {
+		return c.invalidSyntax(errState, `%s. Try \? for help.`, c.lastInputLine)
+	}
+
+	if c.levels >= maxRecursionLevels {
+		c.exitErr = errors.Newf(`\i: too many recursion levels (max %d)`, maxRecursionLevels)
+		fmt.Fprintf(stderr, "%v\n", c.exitErr)
+		return errState
+	}
+
+	if len(c.partialLines) > 0 {
+		return c.invalidSyntax(errState, `cannot use \i during multi-line entry.`)
+	}
+
+	filename := cmd[0]
+	if !filepath.IsAbs(filename) && relative {
+		// In relative mode, the filename is resolved relative to the
+		// surrounding script.
+		filename = filepath.Join(c.includeDir, filename)
+	}
+
+	f, err := os.Open(filename)
+	if err != nil {
+		fmt.Fprintln(stderr, err)
+		c.exitErr = err
+		return errState
+	}
+	// Close the file at the end.
+	defer func() {
+		if err := f.Close(); err != nil {
+			fmt.Fprintf(stderr, "error: closing %s: %v\n", filename, err)
+			c.exitErr = errors.CombineErrors(c.exitErr, err)
+			resState = errState
+		}
+	}()
+
+	newState := cliState{
+		conn:       c.conn,
+		includeDir: filepath.Dir(filename),
+		ins:        noLineEditor,
+		buf:        bufio.NewReader(f),
+		levels:     c.levels + 1,
+	}
+
+	if err := newState.doRunShell(cliStartLine, f); err != nil {
+		// Note: a message was already printed on stderr at the point at
+		// which the error originated. No need to repeat it here.
+		c.exitErr = errors.Wrapf(err, "%v", filename)
+		return errState
+	}
+
+	return contState
+}
+
 func (c *cliState) doPrepareStatementLine(
 	startState, contState, checkState, execState cliStateEnum,
 ) cliStateEnum {
@@ -1293,9 +1366,14 @@ func (c *cliState) doDecidePath() cliStateEnum {
 // runInteractive runs the SQL client interactively, presenting
 // a prompt to the user for each statement.
 func runInteractive(conn *sqlConn, cmdIn *os.File) (exitErr error) {
-	c := cliState{conn: conn}
+	c := cliState{
+		conn:       conn,
+		includeDir: ".",
+	}
+	return c.doRunShell(cliStart, cmdIn)
+}
 
-	state := cliStart
+func (c *cliState) doRunShell(state cliStateEnum, cmdIn *os.File) (exitErr error) {
 	for {
 		if state == cliStop {
 			break

pkg/cli/testdata/i_maxrecursion.sql

@@ -0,0 +1 @@
+\i testdata/i_maxrecursion.sql

pkg/cli/testdata/i_multiline.sql

@@ -0,0 +1,3 @@
+SELECT -- incomplete statement, \i invalid
+\i testdata/i_twolevels2.sql
+123;

pkg/cli/testdata/i_stopmiddle.sql

@@ -0,0 +1,3 @@
+SELECT 123;
+\q
+SELECT 456;

pkg/cli/testdata/i_twolevels1.sql

@@ -0,0 +1,9 @@
+-- test that \echo spills into the second level
+\set echo
+
+\i testdata/i_twolevels2.sql
+
+-- at this point, the second level has disabled echo.
+-- verify this.
+SELECT 456;
+

pkg/cli/testdata/i_twolevels2.sql

@@ -0,0 +1,7 @@
+SELECT 123;
+
+-- check the relative include
+\ir i_twolevels3.sql
+
+-- check that disabling echo here spills into the first level.
+\unset echo

pkg/cli/testdata/i_twolevels3.sql

@@ -0,0 +1 @@
+SELECT 789;