Add option for setting a custom API endpoint (#162)

Useful for testing.
cloudspannerecosystem · Feb 19, 2024 · 4ad906e · 4ad906e
1 parent a80699f
commit 4ad906e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 9 deletions.
diff --git a/cli.go b/cli.go
@@ -67,15 +67,16 @@ type Cli struct {
 	ErrStream   io.Writer
 	Verbose     bool
 	Priority    pb.RequestOptions_Priority
+	Endpoint    string
 }
 
 type command struct {
 	Stmt     Statement
 	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) (*Cli, error) {
-	session, err := createSession(projectId, instanceId, databaseId, credential, priority, role)
+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) (*Cli, error) {
+	session, err := createSession(projectId, instanceId, databaseId, credential, priority, role, endpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -97,6 +98,7 @@ func NewCli(projectId, instanceId, databaseId, prompt, historyFile string, crede
 		OutStream:   outStream,
 		ErrStream:   errStream,
 		Verbose:     verbose,
+		Endpoint:    endpoint,
 	}, nil
 }
 
@@ -146,7 +148,7 @@ func (c *Cli) RunInteractive() int {
 		}
 
 		if s, ok := stmt.(*UseStatement); ok {
-			newSession, err := createSession(c.Session.projectId, c.Session.instanceId, s.Database, c.Credential, c.Priority, s.Role)
+			newSession, err := createSession(c.Session.projectId, c.Session.instanceId, s.Database, c.Credential, c.Priority, s.Role, c.Endpoint)
 			if err != nil {
 				c.PrintInteractiveError(err)
 				continue
@@ -308,13 +310,15 @@ func (c *Cli) getInterpolatedPrompt() string {
 	return prompt
 }
 
-func createSession(projectId string, instanceId string, databaseId string, credential []byte, priority pb.RequestOptions_Priority, role string) (*Session, error) {
+func createSession(projectId string, instanceId string, databaseId string, credential []byte, priority pb.RequestOptions_Priority, role string, endpoint string) (*Session, error) {
+	var opts []option.ClientOption
 	if credential != nil {
-		credentialOption := option.WithCredentialsJSON(credential)
-		return NewSession(projectId, instanceId, databaseId, priority, role, credentialOption)
-	} else {
-		return NewSession(projectId, instanceId, databaseId, priority, role)
+		opts = append(opts, option.WithCredentialsJSON(credential))
+	}
+	if endpoint != "" {
+		opts = append(opts, option.WithEndpoint(endpoint))
 	}
+	return NewSession(projectId, instanceId, databaseId, priority, role, opts...)
 }
 
 func readInteractiveInput(rl *readline.Instance, prompt string) (*inputStatement, error) {

diff --git a/main.go b/main.go
@@ -45,6 +45,7 @@ type spannerOptions struct {
 	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"`
+	Endpoint    string `long:"endpoint" description:"Set the Spanner API endpoint (host:port)"`
 }
 
 func main() {
@@ -85,7 +86,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)
+	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)
 	if err != nil {
 		exitf("Failed to connect to Spanner: %v", err)
 	}