Merge pull request #8363 from dolthub/macneale4/shell-edit

\edit support in `dolt sql` shell

.github/workflows/ci-check-repo.yaml

@@ -81,11 +81,11 @@ jobs:
       - name: Build go deps tool
         working-directory: go/utils/3pdeps
         run: go build .
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
           name: update-godeps-tool
           path: go/utils/3pdeps/3pdeps
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
           name: format-code-script
           path: go/utils/repofmt/_format_repo.sh
@@ -110,11 +110,11 @@ jobs:
         working-directory: ./go
       - name: Install goimports
         run: go install golang.org/x/tools/cmd/goimports@latest
-      - uses: actions/download-artifact@v4.1.7
+      - uses: actions/download-artifact@v4
         with:
           name: format-code-script
           path: go/utils/repofmt
-      - uses: actions/download-artifact@v4.1.7
+      - uses: actions/download-artifact@v4
         with:
           name: update-godeps-tool
           path: go

go/cmd/dolt/commands/commit.go

@@ -38,7 +38,6 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/env/actions"
 	"github.com/dolthub/dolt/go/libraries/utils/argparser"
 	"github.com/dolthub/dolt/go/libraries/utils/config"
-	"github.com/dolthub/dolt/go/libraries/utils/editor"
 	"github.com/dolthub/dolt/go/libraries/utils/iohelp"
 	"github.com/dolthub/dolt/go/libraries/utils/set"
 	"github.com/dolthub/dolt/go/store/datas"
@@ -347,14 +346,6 @@ func handleCommitErr(sqlCtx *sql.Context, queryist cli.Queryist, err error, usag
 // getCommitMessageFromEditor opens editor to ask user for commit message if none defined from command line.
 // suggestedMsg will be returned if no-edit flag is defined or if this function was called from sql dolt_merge command.
 func getCommitMessageFromEditor(sqlCtx *sql.Context, queryist cli.Queryist, suggestedMsg, amendString string, noEdit bool, cliCtx cli.CliContext) (string, error) {
-	if cli.ExecuteWithStdioRestored == nil || noEdit {
-		return suggestedMsg, nil
-	}
-
-	if !checkIsTerminal() {
-		return suggestedMsg, nil
-	}
-
 	var finalMsg string
 	initialMsg, err := buildInitalCommitMsg(sqlCtx, queryist, suggestedMsg)
 	if err != nil {
@@ -364,26 +355,20 @@ func getCommitMessageFromEditor(sqlCtx *sql.Context, queryist cli.Queryist, sugg
 		initialMsg = fmt.Sprintf("%s\n%s", amendString, initialMsg)
 	}
 
-	backupEd := "vim"
-	// try getting default editor on the user system
-	if ed, edSet := os.LookupEnv(dconfig.EnvEditor); edSet {
-		backupEd = ed
+	if cli.ExecuteWithStdioRestored == nil || noEdit {
+		return suggestedMsg, nil
 	}
-	// try getting Dolt config core.editor
-	editorStr := cliCtx.Config().GetStringOrDefault(config.DoltEditor, backupEd)
 
-	cli.ExecuteWithStdioRestored(func() {
-		commitMsg, cErr := editor.OpenTempEditor(editorStr, initialMsg)
-		if cErr != nil {
-			err = cErr
-		}
-		finalMsg = parseCommitMessage(commitMsg)
-	})
+	if !checkIsTerminal() {
+		return suggestedMsg, nil
+	}
 
+	commitMsg, err := execEditor(initialMsg, "", cliCtx)
 	if err != nil {
 		return "", fmt.Errorf("Failed to open commit editor: %v \n Check your `EDITOR` environment variable with `echo $EDITOR` or your dolt config with `dolt config --list` to ensure that your editor is valid", err)
 	}
 
+	finalMsg = parseCommitMessage(commitMsg)
 	return finalMsg, nil
 }
 

go/cmd/dolt/commands/rebase.go

@@ -230,7 +230,7 @@ func getRebasePlan(cliCtx cli.CliContext, sqlCtx *sql.Context, queryist cli.Quer
 
 	var rebaseMsg string
 	cli.ExecuteWithStdioRestored(func() {
-		rebaseMsg, err = editor.OpenTempEditor(editorStr, initialRebaseMsg)
+		rebaseMsg, err = editor.OpenTempEditor(editorStr, initialRebaseMsg, "")
 	})
 	if err != nil {
 		return nil, err

go/cmd/dolt/commands/sql.go

@@ -749,23 +749,17 @@ func execShell(sqlCtx *sql.Context, qryist cli.Queryist, format engine.PrintResu
 
 	initialCtx := sqlCtx.Context
 
+	// Used for the \edit command.
+	lastSqlCmd := ""
+
 	shell.Uninterpreted(func(c *ishell.Context) {
 		query := c.Args[0]
 		query = strings.TrimSpace(query)
 		if len(query) == 0 {
 			return
 		}
 
-		// TODO: there's a bug in the readline library when editing multi-line history entries.
-		// Longer term we need to switch to a new readline library, like in this bug:
-		// https://github.com/cockroachdb/cockroach/issues/15460
-		// For now, we store all history entries as single-line strings to avoid the issue.
-		singleLine := strings.ReplaceAll(query, "\n", " ")
-
-		if err := shell.AddHistory(singleLine); err != nil {
-			// TODO: handle better, like by turning off history writing for the rest of the session
-			shell.Println(color.RedString(err.Error()))
-		}
+		trackHistory(shell, query)
 
 		query = strings.TrimSuffix(query, shell.LineTerminator())
 
@@ -786,13 +780,23 @@ func execShell(sqlCtx *sql.Context, qryist cli.Queryist, format engine.PrintResu
 
 			sqlCtx := sql.NewContext(subCtx, sql.WithSession(sqlCtx.Session))
 
-			subCmd, foundCmd := isSlashQuery(query)
-			if foundCmd {
+			cmdType, subCmd, newQuery, err := preprocessQuery(query, lastSqlCmd, cliCtx)
+			if err != nil {
+				shell.Println(color.RedString(err.Error()))
+				return true
+			}
+
+			if cmdType == DoltCliCommand {
 				err := handleSlashCommand(sqlCtx, subCmd, query, cliCtx)
 				if err != nil {
 					shell.Println(color.RedString(err.Error()))
 				}
 			} else {
+				if cmdType == TransformCommand {
+					query = newQuery
+					trackHistory(shell, query+";")
+				}
+				lastSqlCmd = query
 				var sqlSch sql.Schema
 				var rowIter sql.RowIter
 				if sqlSch, rowIter, _, err = processQuery(sqlCtx, query, qryist); err != nil {
@@ -831,13 +835,55 @@ func execShell(sqlCtx *sql.Context, qryist cli.Queryist, format engine.PrintResu
 	return nil
 }
 
-func isSlashQuery(query string) (cli.Command, bool) {
+func trackHistory(shell *ishell.Shell, query string) {
+	// TODO: there's a bug in the readline library when editing multi-line history entries.
+	// Longer term we need to switch to a new readline library, like in this bug:
+	// https://github.com/cockroachdb/cockroach/issues/15460
+	// For now, we store all history entries as single-line strings to avoid the issue.
+	singleLine := strings.ReplaceAll(query, "\n", " ")
+
+	if err := shell.AddHistory(singleLine); err != nil {
+		// TODO: handle better, like by turning off history writing for the rest of the session
+		shell.Println(color.RedString(err.Error()))
+	}
+}
+
+type CommandType int
+
+// CommandType is used to determine how to handle a query. See preprocessQuery.
+const (
+	DoltCliCommand CommandType = iota
+	SqlShellCommand
+	TransformCommand
+)
+
+// preprocessQuery takes the user's query and returns the command type, the command, and the query to execute. The
+// CommandType returned is going to be used to determine how to handle the query.
+//   - DoltCliCommand: the cli.Command returned should be executed. Query string is empty, and should be ignored.
+//   - TransformCommand: The 'lastQuery' argument will be transformed into something else, using the EDITOR.
+//     The query returned will be the edited query, and should be entered into the user's command history. The cli.Command will be nil.
+//   - SqlShellCommand: cli.Command will be nil. The query returned will be identical to the query passed in.
+func preprocessQuery(query, lastQuery string, cliCtx cli.CliContext) (CommandType, cli.Command, string, error) {
 	// strip leading whitespace
 	query = strings.TrimLeft(query, " \t\n\r\v\f")
 	if strings.HasPrefix(query, "\\") {
-		return findSlashCmd(query[1:])
+		if query == "\\edit" {
+			// \edit is a special case. Maybe we'll generalize this in the future.
+			updatedQuery, err := execEditor(lastQuery, ".sql", cliCtx)
+			if err != nil {
+				return TransformCommand, nil, "", err
+			}
+			// Trailing newlines are common in editors, so may as well trim all whitespace.
+			updatedQuery = strings.TrimRight(updatedQuery, " \t\n\r\v\f")
+			return TransformCommand, nil, updatedQuery, nil
+		}
+
+		cmd, ok := findSlashCmd(query[1:])
+		if ok {
+			return DoltCliCommand, cmd, "", nil
+		}
 	}
-	return nil, false
+	return SqlShellCommand, nil, query, nil
 }
 
 // postCommandUpdate is a helper function that is run after the shell has completed a command. It updates the the database

go/cmd/dolt/commands/sql_slash.go

@@ -39,6 +39,7 @@ var slashCmds = []cli.Command{
 	BranchCmd{},
 	MergeCmd{},
 	SlashHelp{},
+	SlashEdit{},
 }
 
 // parseSlashCmd parses a command line string into a slice of strings, splitting on spaces, but allowing spaces within
@@ -68,6 +69,7 @@ func parseSlashCmd(cmd string) []string {
 	return cmdWords
 }
 
+// handleSlashCommand executes the command given by the fullCmd string. These are commands are direct calls to CLI commands.
 func handleSlashCommand(sqlCtx *sql.Context, subCmd cli.Command, fullCmd string, cliCtx cli.CliContext) error {
 	cliCmd := parseSlashCmd(fullCmd)
 	if len(cliCmd) == 0 {
@@ -109,7 +111,6 @@ func (s SlashHelp) Exec(ctx context.Context, _ string, args []string, _ *env.Dol
 		if ok {
 			foo, _ := cli.HelpAndUsagePrinters(subCmdInst.Docs())
 			foo()
-
 		} else {
 			cli.Println(fmt.Sprintf("Unknown command: %s", subCmd))
 		}
@@ -178,3 +179,39 @@ func findSlashCmd(cmd string) (cli.Command, bool) {
 	}
 	return nil, false
 }
+
+type SlashEdit struct{}
+
+var _ cli.Command = SlashEdit{}
+
+func (s SlashEdit) Name() string {
+	return "edit"
+}
+
+func (s SlashEdit) Description() string {
+	return "Use $EDITOR to edit the last command."
+}
+func (s SlashEdit) Exec(ctx context.Context, commandStr string, args []string, dEnv *env.DoltEnv, cliCtx cli.CliContext) int {
+
+	initialCmd := "select * from my_table;"
+
+	contents, err := execEditor(initialCmd, ".sql", cliCtx)
+	if err != nil {
+		cli.PrintErrln(err.Error())
+		return 1
+	}
+
+	cli.Printf("Edited command: %s", contents)
+
+	return 0
+}
+
+func (s SlashEdit) Docs() *cli.CommandDocumentation {
+	//TODO implement me
+	return &cli.CommandDocumentation{}
+}
+
+func (s SlashEdit) ArgParser() *argparser.ArgParser {
+	// No arguments.
+	return &argparser.ArgParser{}
+}

go/cmd/dolt/commands/utils.go

@@ -19,6 +19,7 @@ import (
 	"crypto/sha1"
 	"fmt"
 	"net"
+	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -33,11 +34,14 @@ import (
 	"github.com/dolthub/dolt/go/cmd/dolt/cli"
 	"github.com/dolthub/dolt/go/cmd/dolt/commands/engine"
 	"github.com/dolthub/dolt/go/cmd/dolt/errhand"
+	"github.com/dolthub/dolt/go/libraries/doltcore/dconfig"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env/actions"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
 	"github.com/dolthub/dolt/go/libraries/utils/argparser"
+	"github.com/dolthub/dolt/go/libraries/utils/config"
+	"github.com/dolthub/dolt/go/libraries/utils/editor"
 	"github.com/dolthub/dolt/go/libraries/utils/filesys"
 	"github.com/dolthub/dolt/go/store/datas"
 	"github.com/dolthub/dolt/go/store/util/outputpager"
@@ -916,3 +920,35 @@ func PrintStagingError(err error) {
 
 	cli.PrintErrln(vErr.Verbose())
 }
+
+// execEditor opens editor to ask user for input.
+func execEditor(initialMsg string, suffix string, cliCtx cli.CliContext) (editedMsg string, err error) {
+	if cli.ExecuteWithStdioRestored == nil {
+		return initialMsg, nil
+	}
+
+	if !checkIsTerminal() {
+		return initialMsg, nil
+	}
+
+	backupEd := "vim"
+	// try getting default editor on the user system
+	if ed, edSet := os.LookupEnv(dconfig.EnvEditor); edSet {
+		backupEd = ed
+	}
+	// try getting Dolt config core.editor
+	editorStr := cliCtx.Config().GetStringOrDefault(config.DoltEditor, backupEd)
+
+	cli.ExecuteWithStdioRestored(func() {
+		editedMsg, err = editor.OpenTempEditor(editorStr, initialMsg, suffix)
+		if err != nil {
+			return
+		}
+	})
+
+	if err != nil {
+		return "", fmt.Errorf("Failed to open commit editor: %v \n Check your `EDITOR` environment variable with `echo $EDITOR` or your dolt config with `dolt config --list` to ensure that your editor is valid", err)
+	}
+
+	return editedMsg, nil
+}

go/libraries/utils/editor/edit.go

@@ -25,8 +25,9 @@ import (
 )
 
 // OpenTempEditor allows user to write/edit message in temporary file
-func OpenTempEditor(ed string, initialContents string) (string, error) {
-	filename := filepath.Join(os.TempDir(), uuid.New().String())
+func OpenTempEditor(ed string, initialContents string, fileSuffix string) (string, error) {
+	fileName := uuid.New().String() + fileSuffix
+	filename := filepath.Join(os.TempDir(), fileName)
 	err := os.WriteFile(filename, []byte(initialContents), os.ModePerm)
 
 	if err != nil {

go/libraries/utils/editor/edit_test.go

@@ -59,7 +59,7 @@ func TestOpenCommitEditor(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		val, err := OpenTempEditor(test.editorStr, test.initialContents)
+		val, err := OpenTempEditor(test.editorStr, test.initialContents, "")
 
 		if err != nil {
 			t.Error(err)