Skip to content
This repository was archived by the owner on Apr 7, 2026. It is now read-only.

Commit c4fd3ae

Browse files
aseeringaseering
committed
Address code-review comments
1 parent 4556e81 commit c4fd3ae

7 files changed

Lines changed: 120 additions & 96 deletions

File tree

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -684,20 +684,20 @@ QueryOptions buildQueryOptions(QueryOptions requestOptions) {
684684
}
685685

686686
RequestOptions buildRequestOptions(Options options) {
687-
RequestOptions requestOptions = options.toRequestOptionsProto(false);
687+
RequestOptions.Builder builder = options.toRequestOptionsProto(false).toBuilder();
688688
RequestOptions.ClientContext defaultClientContext =
689689
session.getSpanner().getOptions().getClientContext();
690690
if (defaultClientContext != null) {
691-
RequestOptions.ClientContext.Builder builder = defaultClientContext.toBuilder();
692-
if (requestOptions.hasClientContext()) {
693-
builder.mergeFrom(requestOptions.getClientContext());
691+
RequestOptions.ClientContext.Builder clientContextBuilder = defaultClientContext.toBuilder();
692+
if (builder.hasClientContext()) {
693+
clientContextBuilder.mergeFrom(builder.getClientContext());
694694
}
695-
requestOptions = requestOptions.toBuilder().setClientContext(builder.build()).build();
695+
builder.setClientContext(clientContextBuilder.build());
696696
}
697697
if (getTransactionTag() != null) {
698-
return requestOptions.toBuilder().setTransactionTag(getTransactionTag()).build();
698+
builder.setTransactionTag(getTransactionTag());
699699
}
700-
return requestOptions;
700+
return builder.build();
701701
}
702702

703703
ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -489,18 +489,20 @@ ApiFuture<Transaction> beginTransactionAsync(
489489
if (sessionReference.getIsMultiplexed() && mutation != null) {
490490
requestBuilder.setMutationKey(mutation);
491491
}
492-
RequestOptions requestOptions = transactionOptions.toRequestOptionsProto(true);
492+
RequestOptions.Builder optionsBuilder =
493+
transactionOptions.toRequestOptionsProto(true).toBuilder();
493494
RequestOptions.ClientContext defaultClientContext = spanner.getOptions().getClientContext();
494495
if (defaultClientContext != null) {
495496
RequestOptions.ClientContext.Builder builder = defaultClientContext.toBuilder();
496-
if (requestOptions.hasClientContext()) {
497-
builder.mergeFrom(requestOptions.getClientContext());
497+
if (optionsBuilder.hasClientContext()) {
498+
builder.mergeFrom(optionsBuilder.getClientContext());
498499
}
499-
requestOptions = requestOptions.toBuilder().setClientContext(builder.build()).build();
500+
optionsBuilder.setClientContext(builder.build());
500501
}
501502
if (!sessionReference.getIsMultiplexed()) {
502-
requestOptions = requestOptions.toBuilder().clearTransactionTag().build();
503+
optionsBuilder.clearTransactionTag();
503504
}
505+
RequestOptions requestOptions = optionsBuilder.build();
504506
if (!requestOptions.equals(RequestOptions.getDefaultInstance())) {
505507
requestBuilder.setRequestOptions(requestOptions);
506508
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1988,7 +1988,7 @@ public Builder setDefaultTransactionOptions(
19881988
}
19891989

19901990
/** Sets the default {@link RequestOptions.ClientContext} for all requests. */
1991-
public Builder setClientContext(RequestOptions.ClientContext clientContext) {
1991+
public Builder setDefaultClientContext(RequestOptions.ClientContext clientContext) {
19921992
this.clientContext = clientContext;
19931993
return this;
19941994
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,14 @@
4646
/** Unit tests for {@link Options}. */
4747
@RunWith(JUnit4.class)
4848
public class OptionsTest {
49+
private static final DirectedReadOptions DIRECTED_READ_OPTIONS =
50+
DirectedReadOptions.newBuilder()
51+
.setIncludeReplicas(
52+
IncludeReplicas.newBuilder()
53+
.addReplicaSelections(
54+
ReplicaSelection.newBuilder().setLocation("us-west1").build()))
55+
.build();
56+
4957
@Test
5058
public void testToRequestOptionsProto() {
5159
RequestOptions.ClientContext clientContext =
@@ -72,14 +80,6 @@ public void testToRequestOptionsProto() {
7280
assertEquals(clientContext, protoForTransaction.getClientContext());
7381
}
7482

75-
private static final DirectedReadOptions DIRECTED_READ_OPTIONS =
76-
DirectedReadOptions.newBuilder()
77-
.setIncludeReplicas(
78-
IncludeReplicas.newBuilder()
79-
.addReplicaSelections(
80-
ReplicaSelection.newBuilder().setLocation("us-west1").build()))
81-
.build();
82-
8383
@Test
8484
public void negativeLimitsNotAllowed() {
8585
IllegalArgumentException e =

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java‎

Lines changed: 36 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -78,42 +78,6 @@
7878
/** Unit tests for {@link com.google.cloud.spanner.SessionImpl}. */
7979
@RunWith(JUnit4.class)
8080
public class SessionImplTest {
81-
@Test
82-
public void testBeginTransactionWithClientContext() {
83-
RequestOptions.ClientContext clientContext =
84-
RequestOptions.ClientContext.newBuilder()
85-
.putSecureContext(
86-
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
87-
.build();
88-
Mockito.when(
89-
rpc.beginTransactionAsync(
90-
Mockito.any(BeginTransactionRequest.class), anyMap(), eq(true)))
91-
.thenReturn(
92-
ApiFutures.immediateFuture(
93-
Transaction.newBuilder().setId(ByteString.copyFromUtf8("tx")).build()));
94-
95-
((SessionImpl) session)
96-
.beginTransactionAsync(
97-
Options.fromTransactionOptions(
98-
Options.priority(Options.RpcPriority.HIGH),
99-
Options.tag("tag"),
100-
Options.clientContext(clientContext)),
101-
true,
102-
Collections.emptyMap(),
103-
null,
104-
null);
105-
106-
ArgumentCaptor<BeginTransactionRequest> requestCaptor =
107-
ArgumentCaptor.forClass(BeginTransactionRequest.class);
108-
Mockito.verify(rpc).beginTransactionAsync(requestCaptor.capture(), anyMap(), eq(true));
109-
BeginTransactionRequest request = requestCaptor.getValue();
110-
RequestOptions requestOptions = request.getRequestOptions();
111-
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, requestOptions.getPriority());
112-
// TransactionTag should NOT be set because session is not multiplexed.
113-
assertEquals("", requestOptions.getTransactionTag());
114-
assertEquals(clientContext, requestOptions.getClientContext());
115-
}
116-
11781
@Mock private SpannerRpc rpc;
11882
@Mock private SpannerOptions spannerOptions;
11983
private com.google.cloud.spanner.Session session;
@@ -204,6 +168,42 @@ private void doNestedRwTransaction() {
204168
});
205169
}
206170

171+
@Test
172+
public void testBeginTransactionWithClientContext() {
173+
RequestOptions.ClientContext clientContext =
174+
RequestOptions.ClientContext.newBuilder()
175+
.putSecureContext(
176+
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
177+
.build();
178+
Mockito.when(
179+
rpc.beginTransactionAsync(
180+
Mockito.any(BeginTransactionRequest.class), anyMap(), eq(true)))
181+
.thenReturn(
182+
ApiFutures.immediateFuture(
183+
Transaction.newBuilder().setId(ByteString.copyFromUtf8("tx")).build()));
184+
185+
((SessionImpl) session)
186+
.beginTransactionAsync(
187+
Options.fromTransactionOptions(
188+
Options.priority(Options.RpcPriority.HIGH),
189+
Options.tag("tag"),
190+
Options.clientContext(clientContext)),
191+
true,
192+
Collections.emptyMap(),
193+
null,
194+
null);
195+
196+
ArgumentCaptor<BeginTransactionRequest> requestCaptor =
197+
ArgumentCaptor.forClass(BeginTransactionRequest.class);
198+
Mockito.verify(rpc).beginTransactionAsync(requestCaptor.capture(), anyMap(), eq(true));
199+
BeginTransactionRequest request = requestCaptor.getValue();
200+
RequestOptions requestOptions = request.getRequestOptions();
201+
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, requestOptions.getPriority());
202+
// TransactionTag should NOT be set because session is not multiplexed.
203+
assertEquals("", requestOptions.getTransactionTag());
204+
assertEquals(clientContext, requestOptions.getClientContext());
205+
}
206+
207207
@Test
208208
public void nestedReadWriteTxnThrows() {
209209
SpannerException e = assertThrows(SpannerException.class, () -> doNestedRwTransaction());

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java‎

Lines changed: 31 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import static org.mockito.ArgumentMatchers.any;
2424
import static org.mockito.ArgumentMatchers.eq;
2525
import static org.mockito.Mockito.doThrow;
26-
import static org.mockito.Mockito.eq;
2726
import static org.mockito.Mockito.mock;
2827
import static org.mockito.Mockito.never;
2928
import static org.mockito.Mockito.times;
@@ -101,37 +100,6 @@ public void release(ScheduledExecutorService exec) {
101100
}
102101
}
103102

104-
@Test
105-
public void testCommitWithClientContext() {
106-
RequestOptions.ClientContext clientContext =
107-
RequestOptions.ClientContext.newBuilder()
108-
.putSecureContext(
109-
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
110-
.build();
111-
when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
112-
when(session.newTransaction(any(Options.class), any())).thenReturn(txn);
113-
Mockito.clearInvocations(session);
114-
transactionRunner =
115-
new TransactionRunnerImpl(
116-
session,
117-
Options.priority(Options.RpcPriority.HIGH),
118-
Options.tag("tag"),
119-
Options.clientContext(clientContext));
120-
transactionRunner.setSpan(span);
121-
122-
transactionRunner.run(
123-
transaction -> {
124-
return null;
125-
});
126-
127-
ArgumentCaptor<Options> optionsCaptor = ArgumentCaptor.forClass(Options.class);
128-
verify(session).newTransaction(optionsCaptor.capture(), any());
129-
Options capturedOptions = optionsCaptor.getValue();
130-
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, capturedOptions.priority());
131-
assertEquals("tag", capturedOptions.tag());
132-
assertEquals(clientContext, capturedOptions.clientContext());
133-
}
134-
135103
@Mock private SpannerRpc rpc;
136104
@Mock private SessionImpl session;
137105
@Mock private TransactionRunnerImpl.TransactionContextImpl txn;
@@ -196,6 +164,37 @@ public void setUp() {
196164
transactionRunner.setSpan(span);
197165
}
198166

167+
@Test
168+
public void testCommitWithClientContext() {
169+
RequestOptions.ClientContext clientContext =
170+
RequestOptions.ClientContext.newBuilder()
171+
.putSecureContext(
172+
"key", com.google.protobuf.Value.newBuilder().setStringValue("value").build())
173+
.build();
174+
when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
175+
when(session.newTransaction(any(Options.class), any())).thenReturn(txn);
176+
Mockito.clearInvocations(session);
177+
transactionRunner =
178+
new TransactionRunnerImpl(
179+
session,
180+
Options.priority(Options.RpcPriority.HIGH),
181+
Options.tag("tag"),
182+
Options.clientContext(clientContext));
183+
transactionRunner.setSpan(span);
184+
185+
transactionRunner.run(
186+
transaction -> {
187+
return null;
188+
});
189+
190+
ArgumentCaptor<Options> optionsCaptor = ArgumentCaptor.forClass(Options.class);
191+
verify(session).newTransaction(optionsCaptor.capture(), any());
192+
Options capturedOptions = optionsCaptor.getValue();
193+
assertEquals(RequestOptions.Priority.PRIORITY_HIGH, capturedOptions.priority());
194+
assertEquals("tag", capturedOptions.tag());
195+
assertEquals(clientContext, capturedOptions.clientContext());
196+
}
197+
199198
@SuppressWarnings("unchecked")
200199
@Test
201200
public void usesPreparedTransaction() {

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ClientContextMockServerTest.java‎

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.google.cloud.spanner.DatabaseId;
2424
import com.google.cloud.spanner.Dialect;
2525
import com.google.cloud.spanner.MockSpannerServiceImpl;
26+
import com.google.cloud.spanner.Mutation;
2627
import com.google.cloud.spanner.ResultSet;
2728
import com.google.cloud.spanner.Spanner;
2829
import com.google.cloud.spanner.SpannerOptions;
@@ -144,8 +145,7 @@ public void testCommit_PropagatesClientContext() {
144145
}
145146

146147
@Test
147-
public void testBeginTransaction_PropagatesClientContext() {
148-
// 1. Test lazy transaction start (default).
148+
public void testBeginTransaction_PropagatesClientContextWithLazyStart() {
149149
// The BeginTransaction option is inlined with the first statement.
150150
try (Connection connection = createConnection()) {
151151
connection.setClientContext(CLIENT_CONTEXT);
@@ -157,8 +157,10 @@ public void testBeginTransaction_PropagatesClientContext() {
157157
assertEquals(CLIENT_CONTEXT, request.getRequestOptions().getClientContext());
158158
assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
159159
}
160+
}
160161

161-
// 2. Test eager transaction start.
162+
@Test
163+
public void testBeginTransaction_PropagatesClientContextWithEagerStartAborted() {
162164
// We can force an explicit BeginTransaction RPC by failing the first statement with an ABORTED
163165
// error. If the statement fails before returning a transaction ID, the retry will use an
164166
// explicit BeginTransaction RPC.
@@ -178,12 +180,11 @@ public void testBeginTransaction_PropagatesClientContext() {
178180
connection.beginTransaction();
179181
connection.executeUpdate(INSERT_STATEMENT);
180182

181-
// We expect multiple ExecuteSqlRequests.
183+
// We expect two ExecuteSqlRequests.
182184
// 1. The first one fails with ABORTED. This request includes the BeginTransaction option.
183185
// 2. The retry.
184-
// Note: precise count depends on Gax retry logic vs Spanner retry logic interaction.
185186
int executeSqlCount = mockSpanner.countRequestsOfType(ExecuteSqlRequest.class);
186-
assertFalse(executeSqlCount < 2);
187+
assertEquals(2, executeSqlCount);
187188

188189
for (ExecuteSqlRequest req : mockSpanner.getRequestsOfType(ExecuteSqlRequest.class)) {
189190
assertEquals(CLIENT_CONTEXT, req.getRequestOptions().getClientContext());
@@ -197,6 +198,28 @@ public void testBeginTransaction_PropagatesClientContext() {
197198
}
198199
}
199200

201+
@Test
202+
public void testBeginTransaction_PropagatesClientContextWithEagerStartMutations() {
203+
// We can also force an explicit BeginTransaction RPC by constructing a transaction
204+
// that only issues mutations. Mutation RPCs cannot start a transaction, so
205+
// if they are the only RPCs in the transaction, then an explicit BeginTransaction
206+
// must be issued.
207+
try (Connection connection = createConnection()) {
208+
connection.setClientContext(CLIENT_CONTEXT);
209+
connection.beginTransaction();
210+
connection.bufferedWrite(Mutation.newInsertBuilder("my-table").set("my-col").to(1L).build());
211+
connection.commit();
212+
213+
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
214+
BeginTransactionRequest request =
215+
mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0);
216+
assertEquals(CLIENT_CONTEXT, request.getRequestOptions().getClientContext());
217+
assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
218+
CommitRequest commitRequest = mockSpanner.getRequestsOfType(CommitRequest.class).get(0);
219+
assertEquals(CLIENT_CONTEXT, commitRequest.getRequestOptions().getClientContext());
220+
}
221+
}
222+
200223
@Test
201224
public void testDatabaseClient_ClientContextMerging() {
202225
String projectId = "test-project";
@@ -215,7 +238,7 @@ public void testDatabaseClient_ClientContextMerging() {
215238
.setProjectId(projectId)
216239
.setHost("http://localhost:" + getPort())
217240
.usePlainText()
218-
.setClientContext(defaultContext)
241+
.setDefaultClientContext(defaultContext)
219242
.build();
220243

221244
try (Spanner spanner = options.getService()) {

0 commit comments

Comments
 (0)