Add --prompt2, --no-prompt2

cloudspannerecosystem · Aug 21, 2024 · 21c0258 · 21c0258
1 parent a0ff8ae
commit 21c0258
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/cli.go b/cli.go
@@ -44,6 +44,7 @@ const (
 	DisplayModeTab
 
 	defaultPrompt      = `spanner\t> `
+	defaultPrompt2     = `-> `
 	defaultHistoryFile = `/tmp/spanner_cli_readline.tmp`
 
 	exitCodeSuccess = 0
@@ -60,6 +61,7 @@ var (
 type Cli struct {
 	Session     *Session
 	Prompt      string
+	Prompt2     string
 	HistoryFile string
 	Credential  []byte
 	InStream    io.ReadCloser
@@ -75,7 +77,7 @@ type command struct {
 	Vertical bool
 }
 
-func NewCli(projectId, instanceId, databaseId, prompt, historyFile string, credential []byte, inStream io.ReadCloser, outStream io.Writer, errStream io.Writer, verbose bool, priority pb.RequestOptions_Priority, role string, endpoint string, directedRead *pb.DirectedReadOptions) (*Cli, error) {
+func NewCli(projectId, instanceId, databaseId, prompt, prompt2, historyFile string, credential []byte, inStream io.ReadCloser, outStream io.Writer, errStream io.Writer, verbose bool, priority pb.RequestOptions_Priority, role string, endpoint string, directedRead *pb.DirectedReadOptions, noPrompt2 bool) (*Cli, error) {
 	session, err := createSession(projectId, instanceId, databaseId, credential, priority, role, endpoint, directedRead)
 	if err != nil {
 		return nil, err
@@ -85,13 +87,20 @@ func NewCli(projectId, instanceId, databaseId, prompt, historyFile string, crede
 		prompt = defaultPrompt
 	}
 
+	if noPrompt2 {
+		prompt2 = ""
+	} else if prompt2 == "" {
+		prompt2 = defaultPrompt2
+	}
+
 	if historyFile == "" {
 		historyFile = defaultHistoryFile
 	}
 
 	return &Cli{
 		Session:     session,
 		Prompt:      prompt,
+		Prompt2:     prompt2,
 		HistoryFile: historyFile,
 		Credential:  credential,
 		InStream:    inStream,
@@ -125,7 +134,7 @@ func (c *Cli) RunInteractive() int {
 		prompt := c.getInterpolatedPrompt()
 		rl.SetPrompt(prompt)
 
-		input, err := readInteractiveInput(rl, prompt)
+		input, err := readInteractiveInput(rl, prompt, c.Prompt2)
 		if err == io.EOF {
 			return c.Exit()
 		}
@@ -321,7 +330,7 @@ func createSession(projectId string, instanceId string, databaseId string, crede
 	return NewSession(projectId, instanceId, databaseId, priority, role, directedRead, opts...)
 }
 
-func readInteractiveInput(rl *readline.Instance, prompt string) (*inputStatement, error) {
+func readInteractiveInput(rl *readline.Instance, prompt string, prompt2 string) (*inputStatement, error) {
 	defer rl.SetPrompt(prompt)
 
 	var input string
@@ -347,10 +356,11 @@ func readInteractiveInput(rl *readline.Instance, prompt string) (*inputStatement
 
 		// show prompt to urge next input
 		var margin string
-		if l := len(prompt); l >= 3 {
-			margin = strings.Repeat(" ", l-3)
+		if l := len(prompt); l >= len(prompt2) {
+			margin = strings.Repeat(" ", l-len(prompt2))
 		}
-		rl.SetPrompt(margin + "-> ")
+		// rl.SetPrompt(margin + "-> ")
+		rl.SetPrompt(margin + prompt2)
 	}
 }
 

diff --git a/cli_test.go b/cli_test.go
@@ -172,7 +172,7 @@ func TestReadInteractiveInput(t *testing.T) {
 				t.Fatalf("unexpected readline.NewEx() error: %v", err)
 			}
 
-			got, err := readInteractiveInput(rl, "")
+			got, err := readInteractiveInput(rl, "", "")
 			if err != nil && !tt.wantError {
 				t.Errorf("readInteractiveInput(%q) got error: %v", tt.input, err)
 			}

diff --git a/main.go b/main.go
@@ -42,6 +42,8 @@ type spannerOptions struct {
 	Verbose      bool   `short:"v" long:"verbose" description:"Display verbose output."`
 	Credential   string `long:"credential" description:"Use the specific credential file"`
 	Prompt       string `long:"prompt" description:"Set the prompt to the specified format"`
+	Prompt2      string `long:"prompt2" description:"Set the prompt2 to the specified format"`
+	NoPrompt2    bool   `long:"no-prompt2" description:"Set the prompt2 to empty"`
 	HistoryFile  string `long:"history" description:"Set the history file to the specified path"`
 	Priority     string `long:"priority" description:"Set default request priority (HIGH|MEDIUM|LOW)"`
 	Role         string `long:"role" description:"Use the specific database role"`
@@ -96,7 +98,7 @@ func main() {
 		}
 	}
 
-	cli, err := NewCli(opts.ProjectId, opts.InstanceId, opts.DatabaseId, opts.Prompt, opts.HistoryFile, cred, os.Stdin, os.Stdout, os.Stderr, opts.Verbose, priority, opts.Role, opts.Endpoint, directedRead)
+	cli, err := NewCli(opts.ProjectId, opts.InstanceId, opts.DatabaseId, opts.Prompt, opts.Prompt2, opts.HistoryFile, cred, os.Stdin, os.Stdout, os.Stderr, opts.Verbose, priority, opts.Role, opts.Endpoint, directedRead, opts.NoPrompt2)
 	if err != nil {
 		exitf("Failed to connect to Spanner: %v", err)
 	}