feat(spanner): add LockHint feature (googleapis#10382)

* feat(spanner): add LockHint feature * fix * revert TestClient_ReadWriteTransaction_Query_ReadOptions * Override lockhint in RO txn * Fix comment * fix reviews * fix review comment --------- Co-authored-by: Sri Harsha CH <[email protected]>
codyoss · Jun 27, 2024 · 64bdcb1 · 64bdcb1
1 parent b53c5fb
commit 64bdcb1
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 0 deletions.
diff --git a/spanner/client.go b/spanner/client.go
@@ -609,6 +609,7 @@ func (c *Client) Single() *ReadOnlyTransaction {
 	}
 	t.txReadOnly.qo.DirectedReadOptions = c.dro
 	t.txReadOnly.ro.DirectedReadOptions = c.dro
+	t.txReadOnly.ro.LockHint = sppb.ReadRequest_LOCK_HINT_UNSPECIFIED
 	t.ct = c.ct
 	t.otConfig = c.otConfig
 	return t
@@ -635,6 +636,7 @@ func (c *Client) ReadOnlyTransaction() *ReadOnlyTransaction {
 	t.txReadOnly.disableRouteToLeader = true
 	t.txReadOnly.qo.DirectedReadOptions = c.dro
 	t.txReadOnly.ro.DirectedReadOptions = c.dro
+	t.txReadOnly.ro.LockHint = sppb.ReadRequest_LOCK_HINT_UNSPECIFIED
 	t.ct = c.ct
 	t.otConfig = c.otConfig
 	return t
@@ -706,6 +708,7 @@ func (c *Client) BatchReadOnlyTransaction(ctx context.Context, tb TimestampBound
 	t.txReadOnly.disableRouteToLeader = true
 	t.txReadOnly.qo.DirectedReadOptions = c.dro
 	t.txReadOnly.ro.DirectedReadOptions = c.dro
+	t.txReadOnly.ro.LockHint = sppb.ReadRequest_LOCK_HINT_UNSPECIFIED
 	t.ct = c.ct
 	t.otConfig = c.otConfig
 	return t, nil
@@ -740,6 +743,7 @@ func (c *Client) BatchReadOnlyTransactionFromID(tid BatchReadOnlyTransactionID)
 	t.txReadOnly.disableRouteToLeader = true
 	t.txReadOnly.qo.DirectedReadOptions = c.dro
 	t.txReadOnly.ro.DirectedReadOptions = c.dro
+	t.txReadOnly.ro.LockHint = sppb.ReadRequest_LOCK_HINT_UNSPECIFIED
 	t.ct = c.ct
 	t.otConfig = c.otConfig
 	return t

diff --git a/spanner/client_test.go b/spanner/client_test.go
@@ -877,6 +877,9 @@ func checkReqsForReadOptions(t *testing.T, server InMemSpannerServer, ro ReadOpt
 	if got, want := sqlReq.OrderBy, ro.OrderBy; got != want {
 		t.Fatalf("OrderBy mismatch, got %v, want %v", got, want)
 	}
+	if got, want := sqlReq.LockHint, ro.LockHint; got != want {
+		t.Fatalf("LockHint mismatch, got %v, want %v", got, want)
+	}
 }
 
 func checkReqsForTransactionOptions(t *testing.T, server InMemSpannerServer, txo TransactionOptions) {
@@ -2143,6 +2146,89 @@ func TestClient_ReadWriteTransaction_Query_QueryOptions(t *testing.T) {
 	}
 }
 
+func TestClient_ReadWriteTransaction_LockHintOptions(t *testing.T) {
+	readOptionsTestCases := []ReadOptionsTestCase{
+		{
+			name:   "Client level",
+			client: &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+			want:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+		},
+		{
+			name:   "Read level",
+			client: &ReadOptions{},
+			read:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+			want:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+		},
+		{
+			name:   "Read level has precedence than client level",
+			client: &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_SHARED},
+			read:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+			want:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+		},
+		{
+			name:   "Client level has precendence when LOCK_HINT_UNSPECIFIED at read level",
+			client: &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+			read:   &ReadOptions{},
+			want:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+		},
+	}
+
+	for _, tt := range readOptionsTestCases {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{ReadOptions: *tt.client})
+			defer teardown()
+
+			_, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *ReadWriteTransaction) error {
+				var iter *RowIterator
+				if tt.read == nil {
+					iter = tx.Read(ctx, "Albums", KeySets(Key{"foo"}), []string{"SingerId", "AlbumId", "AlbumTitle"})
+				} else {
+					iter = tx.ReadWithOptions(ctx, "Albums", KeySets(Key{"foo"}), []string{"SingerId", "AlbumId", "AlbumTitle"}, tt.read)
+				}
+				testReadOptions(t, iter, server.TestSpanner, *tt.want)
+				return nil
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}
+
+func TestClient_ReadOnlyTransaction_LockHintOptions(t *testing.T) {
+	readOptionsTestCases := []ReadOptionsTestCase{
+		{
+			name:   "Client level Lock hint overiden in request level",
+			client: &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+			read:   &ReadOptions{},
+			want:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_UNSPECIFIED},
+		},
+		{
+			name:   "Request level",
+			client: &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_EXCLUSIVE},
+			read:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_SHARED},
+			want:   &ReadOptions{LockHint: sppb.ReadRequest_LOCK_HINT_SHARED},
+		},
+	}
+
+	for _, tt := range readOptionsTestCases {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{ReadOptions: *tt.client})
+			defer teardown()
+
+			for _, tx := range []*ReadOnlyTransaction{
+				client.Single(),
+				client.ReadOnlyTransaction(),
+			} {
+				iter := tx.ReadWithOptions(ctx, "Albums", KeySets(Key{"foo"}), []string{"SingerId", "AlbumId", "AlbumTitle"}, tt.read)
+				testReadOptions(t, iter, server.TestSpanner, *tt.want)
+			}
+
+		})
+	}
+}
 func TestClient_ReadWriteTransaction_Query_ReadOptions(t *testing.T) {
 	for _, tt := range readOptionsTestCases() {
 		t.Run(tt.name, func(t *testing.T) {

diff --git a/spanner/transaction.go b/spanner/transaction.go
@@ -185,6 +185,10 @@ type ReadOptions struct {
 
 	// An option to control the order in which rows are returned from a read.
 	OrderBy sppb.ReadRequest_OrderBy
+
+	// A lock hint mechanism to use for this request. This setting is only applicable for
+	// read-write transaction as as read-only transactions do not take locks.
+	LockHint sppb.ReadRequest_LockHint
 }
 
 // merge combines two ReadOptions that the input parameter will have higher
@@ -198,6 +202,7 @@ func (ro ReadOptions) merge(opts ReadOptions) ReadOptions {
 		DataBoostEnabled:    ro.DataBoostEnabled,
 		DirectedReadOptions: ro.DirectedReadOptions,
 		OrderBy:             ro.OrderBy,
+		LockHint:            ro.LockHint,
 	}
 	if opts.Index != "" {
 		merged.Index = opts.Index
@@ -220,6 +225,9 @@ func (ro ReadOptions) merge(opts ReadOptions) ReadOptions {
 	if opts.OrderBy != sppb.ReadRequest_ORDER_BY_UNSPECIFIED {
 		merged.OrderBy = opts.OrderBy
 	}
+	if opts.LockHint != sppb.ReadRequest_LOCK_HINT_UNSPECIFIED {
+		merged.LockHint = opts.LockHint
+	}
 	return merged
 }
 
@@ -253,6 +261,7 @@ func (t *txReadOnly) ReadWithOptions(ctx context.Context, table string, keys Key
 	dataBoostEnabled := t.ro.DataBoostEnabled
 	directedReadOptions := t.ro.DirectedReadOptions
 	orderBy := t.ro.OrderBy
+	lockHint := t.ro.LockHint
 	if opts != nil {
 		index = opts.Index
 		if opts.Limit > 0 {
@@ -269,6 +278,10 @@ func (t *txReadOnly) ReadWithOptions(ctx context.Context, table string, keys Key
 		if opts.OrderBy != sppb.ReadRequest_ORDER_BY_UNSPECIFIED {
 			orderBy = opts.OrderBy
 		}
+		if opts.LockHint != sppb.ReadRequest_LOCK_HINT_UNSPECIFIED {
+			lockHint = opts.LockHint
+		}
+
 	}
 	var setTransactionID func(transactionID)
 	if _, ok := ts.Selector.(*sppb.TransactionSelector_Begin); ok {
@@ -297,6 +310,7 @@ func (t *txReadOnly) ReadWithOptions(ctx context.Context, table string, keys Key
 					DataBoostEnabled:    dataBoostEnabled,
 					DirectedReadOptions: directedReadOptions,
 					OrderBy:             orderBy,
+					LockHint:            lockHint,
 				})
 			if err != nil {
 				if _, ok := t.getTransactionSelector().GetSelector().(*sppb.TransactionSelector_Begin); ok {