feat(spanner): set client wide ReadOptions, ApplyOptions, and Transac…

…tionOptions (#6486)

* feat(spanner): set client wide ReadOptions, ApplyOptions, and TransactionOptions

* Add missing unit tests

Co-authored-by: rahul2393 <rahulyadavsep92@gmail.com>

spanner/batch.go

@@ -104,13 +104,13 @@ func (t *BatchReadOnlyTransaction) PartitionRead(ctx context.Context, table stri
 // can be configured using PartitionOptions. Pass a ReadOptions to modify the
 // read operation.
 func (t *BatchReadOnlyTransaction) PartitionReadWithOptions(ctx context.Context, table string, keys KeySet, columns []string, opt PartitionOptions, readOptions ReadOptions) ([]*Partition, error) {
-	return t.PartitionReadUsingIndexWithOptions(ctx, table, "", keys, columns, opt, readOptions)
+	return t.PartitionReadUsingIndexWithOptions(ctx, table, "", keys, columns, opt, t.ReadOnlyTransaction.txReadOnly.ro.merge(readOptions))
 }
 
 // PartitionReadUsingIndex returns a list of Partitions that can be used to read
 // rows from the database using an index.
 func (t *BatchReadOnlyTransaction) PartitionReadUsingIndex(ctx context.Context, table, index string, keys KeySet, columns []string, opt PartitionOptions) ([]*Partition, error) {
-	return t.PartitionReadUsingIndexWithOptions(ctx, table, index, keys, columns, opt, ReadOptions{})
+	return t.PartitionReadUsingIndexWithOptions(ctx, table, index, keys, columns, opt, t.ReadOnlyTransaction.txReadOnly.ro)
 }
 
 // PartitionReadUsingIndexWithOptions returns a list of Partitions that can be

spanner/batch_test.go

@@ -22,8 +22,9 @@ import (
 	"testing"
 	"time"
 
-	. "cloud.google.com/go/spanner/internal/testutil"
 	sppb "google.golang.org/genproto/googleapis/spanner/v1"
+
+	. "cloud.google.com/go/spanner/internal/testutil"
 )
 
 func TestPartitionRoundTrip(t *testing.T) {
@@ -120,6 +121,77 @@ func TestPartitionQuery_QueryOptions(t *testing.T) {
 	}
 }
 
+func TestPartitionQuery_ReadOptions(t *testing.T) {
+	testcases := []ReadOptionsTestCase{
+		{
+			name:   "Client level",
+			client: &ReadOptions{Index: "testIndex", Limit: 100, Priority: sppb.RequestOptions_PRIORITY_LOW, RequestTag: "testRequestTag"},
+			// Index and Limit are always ignored
+			want: &ReadOptions{Index: "", Limit: 0, Priority: sppb.RequestOptions_PRIORITY_LOW, RequestTag: "testRequestTag"},
+		},
+		{
+			name:   "Read level",
+			client: &ReadOptions{},
+			read:   &ReadOptions{Index: "testIndex", Limit: 100, Priority: sppb.RequestOptions_PRIORITY_LOW, RequestTag: "testRequestTag"},
+			// Index and Limit are always ignored
+			want: &ReadOptions{Index: "", Limit: 0, Priority: sppb.RequestOptions_PRIORITY_LOW, RequestTag: "testRequestTag"},
+		},
+		{
+			name:   "Read level has precedence than client level",
+			client: &ReadOptions{Index: "clientIndex", Limit: 10, Priority: sppb.RequestOptions_PRIORITY_LOW, RequestTag: "clientRequestTag"},
+			read:   &ReadOptions{Index: "readIndex", Limit: 20, Priority: sppb.RequestOptions_PRIORITY_MEDIUM, RequestTag: "readRequestTag"},
+			// Index and Limit are always ignored
+			want: &ReadOptions{Index: "", Limit: 0, Priority: sppb.RequestOptions_PRIORITY_MEDIUM, RequestTag: "readRequestTag"},
+		},
+	}
+
+	for _, tt := range testcases {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			_, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{ReadOptions: *tt.client})
+			defer teardown()
+
+			var (
+				err error
+				txn *BatchReadOnlyTransaction
+				ps  []*Partition
+			)
+
+			if txn, err = client.BatchReadOnlyTransaction(ctx, StrongRead()); err != nil {
+				t.Fatal(err)
+			}
+			defer txn.Cleanup(ctx)
+
+			if tt.read == nil {
+				ps, err = txn.PartitionRead(ctx, "Albums", KeySets(Key{"foo"}), []string{"SingerId", "AlbumId", "AlbumTitle"}, PartitionOptions{0, 3})
+			} else {
+				ps, err = txn.PartitionReadWithOptions(ctx, "Albums", KeySets(Key{"foo"}), []string{"SingerId", "AlbumId", "AlbumTitle"}, PartitionOptions{0, 3}, *tt.read)
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			for _, p := range ps {
+				req := p.rreq
+				if got, want := req.Index, tt.want.Index; got != want {
+					t.Fatalf("Incorrect index: got %v, want %v", got, want)
+				}
+				if got, want := req.Limit, int64(tt.want.Limit); got != want {
+					t.Fatalf("Incorrect limit: got %v, want %v", got, want)
+				}
+
+				ro := req.RequestOptions
+				if got, want := ro.Priority, tt.want.Priority; got != want {
+					t.Fatalf("Incorrect priority: got %v, want %v", got, want)
+				}
+				if got, want := ro.RequestTag, tt.want.RequestTag; got != want {
+					t.Fatalf("Incorrect request tag: got %v, want %v", got, want)
+				}
+			}
+		})
+	}
+}
+
 func TestPartitionQuery_Parallel(t *testing.T) {
 	ctx := context.Background()
 	server, client, teardown := setupMockedTestServer(t)

spanner/client.go

@@ -25,8 +25,6 @@ import (
 	"time"
 
 	"cloud.google.com/go/internal/trace"
-	vkit "cloud.google.com/go/spanner/apiv1"
-	"cloud.google.com/go/spanner/internal"
 	"google.golang.org/api/option"
 	"google.golang.org/api/option/internaloption"
 	gtransport "google.golang.org/api/transport/grpc"
@@ -35,6 +33,9 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 
+	vkit "cloud.google.com/go/spanner/apiv1"
+	"cloud.google.com/go/spanner/internal"
+
 	// Install google-c2p resolver, which is required for direct path.
 	_ "google.golang.org/grpc/xds/googledirectpath"
 	// Install RLS load balancer policy, which is needed for gRPC RLS.
@@ -86,6 +87,9 @@ type Client struct {
 	idleSessions *sessionPool
 	logger       *log.Logger
 	qo           QueryOptions
+	ro           ReadOptions
+	ao           []ApplyOption
+	txo          TransactionOptions
 	ct           *commonTags
 }
 
@@ -117,6 +121,15 @@ type ClientConfig struct {
 	// QueryOptions is the configuration for executing a sql query.
 	QueryOptions QueryOptions
 
+	// ReadOptions is the configuration for reading rows from a database
+	ReadOptions ReadOptions
+
+	// ApplyOptions is the configuration for applying
+	ApplyOptions []ApplyOption
+
+	// TransactionOptions is the configuration for a transaction.
+	TransactionOptions TransactionOptions
+
 	// CallOptions is the configuration for providing custom retry settings that
 	// override the default values.
 	CallOptions *vkit.CallOptions
@@ -211,6 +224,9 @@ func NewClientWithConfig(ctx context.Context, database string, config ClientConf
 		idleSessions: sp,
 		logger:       config.logger,
 		qo:           getQueryOptions(config.QueryOptions),
+		ro:           config.ReadOptions,
+		ao:           config.ApplyOptions,
+		txo:          config.TransactionOptions,
 		ct:           getCommonTags(sc),
 	}
 	return c, nil
@@ -273,6 +289,7 @@ func (c *Client) Single() *ReadOnlyTransaction {
 	t.txReadOnly.sp = c.idleSessions
 	t.txReadOnly.txReadEnv = t
 	t.txReadOnly.qo = c.qo
+	t.txReadOnly.ro = c.ro
 	t.txReadOnly.replaceSessionFunc = func(ctx context.Context) error {
 		if t.sh == nil {
 			return spannerErrorf(codes.InvalidArgument, "missing session handle on transaction")
@@ -309,6 +326,7 @@ func (c *Client) ReadOnlyTransaction() *ReadOnlyTransaction {
 	t.txReadOnly.sp = c.idleSessions
 	t.txReadOnly.txReadEnv = t
 	t.txReadOnly.qo = c.qo
+	t.txReadOnly.ro = c.ro
 	t.ct = c.ct
 	return t
 }
@@ -378,6 +396,7 @@ func (c *Client) BatchReadOnlyTransaction(ctx context.Context, tb TimestampBound
 	t.txReadOnly.sh = sh
 	t.txReadOnly.txReadEnv = t
 	t.txReadOnly.qo = c.qo
+	t.txReadOnly.ro = c.ro
 	t.ct = c.ct
 	return t, nil
 }
@@ -406,6 +425,7 @@ func (c *Client) BatchReadOnlyTransactionFromID(tid BatchReadOnlyTransactionID)
 	t.txReadOnly.sh = sh
 	t.txReadOnly.txReadEnv = t
 	t.txReadOnly.qo = c.qo
+	t.txReadOnly.ro = c.ro
 	t.ct = c.ct
 	return t
 }
@@ -491,7 +511,8 @@ func (c *Client) rwTransaction(ctx context.Context, f func(context.Context, *Rea
 		t.txReadOnly.sh = sh
 		t.txReadOnly.txReadEnv = t
 		t.txReadOnly.qo = c.qo
-		t.txOpts = options
+		t.txReadOnly.ro = c.ro
+		t.txOpts = c.txo.merge(options)
 		t.ct = c.ct
 
 		trace.TracePrintf(ctx, map[string]interface{}{"transactionID": string(sh.getTransactionID())},
@@ -555,6 +576,11 @@ func Priority(priority sppb.RequestOptions_Priority) ApplyOption {
 // Apply applies a list of mutations atomically to the database.
 func (c *Client) Apply(ctx context.Context, ms []*Mutation, opts ...ApplyOption) (commitTimestamp time.Time, err error) {
 	ao := &applyOption{}
+
+	for _, opt := range c.ao {
+		opt(ao)
+	}
+
 	for _, opt := range opts {
 		opt(ao)
 	}