Implement --proto-descriptor-file

README.md

@@ -37,24 +37,25 @@ Usage:
   spanner-cli [OPTIONS]
 
 spanner:
-  -p, --project=         (required) GCP Project ID. [$SPANNER_PROJECT_ID]
-  -i, --instance=        (required) Cloud Spanner Instance ID [$SPANNER_INSTANCE_ID]
-  -d, --database=        (required) Cloud Spanner Database ID. [$SPANNER_DATABASE_ID]
-  -e, --execute=         Execute SQL statement and quit.
-  -f, --file=            Execute SQL statement from file and quit.
-  -t, --table            Display output in table format for batch mode.
-  -v, --verbose          Display verbose output.
-      --credential=      Use the specific credential file
-      --prompt=          Set the prompt to the specified format
-      --history=         Set the history file to the specified path
-      --priority=        Set default request priority (HIGH|MEDIUM|LOW)
-      --role=            Use the specific database role
-      --endpoint=        Set the Spanner API endpoint (host:port)
-      --directed-read=   Directed read option (replica_location:replica_type). The replicat_type is optional and either READ_ONLY or READ_WRITE
-      --skip-tls-verify  Insecurely skip TLS verify
+  -p, --project=               (required) GCP Project ID. [$SPANNER_PROJECT_ID]
+  -i, --instance=              (required) Cloud Spanner Instance ID [$SPANNER_INSTANCE_ID]
+  -d, --database=              (required) Cloud Spanner Database ID. [$SPANNER_DATABASE_ID]
+  -e, --execute=               Execute SQL statement and quit.
+  -f, --file=                  Execute SQL statement from file and quit.
+  -t, --table                  Display output in table format for batch mode.
+  -v, --verbose                Display verbose output.
+      --credential=            Use the specific credential file
+      --prompt=                Set the prompt to the specified format
+      --history=               Set the history file to the specified path
+      --priority=              Set default request priority (HIGH|MEDIUM|LOW)
+      --role=                  Use the specific database role
+      --endpoint=              Set the Spanner API endpoint (host:port)
+      --directed-read=         Directed read option (replica_location:replica_type). The replicat_type is optional and either READ_ONLY or READ_WRITE
+      --skip-tls-verify        Insecurely skip TLS verify
+      --proto-descriptor-file= Path of a file that contains a protobuf-serialized google.protobuf.FileDescriptorSet message to use in this invocation.
 
 Help Options:
-  -h, --help             Show this help message
+  -h, --help                   Show this help message
 ```
 
 ### Authentication

cli.go

@@ -79,8 +79,11 @@ 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, skipTLSVerify bool) (*Cli, error) {
-	session, err := createSession(projectId, instanceId, databaseId, credential, priority, role, endpoint, directedRead, skipTLSVerify)
+func NewCli(projectId, instanceId, databaseId, prompt, historyFile string, credential []byte,
+	inStream io.ReadCloser, outStream, errStream io.Writer, verbose bool,
+	priority pb.RequestOptions_Priority, role, endpoint string, directedRead *pb.DirectedReadOptions,
+	skipTLSVerify bool, protoDescriptor []byte) (*Cli, error) {
+	session, err := createSession(projectId, instanceId, databaseId, credential, priority, role, endpoint, directedRead, skipTLSVerify, protoDescriptor)
 	if err != nil {
 		return nil, err
 	}
@@ -153,7 +156,8 @@ 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, c.Endpoint, c.Session.directedRead, c.SkipTLSVerify)
+			newSession, err := createSession(c.Session.projectId, c.Session.instanceId, s.Database, c.Credential, c.Priority,
+				s.Role, c.Endpoint, c.Session.directedRead, c.SkipTLSVerify, c.Session.protoDescriptor)
 			if err != nil {
 				c.PrintInteractiveError(err)
 				continue
@@ -315,7 +319,9 @@ func (c *Cli) getInterpolatedPrompt() string {
 	return prompt
 }
 
-func createSession(projectId string, instanceId string, databaseId string, credential []byte, priority pb.RequestOptions_Priority, role string, endpoint string, directedRead *pb.DirectedReadOptions, skipTLSVerify bool) (*Session, error) {
+func createSession(projectId string, instanceId string, databaseId string, credential []byte,
+	priority pb.RequestOptions_Priority, role string, endpoint string, directedRead *pb.DirectedReadOptions,
+	skipTLSVerify bool, protoDescriptor []byte) (*Session, error) {
 	var opts []option.ClientOption
 	if credential != nil {
 		opts = append(opts, option.WithCredentialsJSON(credential))
@@ -327,7 +333,7 @@ func createSession(projectId string, instanceId string, databaseId string, crede
 		creds := credentials.NewTLS(&tls.Config{InsecureSkipVerify: true})
 		opts = append(opts, option.WithGRPCDialOption(grpc.WithTransportCredentials(creds)))
 	}
-	return NewSession(projectId, instanceId, databaseId, priority, role, directedRead, opts...)
+	return NewSession(projectId, instanceId, databaseId, priority, role, directedRead, protoDescriptor, opts...)
 }
 
 func readInteractiveInput(rl *readline.Instance, prompt string) (*inputStatement, error) {

go.mod

@@ -7,7 +7,6 @@ require (
 	cloud.google.com/go/spanner v1.62.0
 	github.com/apstndb/gsqlsep v0.0.0-20230324124551-0e8335710080
 	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e
-	github.com/davecgh/go-spew v1.1.1
 	github.com/google/go-cmp v0.6.0
 	github.com/jessevdk/go-flags v1.4.0
 	github.com/olekukonko/tablewriter v0.0.5

integration_test.go

@@ -86,7 +86,7 @@ func setup(t *testing.T, ctx context.Context, dmls []string) (*Session, string,
 	if testCredential != "" {
 		options = append(options, option.WithCredentialsJSON([]byte(testCredential)))
 	}
-	session, err := NewSession(testProjectId, testInstanceId, testDatabaseId, pb.RequestOptions_PRIORITY_UNSPECIFIED, "", nil, options...)
+	session, err := NewSession(testProjectId, testInstanceId, testDatabaseId, pb.RequestOptions_PRIORITY_UNSPECIFIED, "", nil, nil, options...)
 	if err != nil {
 		t.Fatalf("failed to create test session: err=%s", err)
 	}

main.go

@@ -33,21 +33,22 @@ type globalOptions struct {
 }
 
 type spannerOptions struct {
-	ProjectId     string `short:"p" long:"project" env:"SPANNER_PROJECT_ID" description:"(required) GCP Project ID."`
-	InstanceId    string `short:"i" long:"instance" env:"SPANNER_INSTANCE_ID" description:"(required) Cloud Spanner Instance ID"`
-	DatabaseId    string `short:"d" long:"database" env:"SPANNER_DATABASE_ID" description:"(required) Cloud Spanner Database ID."`
-	Execute       string `short:"e" long:"execute" description:"Execute SQL statement and quit."`
-	File          string `short:"f" long:"file" description:"Execute SQL statement from file and quit."`
-	Table         bool   `short:"t" long:"table" description:"Display output in table format for batch mode."`
-	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"`
-	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)"`
-	DirectedRead  string `long:"directed-read" description:"Directed read option (replica_location:replica_type). The replicat_type is optional and either READ_ONLY or READ_WRITE"`
-	SkipTLSVerify bool   `long:"skip-tls-verify" description:"Insecurely skip TLS verify"`
+	ProjectId           string `short:"p" long:"project" env:"SPANNER_PROJECT_ID" description:"(required) GCP Project ID."`
+	InstanceId          string `short:"i" long:"instance" env:"SPANNER_INSTANCE_ID" description:"(required) Cloud Spanner Instance ID"`
+	DatabaseId          string `short:"d" long:"database" env:"SPANNER_DATABASE_ID" description:"(required) Cloud Spanner Database ID."`
+	Execute             string `short:"e" long:"execute" description:"Execute SQL statement and quit."`
+	File                string `short:"f" long:"file" description:"Execute SQL statement from file and quit."`
+	Table               bool   `short:"t" long:"table" description:"Display output in table format for batch mode."`
+	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"`
+	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)"`
+	DirectedRead        string `long:"directed-read" description:"Directed read option (replica_location:replica_type). The replicat_type is optional and either READ_ONLY or READ_WRITE"`
+	SkipTLSVerify       bool   `long:"skip-tls-verify" description:"Insecurely skip TLS verify"`
+	ProtoDescriptorFile string `long:"proto-descriptor-file" description:"Path of a file that contains a protobuf-serialized google.protobuf.FileDescriptorSet message to use in this invocation."`
 }
 
 func main() {
@@ -97,7 +98,19 @@ 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, opts.SkipTLSVerify)
+	// Don't need to unmarshal into descriptorpb.FileDescriptorSet because the UpdateDDL API just accepts []byte.
+	var protoDescriptor []byte
+	if opts.ProtoDescriptorFile != "" {
+		var err error
+		protoDescriptor, err = os.ReadFile(opts.ProtoDescriptorFile)
+		if err != nil {
+			exitf("Failed to read proto descriptor file: %v\n", err)
+		}
+	}
+
+	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,
+		opts.SkipTLSVerify, protoDescriptor)
 	if err != nil {
 		exitf("Failed to connect to Spanner: %v", err)
 	}

session.go

@@ -58,6 +58,7 @@ type Session struct {
 	directedRead    *pb.DirectedReadOptions
 	tc              *transactionContext
 	tcMutex         sync.Mutex // Guard a critical section for transaction.
+	protoDescriptor []byte
 }
 
 type transactionContext struct {
@@ -68,7 +69,8 @@ type transactionContext struct {
 	roTxn         *spanner.ReadOnlyTransaction
 }
 
-func NewSession(projectId string, instanceId string, databaseId string, priority pb.RequestOptions_Priority, role string, directedRead *pb.DirectedReadOptions, opts ...option.ClientOption) (*Session, error) {
+func NewSession(projectId string, instanceId string, databaseId string, priority pb.RequestOptions_Priority, role string, directedRead *pb.DirectedReadOptions,
+	protoDescriptor []byte, opts ...option.ClientOption) (*Session, error) {
 	ctx := context.Background()
 	dbPath := fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectId, instanceId, databaseId)
 	clientConfig := defaultClientConfig
@@ -99,6 +101,7 @@ func NewSession(projectId string, instanceId string, databaseId string, priority
 		adminClient:     adminClient,
 		defaultPriority: priority,
 		directedRead:    directedRead,
+		protoDescriptor: protoDescriptor,
 	}
 	go session.startHeartbeat()
 

session_test.go

@@ -66,7 +66,7 @@ func TestRequestPriority(t *testing.T) {
 		t.Run(test.desc, func(t *testing.T) {
 			defer recorder.flush()
 
-			session, err := NewSession("project", "instance", "database", test.sessionPriority, "role", nil, option.WithGRPCConn(conn))
+			session, err := NewSession("project", "instance", "database", test.sessionPriority, "role", nil, nil, option.WithGRPCConn(conn))
 			if err != nil {
 				t.Fatalf("failed to create spanner-cli session: %v", err)
 			}

statement.go

@@ -406,6 +406,8 @@ func executeDdlStatements(ctx context.Context, session *Session, ddls []string)
 	op, err := session.adminClient.UpdateDatabaseDdl(ctx, &adminpb.UpdateDatabaseDdlRequest{
 		Database:   session.DatabasePath(),
 		Statements: ddls,
+		// There is no problem to send ProtoDescriptors with any DDL statements
+		ProtoDescriptors: session.protoDescriptor,
 	})
 	if err != nil {
 		return nil, err