feat(spanner): add directed reads feature

go.mod

@@ -3,7 +3,6 @@ module cloud.google.com/go
 go 1.19
 
 require (
-	cloud.google.com/go/iam v1.1.3
 	cloud.google.com/go/storage v1.30.1
 	github.com/golang/protobuf v1.5.3
 	github.com/google/go-cmp v0.6.0
@@ -25,6 +24,7 @@ require (
 require (
 	cloud.google.com/go/compute v1.23.1 // indirect
 	cloud.google.com/go/compute/metadata v0.2.3 // indirect
+	cloud.google.com/go/iam v1.1.3 // indirect
 	github.com/go-logr/logr v1.2.4 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect

spanner/batch.go

@@ -151,14 +151,15 @@ func (t *BatchReadOnlyTransaction) PartitionReadUsingIndexWithOptions(ctx contex
 	}
 	// Prepare ReadRequest.
 	req := &sppb.ReadRequest{
-		Session:          sid,
-		Transaction:      ts,
-		Table:            table,
-		Index:            index,
-		Columns:          columns,
-		KeySet:           kset,
-		RequestOptions:   createRequestOptions(readOptions.Priority, readOptions.RequestTag, ""),
-		DataBoostEnabled: readOptions.DataBoostEnabled,
+		Session:             sid,
+		Transaction:         ts,
+		Table:               table,
+		Index:               index,
+		Columns:             columns,
+		KeySet:              kset,
+		RequestOptions:      createRequestOptions(readOptions.Priority, readOptions.RequestTag, ""),
+		DataBoostEnabled:    readOptions.DataBoostEnabled,
+		DirectedReadOptions: readOptions.DirectedReadOptions,
 	}
 	// Generate partitions.
 	for _, p := range resp.GetPartitions() {
@@ -215,14 +216,15 @@ func (t *BatchReadOnlyTransaction) partitionQuery(ctx context.Context, statement
 
 	// prepare ExecuteSqlRequest
 	r := &sppb.ExecuteSqlRequest{
-		Session:          sid,
-		Transaction:      ts,
-		Sql:              statement.SQL,
-		Params:           params,
-		ParamTypes:       paramTypes,
-		QueryOptions:     qOpts.Options,
-		RequestOptions:   createRequestOptions(qOpts.Priority, qOpts.RequestTag, ""),
-		DataBoostEnabled: qOpts.DataBoostEnabled,
+		Session:             sid,
+		Transaction:         ts,
+		Sql:                 statement.SQL,
+		Params:              params,
+		ParamTypes:          paramTypes,
+		QueryOptions:        qOpts.Options,
+		RequestOptions:      createRequestOptions(qOpts.Priority, qOpts.RequestTag, ""),
+		DataBoostEnabled:    qOpts.DataBoostEnabled,
+		DirectedReadOptions: qOpts.DirectedReadOptions,
 	}
 
 	// generate Partitions
@@ -313,16 +315,17 @@ func (t *BatchReadOnlyTransaction) Execute(ctx context.Context, p *Partition) *R
 	if p.rreq != nil {
 		rpc = func(ctx context.Context, resumeToken []byte) (streamingReceiver, error) {
 			client, err := client.StreamingRead(ctx, &sppb.ReadRequest{
-				Session:          p.rreq.Session,
-				Transaction:      p.rreq.Transaction,
-				Table:            p.rreq.Table,
-				Index:            p.rreq.Index,
-				Columns:          p.rreq.Columns,
-				KeySet:           p.rreq.KeySet,
-				PartitionToken:   p.pt,
-				RequestOptions:   p.rreq.RequestOptions,
-				ResumeToken:      resumeToken,
-				DataBoostEnabled: p.rreq.DataBoostEnabled,
+				Session:             p.rreq.Session,
+				Transaction:         p.rreq.Transaction,
+				Table:               p.rreq.Table,
+				Index:               p.rreq.Index,
+				Columns:             p.rreq.Columns,
+				KeySet:              p.rreq.KeySet,
+				PartitionToken:      p.pt,
+				RequestOptions:      p.rreq.RequestOptions,
+				ResumeToken:         resumeToken,
+				DataBoostEnabled:    p.rreq.DataBoostEnabled,
+				DirectedReadOptions: p.rreq.DirectedReadOptions,
 			})
 			if err != nil {
 				return client, err
@@ -338,16 +341,17 @@ func (t *BatchReadOnlyTransaction) Execute(ctx context.Context, p *Partition) *R
 	} else {
 		rpc = func(ctx context.Context, resumeToken []byte) (streamingReceiver, error) {
 			client, err := client.ExecuteStreamingSql(ctx, &sppb.ExecuteSqlRequest{
-				Session:          p.qreq.Session,
-				Transaction:      p.qreq.Transaction,
-				Sql:              p.qreq.Sql,
-				Params:           p.qreq.Params,
-				ParamTypes:       p.qreq.ParamTypes,
-				QueryOptions:     p.qreq.QueryOptions,
-				PartitionToken:   p.pt,
-				RequestOptions:   p.qreq.RequestOptions,
-				ResumeToken:      resumeToken,
-				DataBoostEnabled: p.qreq.DataBoostEnabled,
+				Session:             p.qreq.Session,
+				Transaction:         p.qreq.Transaction,
+				Sql:                 p.qreq.Sql,
+				Params:              p.qreq.Params,
+				ParamTypes:          p.qreq.ParamTypes,
+				QueryOptions:        p.qreq.QueryOptions,
+				PartitionToken:      p.pt,
+				RequestOptions:      p.qreq.RequestOptions,
+				ResumeToken:         resumeToken,
+				DataBoostEnabled:    p.qreq.DataBoostEnabled,
+				DirectedReadOptions: p.qreq.DirectedReadOptions,
 			})
 			if err != nil {
 				return client, err

spanner/client.go

@@ -103,6 +103,7 @@ type Client struct {
 	bwo                  BatchWriteOptions
 	ct                   *commonTags
 	disableRouteToLeader bool
+	dro                  *sppb.DirectedReadOptions
 }
 
 // DatabaseName returns the full name of a database, e.g.,
@@ -186,6 +187,11 @@ type ClientConfig struct {
 
 	// BatchTimeout specifies the timeout for a batch of sessions managed sessionClient.
 	BatchTimeout time.Duration
+
+	// ClientConfig options used to set the DirectedReadOptions for all ReadRequests
+	// and ExecuteSqlRequests for the Client which indicate which replicas or regions
+	// should be used for non-transactional reads or queries.
+	DirectedReadOptions *sppb.DirectedReadOptions
 }
 
 func contextWithOutgoingMetadata(ctx context.Context, md metadata.MD, disableRouteToLeader bool) context.Context {
@@ -270,6 +276,7 @@ func NewClientWithConfig(ctx context.Context, database string, config ClientConf
 	if config.Compression == gzip.Name {
 		md.Append(requestsCompressionHeader, gzip.Name)
 	}
+
 	// Create a session client.
 	sc := newSessionClient(pool, database, config.UserAgent, sessionLabels, config.DatabaseRole, config.DisableRouteToLeader, md, config.BatchTimeout, config.Logger, config.CallOptions)
 
@@ -291,6 +298,7 @@ func NewClientWithConfig(ctx context.Context, database string, config ClientConf
 		bwo:                  config.BatchWriteOptions,
 		ct:                   getCommonTags(sc),
 		disableRouteToLeader: config.DisableRouteToLeader,
+		dro:                  config.DirectedReadOptions,
 	}
 	return c, nil
 }
@@ -374,6 +382,8 @@ func (c *Client) Single() *ReadOnlyTransaction {
 		t.sh = sh
 		return nil
 	}
+	t.txReadOnly.qo.DirectedReadOptions = c.dro
+	t.txReadOnly.ro.DirectedReadOptions = c.dro
 	t.ct = c.ct
 	return t
 }
@@ -397,6 +407,8 @@ func (c *Client) ReadOnlyTransaction() *ReadOnlyTransaction {
 	t.txReadOnly.qo = c.qo
 	t.txReadOnly.ro = c.ro
 	t.txReadOnly.disableRouteToLeader = true
+	t.txReadOnly.qo.DirectedReadOptions = c.dro
+	t.txReadOnly.ro.DirectedReadOptions = c.dro
 	t.ct = c.ct
 	return t
 }