future) throws SpannerException {
}
throw SpannerExceptionFactory.asSpannerException(e.getCause());
} catch (InterruptedException e) {
- throw SpannerExceptionFactory.propagateInterrupt(e, null /*TODO: requestId*/);
+ throw SpannerExceptionFactory.propagateInterrupt(e);
} catch (CancellationException e) {
- throw SpannerExceptionFactory.newSpannerExceptionForCancellation(
- null, e, null /*TODO: requestId*/);
+ throw SpannerExceptionFactory.newSpannerExceptionForCancellation(null, e);
}
}
}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java
index 0d841d24463..837a008a833 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java
@@ -25,9 +25,8 @@ public class SpannerBatchUpdateException extends SpannerException {
ErrorCode code,
String message,
long[] counts,
- Throwable cause,
- XGoogSpannerRequestId reqId) {
- super(token, code, false, message, cause, null, reqId);
+ Throwable cause) {
+ super(token, code, false, message, cause, null);
updateCounts = counts;
}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java
index fbe60e2a1d6..0829cc35d62 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java
@@ -16,11 +16,11 @@
package com.google.cloud.spanner;
-import com.google.api.core.InternalApi;
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.ErrorDetails;
import com.google.cloud.grpc.BaseGrpcServiceException;
import com.google.common.base.Preconditions;
+import com.google.common.base.Strings;
import com.google.protobuf.util.Durations;
import com.google.rpc.ResourceInfo;
import com.google.rpc.RetryInfo;
@@ -41,9 +41,8 @@ public abstract static class ResourceNotFoundException extends SpannerException
@Nullable String message,
ResourceInfo resourceInfo,
@Nullable Throwable cause,
- @Nullable ApiException apiException,
- @Nullable XGoogSpannerRequestId reqId) {
- super(token, ErrorCode.NOT_FOUND, /* retryable */ false, message, cause, apiException, reqId);
+ @Nullable ApiException apiException) {
+ super(token, ErrorCode.NOT_FOUND, /* retryable */ false, message, cause, apiException);
this.resourceInfo = resourceInfo;
}
@@ -59,7 +58,7 @@ public String getResourceName() {
private final ErrorCode code;
private final ApiException apiException;
- private XGoogSpannerRequestId requestId;
+ private final XGoogSpannerRequestId requestId;
private String statement;
/** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
@@ -80,25 +79,13 @@ public String getResourceName() {
@Nullable String message,
@Nullable Throwable cause,
@Nullable ApiException apiException) {
- this(token, code, retryable, message, cause, apiException, null);
- }
-
- /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
- SpannerException(
- DoNotConstructDirectly token,
- ErrorCode code,
- boolean retryable,
- @Nullable String message,
- @Nullable Throwable cause,
- @Nullable ApiException apiException,
- @Nullable XGoogSpannerRequestId requestId) {
super(message, cause, code.getCode(), retryable);
if (token != DoNotConstructDirectly.ALLOWED) {
throw new AssertionError("Do not construct directly: use SpannerExceptionFactory");
}
this.code = Preconditions.checkNotNull(code);
this.apiException = apiException;
- this.requestId = requestId;
+ this.requestId = extractRequestId(cause);
}
@Override
@@ -109,6 +96,14 @@ public String getMessage() {
return String.format("%s - Statement: '%s'", super.getMessage(), this.statement);
}
+ @Override
+ public String toString() {
+ if (this.requestId == null) {
+ return super.toString();
+ }
+ return super.toString() + " - RequestId: " + this.requestId;
+ }
+
/** Returns the error code associated with this exception. */
public ErrorCode getErrorCode() {
return code;
@@ -150,7 +145,7 @@ static long extractRetryDelay(Throwable cause) {
Metadata trailers = Status.trailersFromThrowable(cause);
if (trailers != null && trailers.containsKey(KEY_RETRY_INFO)) {
RetryInfo retryInfo = trailers.get(KEY_RETRY_INFO);
- if (retryInfo.hasRetryDelay()) {
+ if (retryInfo != null && retryInfo.hasRetryDelay()) {
return Durations.toMillis(retryInfo.getRetryDelay());
}
}
@@ -158,6 +153,20 @@ static long extractRetryDelay(Throwable cause) {
return -1L;
}
+ @Nullable
+ static XGoogSpannerRequestId extractRequestId(Throwable cause) {
+ if (cause != null) {
+ Metadata trailers = Status.trailersFromThrowable(cause);
+ if (trailers != null && trailers.containsKey(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY)) {
+ String requestId = trailers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY);
+ if (!Strings.isNullOrEmpty(requestId)) {
+ return XGoogSpannerRequestId.of(requestId);
+ }
+ }
+ }
+ return null;
+ }
+
/**
* Checks the underlying reason of the exception and if it's {@link ApiException} then return the
* reason otherwise null.
@@ -224,10 +233,4 @@ public ErrorDetails getErrorDetails() {
void setStatement(String statement) {
this.statement = statement;
}
-
- /** Sets the requestId. */
- @InternalApi
- public void setRequestId(XGoogSpannerRequestId reqId) {
- this.requestId = reqId;
- }
}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java
index b9cc1cfb8b1..185f98b5433 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java
@@ -58,37 +58,17 @@ public final class SpannerExceptionFactory {
ProtoUtils.keyForProto(ErrorInfo.getDefaultInstance());
public static SpannerException newSpannerException(ErrorCode code, @Nullable String message) {
- return newSpannerException(code, message, (XGoogSpannerRequestId) (null));
- }
-
- public static SpannerException newSpannerException(
- ErrorCode code,
- @Nullable String message,
- @Nullable Throwable cause,
- @Nullable XGoogSpannerRequestId reqId) {
- return newSpannerExceptionPreformatted(
- code, formatMessage(code, message), cause, (ApiException) (null), reqId);
- }
-
- public static SpannerException newSpannerException(
- ErrorCode code, @Nullable String message, @Nullable XGoogSpannerRequestId reqId) {
- return newSpannerException(code, message, (Throwable) (null), reqId);
+ return newSpannerException(code, message, null);
}
public static SpannerException newSpannerException(
ErrorCode code, @Nullable String message, @Nullable Throwable cause) {
- return newSpannerException(code, message, cause, null);
+ return newSpannerExceptionPreformatted(code, formatMessage(code, message), cause, null);
}
public static SpannerException propagateInterrupt(InterruptedException e) {
- return propagateInterrupt(e, null);
- }
-
- public static SpannerException propagateInterrupt(
- InterruptedException e, XGoogSpannerRequestId reqId) {
Thread.currentThread().interrupt();
- return SpannerExceptionFactory.newSpannerException(
- ErrorCode.CANCELLED, "Interrupted", e, reqId);
+ return SpannerExceptionFactory.newSpannerException(ErrorCode.CANCELLED, "Interrupted", e);
}
/**
@@ -132,27 +112,17 @@ public static SpannerException asSpannerException(Throwable t) {
* #newSpannerException(ErrorCode, String)} instead of this method.
*/
public static SpannerException newSpannerException(Throwable cause) {
- return newSpannerException(null, cause, null);
- }
-
- public static SpannerException newSpannerException(Throwable cause, XGoogSpannerRequestId reqId) {
- return newSpannerException(null, cause, reqId);
+ return newSpannerException(null, cause);
}
public static SpannerBatchUpdateException newSpannerBatchUpdateException(
ErrorCode code, String message, long[] updateCounts) {
- return newSpannerBatchUpdateException(code, message, updateCounts, null);
- }
-
- public static SpannerBatchUpdateException newSpannerBatchUpdateException(
- ErrorCode code, String message, long[] updateCounts, @Nullable XGoogSpannerRequestId reqId) {
DoNotConstructDirectly token = DoNotConstructDirectly.ALLOWED;
SpannerException cause = null;
if (isTransactionMutationLimitException(code, message)) {
- cause =
- new TransactionMutationLimitExceededException(token, code, message, null, null, reqId);
+ cause = new TransactionMutationLimitExceededException(token, code, message, null, null);
}
- return new SpannerBatchUpdateException(token, code, message, updateCounts, cause, reqId);
+ return new SpannerBatchUpdateException(token, code, message, updateCounts, cause);
}
/** Constructs a specific error that */
@@ -205,10 +175,6 @@ public static SpannerBatchUpdateException newSpannerBatchUpdateException(
cause);
}
- public static SpannerException newSpannerException(@Nullable Context context, Throwable cause) {
- return newSpannerException(context, cause, null);
- }
-
/**
* Creates a new exception based on {@code cause}. If {@code cause} indicates cancellation, {@code
* context} will be inspected to establish the type of cancellation.
@@ -216,22 +182,21 @@ public static SpannerException newSpannerException(@Nullable Context context, Th
* Intended for internal library use; user code should use {@link
* #newSpannerException(ErrorCode, String)} instead of this method.
*/
- public static SpannerException newSpannerException(
- @Nullable Context context, Throwable cause, @Nullable XGoogSpannerRequestId reqId) {
+ public static SpannerException newSpannerException(@Nullable Context context, Throwable cause) {
if (cause instanceof SpannerException) {
SpannerException e = (SpannerException) cause;
- return newSpannerExceptionPreformatted(e.getErrorCode(), e.getMessage(), e, null, reqId);
+ return newSpannerExceptionPreformatted(e.getErrorCode(), e.getMessage(), e, null);
} else if (cause instanceof CancellationException) {
- return newSpannerExceptionForCancellation(context, cause, reqId);
+ return newSpannerExceptionForCancellation(context, cause);
} else if (cause instanceof ApiException) {
- return fromApiException((ApiException) cause, reqId);
+ return fromApiException((ApiException) cause);
}
// Extract gRPC status. This will produce "UNKNOWN" for non-gRPC exceptions.
Status status = Status.fromThrowable(cause);
if (status.getCode() == Status.Code.CANCELLED) {
- return newSpannerExceptionForCancellation(context, cause, reqId);
+ return newSpannerExceptionForCancellation(context, cause);
}
- return newSpannerException(ErrorCode.fromGrpcStatus(status), cause.getMessage(), cause, reqId);
+ return newSpannerException(ErrorCode.fromGrpcStatus(status), cause.getMessage(), cause);
}
public static RuntimeException causeAsRunTimeException(ExecutionException executionException) {
@@ -256,11 +221,6 @@ static SpannerException newRetryOnDifferentGrpcChannelException(
static SpannerException newSpannerExceptionForCancellation(
@Nullable Context context, @Nullable Throwable cause) {
- return newSpannerExceptionForCancellation(context, cause, null);
- }
-
- static SpannerException newSpannerExceptionForCancellation(
- @Nullable Context context, @Nullable Throwable cause, @Nullable XGoogSpannerRequestId reqId) {
if (context != null && context.isCancelled()) {
Throwable cancellationCause = context.cancellationCause();
Throwable throwable =
@@ -269,14 +229,13 @@ static SpannerException newSpannerExceptionForCancellation(
: MoreObjects.firstNonNull(cause, cancellationCause);
if (cancellationCause instanceof TimeoutException) {
return newSpannerException(
- ErrorCode.DEADLINE_EXCEEDED, "Current context exceeded deadline", throwable, reqId);
+ ErrorCode.DEADLINE_EXCEEDED, "Current context exceeded deadline", throwable);
} else {
- return newSpannerException(
- ErrorCode.CANCELLED, "Current context was cancelled", throwable, reqId);
+ return newSpannerException(ErrorCode.CANCELLED, "Current context was cancelled", throwable);
}
}
return newSpannerException(
- ErrorCode.CANCELLED, cause == null ? "Cancelled" : cause.getMessage(), cause, reqId);
+ ErrorCode.CANCELLED, cause == null ? "Cancelled" : cause.getMessage(), cause);
}
private static String formatMessage(ErrorCode code, @Nullable String message) {
@@ -355,13 +314,12 @@ static SpannerException newSpannerExceptionPreformatted(
ErrorCode code,
@Nullable String message,
@Nullable Throwable cause,
- @Nullable ApiException apiException,
- @Nullable XGoogSpannerRequestId reqId) {
+ @Nullable ApiException apiException) {
// This is the one place in the codebase that is allowed to call constructors directly.
DoNotConstructDirectly token = DoNotConstructDirectly.ALLOWED;
switch (code) {
case ABORTED:
- return new AbortedException(token, message, cause, apiException, reqId);
+ return new AbortedException(token, message, cause, apiException);
case RESOURCE_EXHAUSTED:
ErrorInfo info = extractErrorInfo(cause, apiException);
if (info != null
@@ -370,8 +328,7 @@ static SpannerException newSpannerExceptionPreformatted(
&& AdminRequestsPerMinuteExceededException.ADMIN_REQUESTS_LIMIT_VALUE.equals(
info.getMetadataMap()
.get(AdminRequestsPerMinuteExceededException.ADMIN_REQUESTS_LIMIT_KEY))) {
- return new AdminRequestsPerMinuteExceededException(
- token, message, cause, apiException, reqId);
+ return new AdminRequestsPerMinuteExceededException(token, message, cause, apiException);
}
case NOT_FOUND:
ResourceInfo resourceInfo = extractResourceInfo(cause);
@@ -379,39 +336,36 @@ static SpannerException newSpannerExceptionPreformatted(
switch (resourceInfo.getResourceType()) {
case SESSION_RESOURCE_TYPE:
return new SessionNotFoundException(
- token, message, resourceInfo, cause, apiException, reqId);
+ token, message, resourceInfo, cause, apiException);
case DATABASE_RESOURCE_TYPE:
return new DatabaseNotFoundException(
- token, message, resourceInfo, cause, apiException, reqId);
+ token, message, resourceInfo, cause, apiException);
case INSTANCE_RESOURCE_TYPE:
return new InstanceNotFoundException(
- token, message, resourceInfo, cause, apiException, reqId);
+ token, message, resourceInfo, cause, apiException);
}
}
case INVALID_ARGUMENT:
if (isTransactionMutationLimitException(cause, apiException)) {
return new TransactionMutationLimitExceededException(
- token, code, message, cause, apiException, reqId);
+ token, code, message, cause, apiException);
}
if (isMissingDefaultSequenceKindException(apiException)) {
- return new MissingDefaultSequenceKindException(
- token, code, message, cause, apiException, reqId);
+ return new MissingDefaultSequenceKindException(token, code, message, cause, apiException);
}
// Fall through to the default.
default:
return new SpannerException(
- token, code, isRetryable(code, cause), message, cause, apiException, reqId);
+ token, code, isRetryable(code, cause), message, cause, apiException);
}
}
static SpannerException newSpannerExceptionPreformatted(
ErrorCode code, @Nullable String message, @Nullable Throwable cause) {
- return newSpannerExceptionPreformatted(
- code, message, cause, null, (XGoogSpannerRequestId) (null));
+ return newSpannerExceptionPreformatted(code, message, cause, null);
}
- private static SpannerException fromApiException(
- ApiException exception, @Nullable XGoogSpannerRequestId reqId) {
+ private static SpannerException fromApiException(ApiException exception) {
Status.Code code;
if (exception.getStatusCode() instanceof GrpcStatusCode) {
code = ((GrpcStatusCode) exception.getStatusCode()).getTransportCode();
@@ -426,8 +380,7 @@ private static SpannerException fromApiException(
errorCode,
formatMessage(errorCode, exception.getMessage()),
exception.getCause(),
- exception,
- reqId);
+ exception);
}
private static boolean isRetryable(ErrorCode code, @Nullable Throwable cause) {
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
index 16f36144723..b4eef3a0f51 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
@@ -459,6 +459,14 @@ public boolean isClosed() {
}
}
+ void resetRequestIdCounters() {
+ gapicRpc.getRequestIdCreator().reset();
+ }
+
+ long getRequestIdClientId() {
+ return gapicRpc.getRequestIdCreator().getClientId();
+ }
+
/** Helper class for gRPC calls that can return paginated results. */
abstract static class PageFetcher implements NextPageFetcher {
private String nextPageToken;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerRetryHelper.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerRetryHelper.java
index 6ca1a4e02e8..0dabcbd0094 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerRetryHelper.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerRetryHelper.java
@@ -116,8 +116,7 @@ public TimedAttemptSettings createNextAttempt(
public boolean shouldRetry(Throwable prevThrowable, T prevResponse)
throws CancellationException {
if (Context.current().isCancelled()) {
- throw SpannerExceptionFactory.newSpannerExceptionForCancellation(
- Context.current(), null, null);
+ throw SpannerExceptionFactory.newSpannerExceptionForCancellation(Context.current(), null);
}
return prevThrowable instanceof AbortedException
|| prevThrowable instanceof com.google.api.gax.rpc.AbortedException;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java
index f6b0a4efd22..de215c5caee 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java
@@ -37,9 +37,8 @@ public class TransactionMutationLimitExceededException extends SpannerException
ErrorCode errorCode,
String message,
Throwable cause,
- @Nullable ApiException apiException,
- @Nullable XGoogSpannerRequestId reqId) {
- super(token, errorCode, /* retryable= */ false, message, cause, apiException, reqId);
+ @Nullable ApiException apiException) {
+ super(token, errorCode, /* retryable= */ false, message, cause, apiException);
}
static boolean isTransactionMutationLimitException(ErrorCode code, String message) {
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
index 6ce8cc11aa6..7afccce194c 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
@@ -449,8 +449,6 @@ private final class CommitRunnable implements Runnable {
@Override
public void run() {
- XGoogSpannerRequestId reqId =
- session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
try {
prev.get();
if (transactionId == null && transactionIdFuture == null) {
@@ -493,8 +491,7 @@ public void run() {
final ApiFuture commitFuture;
final ISpan opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span);
try (IScope ignore = tracer.withSpan(opSpan)) {
- commitFuture =
- rpc.commitAsync(commitRequest, reqId.withOptions(getTransactionChannelHint()));
+ commitFuture = rpc.commitAsync(commitRequest, getTransactionChannelHint());
}
session.markUsed(clock.instant());
commitFuture.addListener(
@@ -505,7 +502,7 @@ public void run() {
// future, but we add a result here as well as a safety precaution.
res.setException(
SpannerExceptionFactory.newSpannerException(
- ErrorCode.INTERNAL, "commitFuture is not done", reqId));
+ ErrorCode.INTERNAL, "commitFuture is not done"));
return;
}
com.google.spanner.v1.CommitResponse proto = commitFuture.get();
@@ -535,9 +532,7 @@ public void run() {
}
if (!proto.hasCommitTimestamp()) {
throw newSpannerException(
- ErrorCode.INTERNAL,
- "Missing commitTimestamp:\n" + session.getName(),
- reqId);
+ ErrorCode.INTERNAL, "Missing commitTimestamp:\n" + session.getName());
}
span.addAnnotation("Commit Done");
opSpan.end();
@@ -577,8 +572,7 @@ public void run() {
res.setException(SpannerExceptionFactory.propagateTimeout(e));
} catch (Throwable e) {
res.setException(
- SpannerExceptionFactory.newSpannerException(
- e.getCause() == null ? e : e.getCause(), reqId));
+ SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()));
}
}
}
@@ -697,8 +691,7 @@ options, getPreviousTransactionId())))
}
throw se;
} catch (InterruptedException e) {
- throw SpannerExceptionFactory.newSpannerExceptionForCancellation(
- null, e, null /*TODO: requestId*/);
+ throw SpannerExceptionFactory.newSpannerExceptionForCancellation(null, e);
}
}
// There is already a transactionId available. Include that id as the transaction to use.
@@ -938,12 +931,9 @@ private ResultSet internalExecuteUpdate(
final ExecuteSqlRequest.Builder builder =
getExecuteSqlRequestBuilder(
statement, queryMode, options, /* withTransactionSelector= */ true);
- XGoogSpannerRequestId reqId =
- session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
try {
com.google.spanner.v1.ResultSet resultSet =
- rpc.executeQuery(
- builder.build(), reqId.withOptions(getTransactionChannelHint()), isRouteToLeader());
+ rpc.executeQuery(builder.build(), getTransactionChannelHint(), isRouteToLeader());
session.markUsed(clock.instant());
if (resultSet.getMetadata().hasTransaction()) {
onTransactionMetadata(
@@ -1076,11 +1066,9 @@ public long[] batchUpdate(Iterable statements, UpdateOption... update
}
final ExecuteBatchDmlRequest.Builder builder =
getExecuteBatchDmlRequestBuilder(statements, options);
- XGoogSpannerRequestId reqId =
- session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
try {
com.google.spanner.v1.ExecuteBatchDmlResponse response =
- rpc.executeBatchDml(builder.build(), reqId.withOptions(getTransactionChannelHint()));
+ rpc.executeBatchDml(builder.build(), getTransactionChannelHint());
session.markUsed(clock.instant());
long[] results = new long[response.getResultSetsCount()];
for (int i = 0; i < response.getResultSetsCount(); ++i) {
@@ -1104,8 +1092,7 @@ public long[] batchUpdate(Iterable statements, UpdateOption... update
throw newSpannerBatchUpdateException(
ErrorCode.fromRpcStatus(response.getStatus()),
response.getStatus().getMessage(),
- results,
- reqId);
+ results);
}
return results;
} catch (Throwable e) {
@@ -1140,15 +1127,11 @@ public ApiFuture batchUpdateAsync(
final ExecuteBatchDmlRequest.Builder builder =
getExecuteBatchDmlRequestBuilder(statements, options);
ApiFuture response;
- XGoogSpannerRequestId reqId =
- session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
try {
// Register the update as an async operation that must finish before the transaction may
// commit.
increaseAsyncOperations();
- response =
- rpc.executeBatchDmlAsync(
- builder.build(), reqId.withOptions(getTransactionChannelHint()));
+ response = rpc.executeBatchDmlAsync(builder.build(), getTransactionChannelHint());
session.markUsed(clock.instant());
} catch (Throwable t) {
decreaseAsyncOperations();
@@ -1178,8 +1161,7 @@ public ApiFuture batchUpdateAsync(
throw newSpannerBatchUpdateException(
ErrorCode.fromRpcStatus(batchDmlResponse.getStatus()),
batchDmlResponse.getStatus().getMessage(),
- results,
- reqId);
+ results);
}
return results;
},
@@ -1233,7 +1215,7 @@ public ListenableAsyncResultSet executeQueryAsync(
private final SessionImpl session;
private final Options options;
private ISpan span;
- private TraceWrapper tracer;
+ private final TraceWrapper tracer;
private TransactionContextImpl txn;
private volatile boolean isValid = true;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java
index 47ccb05231c..d858fdb9273 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java
@@ -17,13 +17,11 @@
package com.google.cloud.spanner;
import com.google.api.core.InternalApi;
-import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.common.annotations.VisibleForTesting;
+import io.grpc.CallOptions;
import io.grpc.Metadata;
import java.math.BigInteger;
import java.security.SecureRandom;
-import java.util.HashMap;
-import java.util.Map;
import java.util.Objects;
import java.util.regex.MatchResult;
import java.util.regex.Matcher;
@@ -31,13 +29,15 @@
@InternalApi
public class XGoogSpannerRequestId {
- // 1. Generate the random process Id singleton.
+ // 1. Generate the random process ID singleton.
@VisibleForTesting
static final String RAND_PROCESS_ID = XGoogSpannerRequestId.generateRandProcessId();
- public static String REQUEST_ID = "x-goog-spanner-request-id";
- public static final Metadata.Key REQUEST_HEADER_KEY =
- Metadata.Key.of(REQUEST_ID, Metadata.ASCII_STRING_MARSHALLER);
+ public static String REQUEST_ID_HEADER_NAME = "x-goog-spanner-request-id";
+ public static final Metadata.Key REQUEST_ID_HEADER_KEY =
+ Metadata.Key.of(REQUEST_ID_HEADER_NAME, Metadata.ASCII_STRING_MARSHALLER);
+ public static final CallOptions.Key REQUEST_ID_CALL_OPTIONS_KEY =
+ CallOptions.Key.create("XGoogSpannerRequestId");
@VisibleForTesting
static final long VERSION = 1; // The version of the specification being implemented.
@@ -59,6 +59,20 @@ public static XGoogSpannerRequestId of(
return new XGoogSpannerRequestId(nthClientId, nthChannelId, nthRequest, attempt);
}
+ @VisibleForTesting
+ long getNthClientId() {
+ return nthClientId;
+ }
+
+ @VisibleForTesting
+ long getNthChannelId() {
+ return nthChannelId;
+ }
+
+ boolean hasChannelId() {
+ return nthChannelId > 0;
+ }
+
@VisibleForTesting
long getAttempt() {
return this.attempt;
@@ -95,8 +109,8 @@ private static String generateRandProcessId() {
return String.format("%016x", bigInt);
}
- @Override
- public String toString() {
+ /** Returns the string representation of this RequestId as it should be sent to Spanner. */
+ public String getHeaderValue() {
return String.format(
"%d.%s.%d.%d.%d.%d",
XGoogSpannerRequestId.VERSION,
@@ -107,6 +121,18 @@ public String toString() {
this.attempt);
}
+ @Override
+ public String toString() {
+ return String.format(
+ "%d.%s.%d.%s.%d.%d",
+ XGoogSpannerRequestId.VERSION,
+ XGoogSpannerRequestId.RAND_PROCESS_ID,
+ this.nthClientId,
+ this.nthChannelId < 0 ? "x" : String.valueOf(this.nthChannelId),
+ this.nthRequest,
+ this.attempt);
+ }
+
public String debugToString() {
return String.format(
"%d.%s.nth_client=%d.nth_chan=%d.nth_req=%d.attempt=%d",
@@ -151,31 +177,38 @@ public void incrementAttempt() {
this.attempt++;
}
- Map withOptions(Map options) {
- Map copyOptions = new HashMap<>();
- if (options != null) {
- copyOptions.putAll(options);
- }
- copyOptions.put(SpannerRpc.Option.REQUEST_ID, this);
- return copyOptions;
- }
-
@Override
public int hashCode() {
return Objects.hash(this.nthClientId, this.nthChannelId, this.nthRequest, this.attempt);
}
- interface RequestIdCreator {
- XGoogSpannerRequestId nextRequestId(long channelId, int attempt);
+ @InternalApi
+ public interface RequestIdCreator {
+ long getClientId();
+
+ XGoogSpannerRequestId nextRequestId(long channelId);
+
+ void reset();
}
- static class NoopRequestIdCreator implements RequestIdCreator {
- NoopRequestIdCreator() {}
+ // TODO: Move this class into test code.
+ static final class NoopRequestIdCreator implements RequestIdCreator {
+ static final NoopRequestIdCreator INSTANCE = new NoopRequestIdCreator();
+
+ private NoopRequestIdCreator() {}
@Override
- public XGoogSpannerRequestId nextRequestId(long channelId, int attempt) {
- return XGoogSpannerRequestId.of(1, 1, 1, 0);
+ public long getClientId() {
+ return 1L;
}
+
+ @Override
+ public XGoogSpannerRequestId nextRequestId(long channelId) {
+ return XGoogSpannerRequestId.of(1, channelId, 1, 0);
+ }
+
+ @Override
+ public void reset() {}
}
public void setChannelId(long channelId) {
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java
index c448695afd2..4a8d643b79c 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java
@@ -266,7 +266,7 @@ public ApiFuture runBatchAsync(CallType callType) {
} catch (SpannerException e) {
long[] updateCounts = extractUpdateCounts(operationReference.get());
throw SpannerExceptionFactory.newSpannerBatchUpdateException(
- e.getErrorCode(), e.getMessage(), updateCounts, null /* TODO: requestId */);
+ e.getErrorCode(), e.getMessage(), updateCounts);
}
} catch (Throwable t) {
span.setStatus(StatusCode.ERROR);
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
index 371d73a0af2..306bd9842ac 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
@@ -19,6 +19,7 @@
import static com.google.cloud.spanner.SpannerExceptionFactory.asSpannerException;
import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
import static com.google.cloud.spanner.ThreadFactoryUtil.tryCreateVirtualThreadPerTaskExecutor;
+import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_CALL_OPTIONS_KEY;
import com.google.api.core.ApiFunction;
import com.google.api.core.ApiFuture;
@@ -73,6 +74,7 @@
import com.google.cloud.spanner.SpannerOptions.CallContextConfigurator;
import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider;
import com.google.cloud.spanner.XGoogSpannerRequestId;
+import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
import com.google.cloud.spanner.admin.database.v1.stub.GrpcDatabaseAdminCallableFactory;
@@ -88,7 +90,6 @@
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.io.Resources;
import com.google.common.util.concurrent.RateLimiter;
@@ -241,6 +242,7 @@ public class GapicSpannerRpc implements SpannerRpc {
public static boolean DIRECTPATH_CHANNEL_CREATED = false;
private static final String API_FILE = "grpc-gcp-apiconfig.json";
+ private final RequestIdCreator requestIdCreator = new RequestIdCreatorImpl();
private boolean rpcIsClosed;
private final SpannerStub spannerStub;
private final RetrySettings executeQueryRetrySettings;
@@ -825,7 +827,7 @@ public OperationFuture call() {
isRetry = true;
if (operationName == null) {
- GrpcCallContext context = newCallContext(null, instanceName, initialRequest, method);
+ GrpcCallContext context = newAdminCallContext(instanceName, initialRequest, method);
return operationCallable.futureCall(initialRequest, context);
} else {
return operationCallable.resumeFutureCall(operationName);
@@ -916,8 +918,7 @@ public Paginated listInstanceConfigs(int pageSize, @Nullable Str
ListInstanceConfigsRequest request = requestBuilder.build();
GrpcCallContext context =
- newCallContext(
- null, projectName, request, InstanceAdminGrpc.getListInstanceConfigsMethod());
+ newAdminCallContext(projectName, request, InstanceAdminGrpc.getListInstanceConfigsMethod());
ListInstanceConfigsResponse response =
get(instanceAdminStub.listInstanceConfigsCallable().futureCall(request, context));
return new Paginated<>(response.getInstanceConfigsList(), response.getNextPageToken());
@@ -940,7 +941,7 @@ public OperationFuture createInsta
}
CreateInstanceConfigRequest request = builder.build();
GrpcCallContext context =
- newCallContext(null, parent, request, InstanceAdminGrpc.getCreateInstanceConfigMethod());
+ newAdminCallContext(parent, request, InstanceAdminGrpc.getCreateInstanceConfigMethod());
return instanceAdminStub.createInstanceConfigOperationCallable().futureCall(request, context);
}
@@ -957,11 +958,8 @@ public OperationFuture updateInsta
}
UpdateInstanceConfigRequest request = builder.build();
GrpcCallContext context =
- newCallContext(
- null,
- instanceConfig.getName(),
- request,
- InstanceAdminGrpc.getUpdateInstanceConfigMethod());
+ newAdminCallContext(
+ instanceConfig.getName(), request, InstanceAdminGrpc.getUpdateInstanceConfigMethod());
return instanceAdminStub.updateInstanceConfigOperationCallable().futureCall(request, context);
}
@@ -971,7 +969,7 @@ public InstanceConfig getInstanceConfig(String instanceConfigName) throws Spanne
GetInstanceConfigRequest.newBuilder().setName(instanceConfigName).build();
GrpcCallContext context =
- newCallContext(null, projectName, request, InstanceAdminGrpc.getGetInstanceConfigMethod());
+ newAdminCallContext(projectName, request, InstanceAdminGrpc.getGetInstanceConfigMethod());
return get(instanceAdminStub.getInstanceConfigCallable().futureCall(request, context));
}
@@ -990,8 +988,8 @@ public void deleteInstanceConfig(
}
DeleteInstanceConfigRequest request = requestBuilder.build();
GrpcCallContext context =
- newCallContext(
- null, instanceConfigName, request, InstanceAdminGrpc.getDeleteInstanceConfigMethod());
+ newAdminCallContext(
+ instanceConfigName, request, InstanceAdminGrpc.getDeleteInstanceConfigMethod());
get(instanceAdminStub.deleteInstanceConfigCallable().futureCall(request, context));
}
@@ -1012,8 +1010,8 @@ public Paginated listInstanceConfigOperations(
final ListInstanceConfigOperationsRequest request = requestBuilder.build();
final GrpcCallContext context =
- newCallContext(
- null, projectName, request, InstanceAdminGrpc.getListInstanceConfigOperationsMethod());
+ newAdminCallContext(
+ projectName, request, InstanceAdminGrpc.getListInstanceConfigOperationsMethod());
ListInstanceConfigOperationsResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() ->
@@ -1038,7 +1036,7 @@ public Paginated listInstances(
ListInstancesRequest request = requestBuilder.build();
GrpcCallContext context =
- newCallContext(null, projectName, request, InstanceAdminGrpc.getListInstancesMethod());
+ newAdminCallContext(projectName, request, InstanceAdminGrpc.getListInstancesMethod());
ListInstancesResponse response =
get(instanceAdminStub.listInstancesCallable().futureCall(request, context));
return new Paginated<>(response.getInstancesList(), response.getNextPageToken());
@@ -1054,7 +1052,7 @@ public OperationFuture createInstance(
.setInstance(instance)
.build();
GrpcCallContext context =
- newCallContext(null, parent, request, InstanceAdminGrpc.getCreateInstanceMethod());
+ newAdminCallContext(parent, request, InstanceAdminGrpc.getCreateInstanceMethod());
return instanceAdminStub.createInstanceOperationCallable().futureCall(request, context);
}
@@ -1064,8 +1062,8 @@ public OperationFuture updateInstance(
UpdateInstanceRequest request =
UpdateInstanceRequest.newBuilder().setInstance(instance).setFieldMask(fieldMask).build();
GrpcCallContext context =
- newCallContext(
- null, instance.getName(), request, InstanceAdminGrpc.getUpdateInstanceMethod());
+ newAdminCallContext(
+ instance.getName(), request, InstanceAdminGrpc.getUpdateInstanceMethod());
return instanceAdminStub.updateInstanceOperationCallable().futureCall(request, context);
}
@@ -1074,7 +1072,7 @@ public Instance getInstance(String instanceName) throws SpannerException {
GetInstanceRequest request = GetInstanceRequest.newBuilder().setName(instanceName).build();
GrpcCallContext context =
- newCallContext(null, instanceName, request, InstanceAdminGrpc.getGetInstanceMethod());
+ newAdminCallContext(instanceName, request, InstanceAdminGrpc.getGetInstanceMethod());
return get(instanceAdminStub.getInstanceCallable().futureCall(request, context));
}
@@ -1084,7 +1082,7 @@ public void deleteInstance(String instanceName) throws SpannerException {
DeleteInstanceRequest.newBuilder().setName(instanceName).build();
GrpcCallContext context =
- newCallContext(null, instanceName, request, InstanceAdminGrpc.getDeleteInstanceMethod());
+ newAdminCallContext(instanceName, request, InstanceAdminGrpc.getDeleteInstanceMethod());
get(instanceAdminStub.deleteInstanceCallable().futureCall(request, context));
}
@@ -1103,8 +1101,8 @@ public Paginated listBackupOperations(
final ListBackupOperationsRequest request = requestBuilder.build();
final GrpcCallContext context =
- newCallContext(
- null, instanceName, request, DatabaseAdminGrpc.getListBackupOperationsMethod());
+ newAdminCallContext(
+ instanceName, request, DatabaseAdminGrpc.getListBackupOperationsMethod());
ListBackupOperationsResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() ->
@@ -1128,8 +1126,8 @@ public Paginated listDatabaseOperations(
final ListDatabaseOperationsRequest request = requestBuilder.build();
final GrpcCallContext context =
- newCallContext(
- null, instanceName, request, DatabaseAdminGrpc.getListDatabaseOperationsMethod());
+ newAdminCallContext(
+ instanceName, request, DatabaseAdminGrpc.getListDatabaseOperationsMethod());
ListDatabaseOperationsResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() ->
@@ -1154,7 +1152,7 @@ public Paginated listDatabaseRoles(
final ListDatabaseRolesRequest request = requestBuilder.build();
final GrpcCallContext context =
- newCallContext(null, databaseName, request, DatabaseAdminGrpc.getListDatabaseRolesMethod());
+ newAdminCallContext(databaseName, request, DatabaseAdminGrpc.getListDatabaseRolesMethod());
ListDatabaseRolesResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.listDatabaseRolesCallable().futureCall(request, context)));
@@ -1178,7 +1176,7 @@ public Paginated listBackups(
final ListBackupsRequest request = requestBuilder.build();
final GrpcCallContext context =
- newCallContext(null, instanceName, request, DatabaseAdminGrpc.getListBackupsMethod());
+ newAdminCallContext(instanceName, request, DatabaseAdminGrpc.getListBackupsMethod());
ListBackupsResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.listBackupsCallable().futureCall(request, context)));
@@ -1198,7 +1196,7 @@ public Paginated listDatabases(
final ListDatabasesRequest request = requestBuilder.build();
final GrpcCallContext context =
- newCallContext(null, instanceName, request, DatabaseAdminGrpc.getListDatabasesMethod());
+ newAdminCallContext(instanceName, request, DatabaseAdminGrpc.getListDatabasesMethod());
ListDatabasesResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.listDatabasesCallable().futureCall(request, context)));
@@ -1300,8 +1298,7 @@ public OperationFuture updateDatabaseDdl(
}
final UpdateDatabaseDdlRequest request = requestBuilder.build();
final GrpcCallContext context =
- newCallContext(
- null,
+ newAdminCallContext(
databaseInfo.getId().getName(),
request,
DatabaseAdminGrpc.getUpdateDatabaseDdlMethod());
@@ -1341,7 +1338,7 @@ public void dropDatabase(String databaseName) throws SpannerException {
DropDatabaseRequest.newBuilder().setDatabase(databaseName).build();
final GrpcCallContext context =
- newCallContext(null, databaseName, request, DatabaseAdminGrpc.getDropDatabaseMethod());
+ newAdminCallContext(databaseName, request, DatabaseAdminGrpc.getDropDatabaseMethod());
runWithRetryOnAdministrativeRequestsExceeded(
() -> {
get(databaseAdminStub.dropDatabaseCallable().futureCall(request, context));
@@ -1356,7 +1353,7 @@ public Database getDatabase(String databaseName) throws SpannerException {
GetDatabaseRequest.newBuilder().setName(databaseName).build();
final GrpcCallContext context =
- newCallContext(null, databaseName, request, DatabaseAdminGrpc.getGetDatabaseMethod());
+ newAdminCallContext(databaseName, request, DatabaseAdminGrpc.getGetDatabaseMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.getDatabaseCallable().futureCall(request, context)));
}
@@ -1367,8 +1364,8 @@ public OperationFuture updateDatabase(
UpdateDatabaseRequest request =
UpdateDatabaseRequest.newBuilder().setDatabase(database).setUpdateMask(updateMask).build();
GrpcCallContext context =
- newCallContext(
- null, database.getName(), request, DatabaseAdminGrpc.getUpdateDatabaseMethod());
+ newAdminCallContext(
+ database.getName(), request, DatabaseAdminGrpc.getUpdateDatabaseMethod());
return databaseAdminStub.updateDatabaseOperationCallable().futureCall(request, context);
}
@@ -1379,7 +1376,7 @@ public GetDatabaseDdlResponse getDatabaseDdl(String databaseName) throws Spanner
GetDatabaseDdlRequest.newBuilder().setDatabase(databaseName).build();
final GrpcCallContext context =
- newCallContext(null, databaseName, request, DatabaseAdminGrpc.getGetDatabaseDdlMethod());
+ newAdminCallContext(databaseName, request, DatabaseAdminGrpc.getGetDatabaseDdlMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.getDatabaseDdlCallable().futureCall(request, context)));
}
@@ -1563,7 +1560,7 @@ public Backup updateBackup(Backup backup, FieldMask updateMask) {
final UpdateBackupRequest request =
UpdateBackupRequest.newBuilder().setBackup(backup).setUpdateMask(updateMask).build();
final GrpcCallContext context =
- newCallContext(null, backup.getName(), request, DatabaseAdminGrpc.getUpdateBackupMethod());
+ newAdminCallContext(backup.getName(), request, DatabaseAdminGrpc.getUpdateBackupMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> databaseAdminStub.updateBackupCallable().call(request, context));
}
@@ -1574,7 +1571,7 @@ public void deleteBackup(String backupName) {
final DeleteBackupRequest request =
DeleteBackupRequest.newBuilder().setName(backupName).build();
final GrpcCallContext context =
- newCallContext(null, backupName, request, DatabaseAdminGrpc.getDeleteBackupMethod());
+ newAdminCallContext(backupName, request, DatabaseAdminGrpc.getDeleteBackupMethod());
runWithRetryOnAdministrativeRequestsExceeded(
() -> {
databaseAdminStub.deleteBackupCallable().call(request, context);
@@ -1587,7 +1584,7 @@ public Backup getBackup(String backupName) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
final GetBackupRequest request = GetBackupRequest.newBuilder().setName(backupName).build();
final GrpcCallContext context =
- newCallContext(null, backupName, request, DatabaseAdminGrpc.getGetBackupMethod());
+ newAdminCallContext(backupName, request, DatabaseAdminGrpc.getGetBackupMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.getBackupCallable().futureCall(request, context)));
}
@@ -1597,7 +1594,7 @@ public Operation getOperation(String name) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
final GetOperationRequest request = GetOperationRequest.newBuilder().setName(name).build();
final GrpcCallContext context =
- newCallContext(null, name, request, OperationsGrpc.getGetOperationMethod());
+ newAdminCallContext(name, request, OperationsGrpc.getGetOperationMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() ->
get(
@@ -1613,7 +1610,7 @@ public void cancelOperation(String name) throws SpannerException {
final CancelOperationRequest request =
CancelOperationRequest.newBuilder().setName(name).build();
final GrpcCallContext context =
- newCallContext(null, name, request, OperationsGrpc.getCancelOperationMethod());
+ newAdminCallContext(name, request, OperationsGrpc.getCancelOperationMethod());
runWithRetryOnAdministrativeRequestsExceeded(
() -> {
get(
@@ -1718,10 +1715,16 @@ public StreamingCall read(
ReadRequest request,
ResultStreamConsumer consumer,
@Nullable Map options,
+ XGoogSpannerRequestId requestId,
boolean routeToLeader) {
GrpcCallContext context =
newCallContext(
- options, request.getSession(), request, SpannerGrpc.getReadMethod(), routeToLeader);
+ options,
+ requestId,
+ request.getSession(),
+ request,
+ SpannerGrpc.getReadMethod(),
+ routeToLeader);
SpannerResponseObserver responseObserver = new SpannerResponseObserver(consumer);
spannerStub.streamingReadCallable().call(request, responseObserver, context);
return new GrpcStreamingCall(context, responseObserver.getController());
@@ -1772,10 +1775,14 @@ public RetrySettings getPartitionedDmlRetrySettings() {
@Override
public ServerStream executeStreamingPartitionedDml(
- ExecuteSqlRequest request, Map options, Duration timeout) {
+ ExecuteSqlRequest request,
+ Map options,
+ XGoogSpannerRequestId requestId,
+ Duration timeout) {
GrpcCallContext context =
newCallContext(
options,
+ requestId,
request.getSession(),
request,
SpannerGrpc.getExecuteStreamingSqlMethod(),
@@ -1798,10 +1805,12 @@ public StreamingCall executeQuery(
ExecuteSqlRequest request,
ResultStreamConsumer consumer,
@Nullable Map options,
+ XGoogSpannerRequestId requestId,
boolean routeToLeader) {
GrpcCallContext context =
newCallContext(
options,
+ requestId,
request.getSession(),
request,
SpannerGrpc.getExecuteStreamingSqlMethod(),
@@ -1987,7 +1996,7 @@ private static T get(final Future future) throws SpannerException {
} catch (ExecutionException e) {
throw asSpannerException(e.getCause());
} catch (CancellationException e) {
- throw newSpannerException(context, e, null);
+ throw newSpannerException(context, e);
} catch (Exception exception) {
throw asSpannerException(exception);
}
@@ -2018,6 +2027,11 @@ GrpcCallContext newCallContext(@Nullable Map options, String resource
return newCallContext(options, resource, null, null);
}
+ private GrpcCallContext newAdminCallContext(
+ String resource, ReqT request, MethodDescriptor method) {
+ return newCallContext(null, resource, request, method, false);
+ }
+
@VisibleForTesting
GrpcCallContext newCallContext(
@Nullable Map options,
@@ -2034,6 +2048,17 @@ GrpcCallContext newCallContext(
ReqT request,
MethodDescriptor method,
boolean routeToLeader) {
+ return newCallContext(options, /* requestId= */ null, resource, request, method, routeToLeader);
+ }
+
+ @VisibleForTesting
+ GrpcCallContext newCallContext(
+ @Nullable Map options,
+ @Nullable XGoogSpannerRequestId requestId,
+ String resource,
+ ReqT request,
+ MethodDescriptor method,
+ boolean routeToLeader) {
GrpcCallContext context = this.baseGrpcCallContext;
Long affinity = options == null ? null : Option.CHANNEL_HINT.getLong(options);
if (affinity != null) {
@@ -2052,10 +2077,15 @@ GrpcCallContext newCallContext(
context = context.withChannelAffinity(affinity.intValue());
}
}
- if (options != null) {
- // TODO(@odeke-em): Infer the affinity if it doesn't match up with in the request-id.
- context = withRequestId(context, options);
+ int requestIdChannel = convertToRequestIdChannelNumber(affinity);
+ if (requestId == null) {
+ requestId = requestIdCreator.nextRequestId(requestIdChannel);
+ } else {
+ requestId.setChannelId(requestIdChannel);
}
+ context =
+ context.withCallOptions(
+ context.getCallOptions().withOption(REQUEST_ID_CALL_OPTIONS_KEY, requestId));
context = context.withExtraHeaders(metadataProvider.newExtraHeaders(resource, projectName));
if (routeToLeader && leaderAwareRoutingEnabled) {
context = context.withExtraHeaders(metadataProvider.newRouteToLeaderHeader());
@@ -2075,17 +2105,19 @@ GrpcCallContext newCallContext(
return (GrpcCallContext) context.merge(apiCallContextFromContext);
}
- GrpcCallContext withRequestId(GrpcCallContext context, Map options) {
- XGoogSpannerRequestId reqId = (XGoogSpannerRequestId) options.get(Option.REQUEST_ID);
- if (reqId == null) {
- return context;
- }
+ @Override
+ public RequestIdCreator getRequestIdCreator() {
+ return this.requestIdCreator;
+ }
- Map> withReqId =
- ImmutableMap.of(
- XGoogSpannerRequestId.REQUEST_HEADER_KEY.name(),
- Collections.singletonList(reqId.toString()));
- return context.withExtraHeaders(withReqId);
+ private int convertToRequestIdChannelNumber(@Nullable Long affinity) {
+ if (affinity == null) {
+ return 0;
+ }
+ int requestIdChannel = affinity.intValue();
+ requestIdChannel = requestIdChannel == Integer.MAX_VALUE ? 0 : Math.abs(requestIdChannel);
+ // Start counting at 1, to distinguish between '0 == Unknown and >0 == known'.
+ return requestIdChannel % this.numChannels + 1;
}
void registerResponseObserver(SpannerResponseObserver responseObserver) {
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java
index 985ad68a429..861e839a036 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java
@@ -188,7 +188,7 @@ private void recordSpan(Span span, String requestId) {
if (afeLatency != null) {
span.setAttribute("afe_latency", afeLatency.toString());
}
- span.setAttribute(XGoogSpannerRequestId.REQUEST_ID, requestId);
+ span.setAttribute(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME, requestId);
}
}
@@ -260,7 +260,7 @@ private DatabaseName extractDatabaseName(Metadata headers) throws ExecutionExcep
}
private String extractRequestId(Metadata headers) throws ExecutionException {
- return headers.get(XGoogSpannerRequestId.REQUEST_HEADER_KEY);
+ return headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY);
}
private TagContext getTagContext(String key, String method, DatabaseName databaseName)
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdCreatorImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdCreatorImpl.java
new file mode 100644
index 00000000000..5904fa581fd
--- /dev/null
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdCreatorImpl.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.spi.v1;
+
+import com.google.cloud.spanner.XGoogSpannerRequestId;
+import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator;
+import java.util.concurrent.atomic.AtomicLong;
+
+class RequestIdCreatorImpl implements RequestIdCreator {
+ private static final AtomicLong NEXT_CLIENT_ID = new AtomicLong();
+
+ private final long clientId = NEXT_CLIENT_ID.incrementAndGet();
+ private final AtomicLong requestId = new AtomicLong();
+
+ @Override
+ public long getClientId() {
+ return this.clientId;
+ }
+
+ @Override
+ public XGoogSpannerRequestId nextRequestId(long channelId) {
+ return XGoogSpannerRequestId.of(clientId, channelId, requestId.incrementAndGet(), 0);
+ }
+
+ @Override
+ public void reset() {
+ requestId.set(0);
+ }
+}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdInterceptor.java
new file mode 100644
index 00000000000..c32db2a2704
--- /dev/null
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdInterceptor.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.spi.v1;
+
+import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_CALL_OPTIONS_KEY;
+import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY;
+
+import com.google.cloud.spanner.XGoogSpannerRequestId;
+import io.grpc.CallOptions;
+import io.grpc.Channel;
+import io.grpc.ClientCall;
+import io.grpc.ClientInterceptor;
+import io.grpc.ForwardingClientCall;
+import io.grpc.Metadata;
+import io.grpc.Metadata.Key;
+import io.grpc.MethodDescriptor;
+import java.util.concurrent.atomic.AtomicLong;
+
+class RequestIdInterceptor implements ClientInterceptor {
+ static final CallOptions.Key ATTEMPT_KEY = CallOptions.Key.create("Attempt");
+ private static final String RESPONSE_ENCODING_KEY_NAME = "x-response-encoding";
+ private static final Key RESPONSE_ENCODING_KEY =
+ Key.of(RESPONSE_ENCODING_KEY_NAME, Metadata.ASCII_STRING_MARSHALLER);
+
+ RequestIdInterceptor() {}
+
+ @Override
+ public ClientCall interceptCall(
+ MethodDescriptor method, CallOptions callOptions, Channel next) {
+ return new ForwardingClientCall.SimpleForwardingClientCall(
+ next.newCall(method, callOptions)) {
+ @Override
+ public void start(Listener responseListener, Metadata headers) {
+ XGoogSpannerRequestId requestId = callOptions.getOption(REQUEST_ID_CALL_OPTIONS_KEY);
+ if (requestId != null) {
+ requestId.incrementAttempt();
+ headers.put(REQUEST_ID_HEADER_KEY, requestId.getHeaderValue());
+ }
+ super.start(responseListener, headers);
+ }
+ };
+ }
+}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerErrorInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerErrorInterceptor.java
index 549ea18a97f..9c3b2af2b06 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerErrorInterceptor.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerErrorInterceptor.java
@@ -17,6 +17,8 @@
package com.google.cloud.spanner.spi.v1;
import com.google.cloud.spanner.IsRetryableInternalError;
+import com.google.cloud.spanner.XGoogSpannerRequestId;
+import com.google.common.base.Strings;
import com.google.rpc.BadRequest;
import com.google.rpc.Help;
import com.google.rpc.LocalizedMessage;
@@ -35,6 +37,7 @@
import io.grpc.Status;
import io.grpc.Status.Code;
import io.grpc.protobuf.ProtoUtils;
+import java.util.Objects;
import java.util.logging.Level;
import java.util.logging.Logger;
@@ -70,7 +73,23 @@ public void start(Listener responseListener, Metadata headers) {
new SimpleForwardingClientCallListener(responseListener) {
@Override
public void onClose(Status status, Metadata trailers) {
+ // Return quickly if there is no error.
+ if (status.isOk()) {
+ super.onClose(status, trailers);
+ return;
+ }
try {
+ if (headers.containsKey(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY)) {
+ String requestId = headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY);
+ if (!Strings.isNullOrEmpty(requestId)) {
+ if (!trailers.containsKey(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY)) {
+ trailers.put(
+ XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY,
+ Objects.requireNonNull(
+ headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY)));
+ }
+ }
+ }
// Translate INTERNAL errors that should be retried to a retryable error code.
if (IsRetryableInternalError.INSTANCE.isRetryableInternalError(status)) {
status =
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java
index ec7a4e54a03..e8d6c3ebddb 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java
@@ -64,6 +64,7 @@ public static SpannerInterceptorProvider createDefault(
defaultInterceptorList.add(
new LoggingInterceptor(Logger.getLogger(GapicSpannerRpc.class.getName()), Level.FINER));
defaultInterceptorList.add(new HeaderInterceptor(new SpannerRpcMetrics(openTelemetry)));
+ defaultInterceptorList.add(new RequestIdInterceptor());
return new SpannerInterceptorProvider(ImmutableList.copyOf(defaultInterceptorList));
}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java
index 4b5682bb2b0..304729247f7 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java
@@ -27,6 +27,8 @@
import com.google.cloud.spanner.BackupId;
import com.google.cloud.spanner.Restore;
import com.google.cloud.spanner.SpannerException;
+import com.google.cloud.spanner.XGoogSpannerRequestId;
+import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
@@ -78,8 +80,7 @@
public interface SpannerRpc extends ServiceRpc {
/** Options passed in {@link SpannerRpc} methods to control how an RPC is issued. */
enum Option {
- CHANNEL_HINT("Channel Hint"),
- REQUEST_ID("Request Id");
+ CHANNEL_HINT("Channel Hint");
private final String value;
@@ -189,6 +190,10 @@ interface StreamingCall {
void cancel(@Nullable String message);
}
+ default RequestIdCreator getRequestIdCreator() {
+ throw new UnsupportedOperationException("Not implemented");
+ }
+
// Instance admin APIs.
Paginated listInstanceConfigs(int pageSize, @Nullable String pageToken)
throws SpannerException;
@@ -389,6 +394,7 @@ StreamingCall read(
ReadRequest request,
ResultStreamConsumer consumer,
@Nullable Map options,
+ XGoogSpannerRequestId requestId,
boolean routeToLeader);
/** Returns the retry settings for streaming query operations. */
@@ -428,7 +434,10 @@ ApiFuture executeQueryAsync(
RetrySettings getPartitionedDmlRetrySettings();
ServerStream executeStreamingPartitionedDml(
- ExecuteSqlRequest request, @Nullable Map options, Duration timeout);
+ ExecuteSqlRequest request,
+ @Nullable Map options,
+ XGoogSpannerRequestId requestId,
+ Duration timeout);
ServerStream batchWriteAtLeastOnce(
BatchWriteRequest request, @Nullable Map options);
@@ -445,6 +454,7 @@ StreamingCall executeQuery(
ExecuteSqlRequest request,
ResultStreamConsumer consumer,
@Nullable Map options,
+ XGoogSpannerRequestId requestId,
boolean routeToLeader);
ExecuteBatchDmlResponse executeBatchDml(ExecuteBatchDmlRequest build, Map options);
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java
index 939114a7f60..fc177ebb420 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java
@@ -27,6 +27,7 @@
import com.google.api.core.ApiFutures;
import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
import com.google.cloud.spanner.Options.TransactionOption;
+import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator;
import com.google.cloud.spanner.spi.v1.SpannerRpc.Option;
import com.google.protobuf.Empty;
import com.google.protobuf.Timestamp;
@@ -72,6 +73,7 @@ SessionImpl mockSession() {
final SessionImpl session = mock(SessionImpl.class);
Map options = new HashMap<>();
options.put(Option.CHANNEL_HINT, channelHint.getAndIncrement());
+ when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
when(session.getOptions()).thenReturn(options);
when(session.getName())
.thenReturn(
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
index f7842fc7a40..7b85c8e772e 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
@@ -42,7 +42,6 @@
import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
-import com.google.api.gax.grpc.testing.LocalChannelProvider;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.ApiCallContext;
import com.google.api.gax.rpc.ServerStream;
@@ -104,16 +103,18 @@
import com.google.spanner.v1.TypeAnnotationCode;
import com.google.spanner.v1.TypeCode;
import io.grpc.Context;
+import io.grpc.ManagedChannelBuilder;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import io.grpc.Server;
import io.grpc.ServerInterceptors;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
-import io.grpc.inprocess.InProcessServerBuilder;
+import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
import io.grpc.protobuf.lite.ProtoLiteUtils;
import io.opencensus.trace.Tracing;
import io.opentelemetry.api.OpenTelemetry;
+import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.Instant;
@@ -128,7 +129,6 @@
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
@@ -159,7 +159,6 @@ public class DatabaseClientImplTest {
private static XGoogSpannerRequestIdTest.ServerHeaderEnforcer xGoogReqIdInterceptor;
private static MockSpannerServiceImpl mockSpanner;
private static Server server;
- private static LocalChannelProvider channelProvider;
private static final Statement UPDATE_STATEMENT =
Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2");
private static final Statement INVALID_UPDATE_STATEMENT =
@@ -247,15 +246,12 @@ public static void startStaticServer() throws Exception {
"google.spanner.v1.Spanner/StreamingRead"));
xGoogReqIdInterceptor = new XGoogSpannerRequestIdTest.ServerHeaderEnforcer(checkMethods);
executor = Executors.newSingleThreadExecutor();
- String uniqueName = InProcessServerBuilder.generateName();
+ InetSocketAddress address = new InetSocketAddress("localhost", 0);
server =
- InProcessServerBuilder.forName(uniqueName)
- // We need to use a real executor for timeouts to occur.
- .scheduledExecutorService(new ScheduledThreadPoolExecutor(1))
+ NettyServerBuilder.forAddress(address)
.addService(ServerInterceptors.intercept(mockSpanner, xGoogReqIdInterceptor))
.build()
.start();
- channelProvider = LocalChannelProvider.create(uniqueName);
}
@AfterClass
@@ -267,11 +263,13 @@ public static void stopServer() throws InterruptedException {
@Before
public void setUp() {
+ String endpoint = "localhost:" + server.getPort();
spanner =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://" + endpoint)
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build())
.build()
@@ -321,7 +319,8 @@ public void tearDown() {
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -401,7 +400,8 @@ public void tearDown() {
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -473,7 +473,8 @@ public void testPoolMaintainer_whenLongRunningPartitionedUpdateRequest_takeNoAct
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -539,7 +540,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -631,7 +633,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -703,7 +706,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -765,7 +769,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -830,7 +835,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -909,7 +915,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -976,7 +983,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -1047,7 +1055,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -1114,7 +1123,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -1181,7 +1191,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -1250,7 +1261,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -1319,7 +1331,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
.setDatabaseRole(TEST_DATABASE_ROLE)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(sessionPoolOptions)
.build()
@@ -2905,7 +2918,8 @@ public void testPartitionedDmlDoesNotTimeout() {
SpannerOptions.Builder builder =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance());
// Set normal DML timeout value.
builder.getSpannerStubSettingsBuilder().executeSqlSettings().setRetrySettings(retrySettings);
@@ -2974,7 +2988,8 @@ public void testPartitionedDmlWithLowerTimeout() {
SpannerOptions.Builder builder =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance());
// Set PDML timeout value.
builder.setPartitionedDmlTimeoutDuration(Duration.ofMillis(10L));
@@ -3008,7 +3023,8 @@ public void testPartitionedDmlWithHigherTimeout() {
SpannerOptions.Builder builder =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance());
// Set PDML timeout value to a value that should allow the statement to be executed.
builder.setPartitionedDmlTimeoutDuration(Duration.ofMillis(5000L));
@@ -3088,7 +3104,8 @@ public void testPartitionedDmlRetriesOnUnavailable() {
SpannerOptions.Builder builder =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance());
try (Spanner spanner = builder.build().getService()) {
DatabaseClient client =
@@ -3111,7 +3128,8 @@ public void testDatabaseOrInstanceDoesNotExistOnInitialization() throws Exceptio
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.build()
.getService()) {
@@ -3156,7 +3174,8 @@ public void testDatabaseOrInstanceDoesNotExistOnCreate() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(
SessionPoolOptions.newBuilder()
@@ -3227,7 +3246,8 @@ public void testDatabaseOrInstanceDoesNotExistOnReplenish() throws Exception {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.build()
.getService()) {
@@ -3283,7 +3303,8 @@ public void testDatabaseOrInstanceIsDeletedAndThenRecreated() throws Exception {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.build()
.getService()) {
@@ -3372,7 +3393,8 @@ public void testGetInvalidatedClientMultipleTimes() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build())
.build()
@@ -3511,7 +3533,8 @@ public void testBackendQueryOptions() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId("[PROJECT]")
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build())
.build()
@@ -3552,7 +3575,8 @@ public void testBackendQueryOptionsWithAnalyzeQuery() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId("[PROJECT]")
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build())
.build()
@@ -3595,7 +3619,8 @@ public void testBackendPartitionQueryOptions() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId("[PROJECT]")
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build())
.setDirectedReadOptions(DIRECTED_READ_OPTIONS2)
@@ -3653,7 +3678,8 @@ public void testBackendPartitionQueryOptions() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId("[PROJECT]")
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build())
.setDirectedReadOptions(DIRECTED_READ_OPTIONS2)
@@ -3709,7 +3735,8 @@ public void testBackendPartitionReadOptions() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId("[PROJECT]")
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build())
.setDirectedReadOptions(DIRECTED_READ_OPTIONS2)
@@ -3759,7 +3786,8 @@ public void testBackendPartitionReadOptions() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId("[PROJECT]")
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build())
.setDirectedReadOptions(DIRECTED_READ_OPTIONS2)
@@ -3858,7 +3886,8 @@ public void testClientIdReusedOnDatabaseNotFound() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId("my-project")
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.build()
.getService()) {
@@ -3894,7 +3923,8 @@ public void testBatchCreateSessionsPermissionDenied() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId("my-project")
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(
SessionPoolOptions.newBuilder()
@@ -5271,7 +5301,8 @@ public void testRetryOnResourceExhausted() {
SpannerOptions.Builder builder =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance());
RetryInfo retryInfo =
RetryInfo.newBuilder()
@@ -5371,7 +5402,8 @@ public void testSessionPoolExhaustedError_containsStackTraces() {
try (Spanner spanner =
SpannerOptions.newBuilder()
.setProjectId(TEST_PROJECT)
- .setChannelProvider(channelProvider)
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
.setCredentials(NoCredentials.getInstance())
.setSessionPoolOption(
SessionPoolOptions.newBuilder()
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java
index 8f71b1278f3..c6155f0cbb6 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java
@@ -33,6 +33,7 @@
import com.google.api.gax.rpc.ServerStream;
import com.google.api.gax.rpc.UnavailableException;
import com.google.cloud.spanner.Options.RpcPriority;
+import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.common.base.Ticker;
import com.google.common.collect.ImmutableList;
@@ -96,11 +97,9 @@ public class PartitionedDmlTransactionTest {
public void setup() {
MockitoAnnotations.initMocks(this);
when(session.getName()).thenReturn(sessionId);
- when(session.getRequestIdCreator())
- .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator());
+ when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
when(session.getOptions()).thenReturn(Collections.EMPTY_MAP);
- when(session.getRequestIdCreator())
- .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator());
+ when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
when(rpc.beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)))
.thenReturn(Transaction.newBuilder().setId(txId).build());
@@ -115,7 +114,7 @@ public void testExecuteStreamingPartitionedUpdate() {
ServerStream stream = mock(ServerStream.class);
when(stream.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator());
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream);
long count = tx.executeStreamingPartitionedUpdate(Statement.of(sql), Duration.ofMinutes(10));
@@ -124,7 +123,7 @@ public void testExecuteStreamingPartitionedUpdate() {
verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
@@ -135,7 +134,7 @@ public void testExecuteStreamingPartitionedUpdateWithUpdateOptions() {
ServerStream stream = mock(ServerStream.class);
when(stream.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator());
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithRequestOptions), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithRequestOptions), anyMap(), any(), any(Duration.class)))
.thenReturn(stream);
long count =
@@ -146,7 +145,7 @@ public void testExecuteStreamingPartitionedUpdateWithUpdateOptions() {
verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithRequestOptions), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithRequestOptions), anyMap(), any(), any(Duration.class));
}
@Test
@@ -166,7 +165,7 @@ public void testExecuteStreamingPartitionedUpdateAborted() {
ServerStream stream2 = mock(ServerStream.class);
when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator());
when(rpc.executeStreamingPartitionedDml(
- any(ExecuteSqlRequest.class), anyMap(), any(Duration.class)))
+ any(ExecuteSqlRequest.class), anyMap(), any(), any(Duration.class)))
.thenReturn(stream1, stream2);
long count = tx.executeStreamingPartitionedUpdate(Statement.of(sql), Duration.ofMinutes(10));
@@ -175,7 +174,7 @@ public void testExecuteStreamingPartitionedUpdateAborted() {
verify(rpc, times(2)).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc, times(2))
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
@@ -195,10 +194,10 @@ public void testExecuteStreamingPartitionedUpdateUnavailable() {
ServerStream stream2 = mock(ServerStream.class);
when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator());
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream1);
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream2);
long count = tx.executeStreamingPartitionedUpdate(Statement.of(sql), Duration.ofMinutes(10));
@@ -207,10 +206,10 @@ public void testExecuteStreamingPartitionedUpdateUnavailable() {
verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
@@ -226,7 +225,7 @@ public void testExecuteStreamingPartitionedUpdateUnavailableAndThenDeadlineExcee
"temporary unavailable", null, GrpcStatusCode.of(Code.UNAVAILABLE), true));
when(stream1.iterator()).thenReturn(iterator);
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream1);
when(ticker.read()).thenReturn(0L, 1L, TimeUnit.NANOSECONDS.convert(10L, TimeUnit.MINUTES));
@@ -238,7 +237,7 @@ public void testExecuteStreamingPartitionedUpdateUnavailableAndThenDeadlineExcee
verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
@@ -254,7 +253,7 @@ public void testExecuteStreamingPartitionedUpdateAbortedAndThenDeadlineExceeded(
"transaction aborted", null, GrpcStatusCode.of(Code.ABORTED), true));
when(stream1.iterator()).thenReturn(iterator);
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream1);
when(ticker.read()).thenReturn(0L, 1L, TimeUnit.NANOSECONDS.convert(10L, TimeUnit.MINUTES));
@@ -266,7 +265,7 @@ public void testExecuteStreamingPartitionedUpdateAbortedAndThenDeadlineExceeded(
verify(rpc, times(2)).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
@@ -282,7 +281,7 @@ public void testExecuteStreamingPartitionedUpdateMultipleAbortsUntilDeadlineExce
"transaction aborted", null, GrpcStatusCode.of(Code.ABORTED), true));
when(stream1.iterator()).thenReturn(iterator);
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream1);
when(ticker.read())
.thenAnswer(
@@ -306,7 +305,7 @@ public Long answer(InvocationOnMock invocation) {
// means that the execute method is only executed 9 times.
verify(rpc, times(9))
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
@@ -329,10 +328,10 @@ public void testExecuteStreamingPartitionedUpdateUnexpectedEOS() {
ServerStream stream2 = mock(ServerStream.class);
when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator());
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream1);
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream2);
PartitionedDmlTransaction tx = new PartitionedDmlTransaction(session, rpc, ticker);
@@ -342,10 +341,10 @@ public void testExecuteStreamingPartitionedUpdateUnexpectedEOS() {
verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
@@ -368,10 +367,10 @@ public void testExecuteStreamingPartitionedUpdateRSTstream() {
ServerStream stream2 = mock(ServerStream.class);
when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator());
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream1);
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream2);
PartitionedDmlTransaction tx = new PartitionedDmlTransaction(session, rpc, ticker);
@@ -381,10 +380,10 @@ public void testExecuteStreamingPartitionedUpdateRSTstream() {
verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
@@ -400,7 +399,7 @@ public void testExecuteStreamingPartitionedUpdateGenericInternalException() {
"INTERNAL: Error", null, GrpcStatusCode.of(Code.INTERNAL), false));
when(stream1.iterator()).thenReturn(iterator);
when(rpc.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)))
.thenReturn(stream1);
PartitionedDmlTransaction tx = new PartitionedDmlTransaction(session, rpc, ticker);
@@ -412,7 +411,7 @@ public void testExecuteStreamingPartitionedUpdateGenericInternalException() {
verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true));
verify(rpc)
.executeStreamingPartitionedDml(
- Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+ Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class));
}
@Test
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RequestIdMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RequestIdMockServerTest.java
new file mode 100644
index 00000000000..c32a633331b
--- /dev/null
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RequestIdMockServerTest.java
@@ -0,0 +1,727 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner;
+
+import static com.google.cloud.spanner.DisableDefaultMtlsProvider.disableDefaultMtlsProvider;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeTrue;
+
+import com.google.cloud.NoCredentials;
+import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.common.collect.ImmutableList;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.ListValue;
+import com.google.protobuf.Value;
+import com.google.rpc.RetryInfo;
+import com.google.spanner.v1.BeginTransactionRequest;
+import com.google.spanner.v1.CommitRequest;
+import com.google.spanner.v1.CreateSessionRequest;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.ReadRequest;
+import com.google.spanner.v1.ResultSetMetadata;
+import com.google.spanner.v1.ResultSetStats;
+import com.google.spanner.v1.StructType;
+import com.google.spanner.v1.StructType.Field;
+import com.google.spanner.v1.Type;
+import com.google.spanner.v1.TypeCode;
+import io.grpc.Context;
+import io.grpc.Contexts;
+import io.grpc.ManagedChannelBuilder;
+import io.grpc.Metadata;
+import io.grpc.Server;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.Status;
+import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
+import io.grpc.protobuf.ProtoUtils;
+import java.net.InetSocketAddress;
+import java.util.List;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.TimeUnit;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.threeten.bp.Duration;
+
+@SuppressWarnings({"StatementWithEmptyBody", "resource"})
+@RunWith(JUnit4.class)
+public class RequestIdMockServerTest {
+ private static MockSpannerServiceImpl mockSpanner;
+ private static Server server;
+ private static Spanner spanner;
+
+ private static final Statement SELECT1 = Statement.of("SELECT 1");
+ private static final com.google.spanner.v1.ResultSet SELECT1_RESULT_SET =
+ com.google.spanner.v1.ResultSet.newBuilder()
+ .setMetadata(
+ ResultSetMetadata.newBuilder()
+ .setRowType(
+ StructType.newBuilder()
+ .addFields(
+ Field.newBuilder()
+ .setName("c")
+ .setType(Type.newBuilder().setCode(TypeCode.INT64).build())
+ .build())
+ .build())
+ .build())
+ .addRows(
+ ListValue.newBuilder()
+ .addValues(Value.newBuilder().setStringValue("1").build())
+ .build())
+ .build();
+ private static final Statement DML = Statement.of("insert into test_table (id) values (1)");
+
+ private static final ConcurrentLinkedQueue requestIds =
+ new ConcurrentLinkedQueue<>();
+
+ @BeforeClass
+ public static void setup() throws Exception {
+ assumeTrue(System.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS") == null);
+ assumeTrue(System.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW") == null);
+
+ disableDefaultMtlsProvider();
+ mockSpanner = new MockSpannerServiceImpl();
+ mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions.
+
+ InetSocketAddress address = new InetSocketAddress("localhost", 0);
+ server =
+ NettyServerBuilder.forAddress(address)
+ .addService(mockSpanner)
+ .intercept(
+ new ServerInterceptor() {
+ @Override
+ public ServerCall.Listener interceptCall(
+ ServerCall call,
+ Metadata headers,
+ ServerCallHandler next) {
+ try {
+ String requestId = headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY);
+ if (requestId != null) {
+ requestIds.add(XGoogSpannerRequestId.of(requestId));
+ } else {
+ requestIds.add(XGoogSpannerRequestId.of(0, 0, 0, 0));
+ }
+ } catch (Throwable t) {
+ // Ignore and continue
+ }
+ return Contexts.interceptCall(Context.current(), call, headers, next);
+ }
+ })
+ .build()
+ .start();
+ spanner = createSpanner();
+
+ setupResults();
+ }
+
+ private static Spanner createSpanner() {
+ return SpannerOptions.newBuilder()
+ .setProjectId("test-project")
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .setHost("http://localhost:" + server.getPort())
+ .setCredentials(NoCredentials.getInstance())
+ .setSessionPoolOption(
+ SessionPoolOptions.newBuilder()
+ .setFailOnSessionLeak()
+ .setMinSessions(0)
+ .setSkipVerifyingBeginTransactionForMuxRW(true)
+ .setWaitForMinSessions(Duration.ofSeconds(5))
+ .build())
+ .build()
+ .getService();
+ }
+
+ private static void setupResults() {
+ mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULT_SET));
+ mockSpanner.putStatementResult(StatementResult.update(DML, 1L));
+ }
+
+ static Metadata createMinimalRetryInfo() {
+ Metadata trailers = new Metadata();
+ RetryInfo retryInfo =
+ RetryInfo.newBuilder()
+ .setRetryDelay(
+ com.google.protobuf.Duration.newBuilder()
+ .setNanos((int) TimeUnit.MILLISECONDS.toNanos(1L))
+ .setSeconds(0L))
+ .build();
+ trailers.put(ProtoUtils.keyForProto(RetryInfo.getDefaultInstance()), retryInfo);
+ return trailers;
+ }
+
+ @AfterClass
+ public static void teardown() throws InterruptedException {
+ if (spanner != null) {
+ spanner.close();
+ }
+ if (server != null) {
+ server.shutdown();
+ server.awaitTermination();
+ }
+ }
+
+ @Before
+ public void prepareTest() {
+ // Call getClient() to make sure the multiplexed session has been created.
+ // Then clear all requests that were received as part of that so we don't need to include
+ // that in the test verifications.
+ getClient();
+ mockSpanner.reset();
+ requestIds.clear();
+ ((SpannerImpl) spanner).resetRequestIdCounters();
+ }
+
+ private DatabaseClient getClient() {
+ return spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+ }
+
+ private long getClientId() {
+ return ((SpannerImpl) spanner).getRequestIdClientId();
+ }
+
+ @Test
+ public void testSingleUseQuery() {
+ try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+
+ assertEquals(ImmutableList.of(ExecuteSqlRequest.class), mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(ImmutableList.of(XGoogSpannerRequestId.of(getClientId(), -1, 1, 1)), actual);
+ }
+
+ @Test
+ public void testQueryError() {
+ Statement query = Statement.of("select * from invalid_table");
+ mockSpanner.putStatementResult(
+ StatementResult.exception(
+ query, Status.NOT_FOUND.withDescription("Table not found").asRuntimeException()));
+
+ XGoogSpannerRequestId requestIdFromException;
+ try (ResultSet resultSet = getClient().singleUse().executeQuery(query)) {
+ SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
+ assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode());
+ assertNotNull(exception.getRequestId());
+ assertNotEquals("Request ID should not be empty", "", exception.getRequestId());
+ requestIdFromException = XGoogSpannerRequestId.of(exception.getRequestId());
+ }
+
+ assertEquals(ImmutableList.of(ExecuteSqlRequest.class), mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(ImmutableList.of(XGoogSpannerRequestId.of(getClientId(), -1, 1, 1)), actual);
+ assertEquals(actual.get(0), requestIdFromException);
+ }
+
+ @Test
+ public void testMultiUseReadOnlyTransaction() {
+ try (ReadOnlyTransaction transaction = getClient().readOnlyTransaction()) {
+ for (int i = 0; i < 2; i++) {
+ try (ResultSet resultSet = transaction.executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+ }
+ }
+
+ assertEquals(
+ ImmutableList.of(
+ BeginTransactionRequest.class, ExecuteSqlRequest.class, ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 3, 1)),
+ actual);
+ verifySameChannelId(actual);
+ }
+
+ @Test
+ public void testDml() {
+ getClient().readWriteTransaction().run(transaction -> transaction.executeUpdate(DML));
+
+ assertEquals(
+ ImmutableList.of(ExecuteSqlRequest.class, CommitRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)),
+ actual);
+ verifySameChannelId(actual);
+ }
+
+ @Test
+ public void testDmlError() {
+ Statement invalidDml = Statement.of("insert into invalid_table (id) values (1)");
+ mockSpanner.putStatementResult(
+ StatementResult.exception(
+ invalidDml, Status.NOT_FOUND.withDescription("Table not found").asRuntimeException()));
+
+ SpannerException exception =
+ assertThrows(
+ SpannerException.class,
+ () ->
+ getClient()
+ .readWriteTransaction()
+ .run(transaction -> transaction.executeUpdate(invalidDml)));
+ assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode());
+ assertNotNull(exception.getRequestId());
+ assertNotEquals("Request ID should not be empty", "", exception.getRequestId());
+ XGoogSpannerRequestId requestIdFromException =
+ XGoogSpannerRequestId.of(exception.getRequestId());
+
+ assertEquals(ImmutableList.of(ExecuteSqlRequest.class), mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(ImmutableList.of(XGoogSpannerRequestId.of(getClientId(), -1, 1, 1)), actual);
+ assertEquals(actual.get(0), requestIdFromException);
+ }
+
+ @Test
+ public void testAbortedTransaction() {
+ mockSpanner.setCommitExecutionTime(
+ SimulatedExecutionTime.ofException(
+ Status.ABORTED.asRuntimeException(createMinimalRetryInfo())));
+ getClient()
+ .readWriteTransaction()
+ .run(
+ transaction -> {
+ try (ResultSet resultSet = transaction.executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+ return transaction.executeUpdate(DML);
+ });
+
+ assertEquals(
+ ImmutableList.of(
+ ExecuteSqlRequest.class,
+ ExecuteSqlRequest.class,
+ CommitRequest.class,
+ ExecuteSqlRequest.class,
+ ExecuteSqlRequest.class,
+ CommitRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ int requestId = 0;
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1)),
+ actual);
+ verifySameChannelId(actual.subList(0, 3));
+ verifySameChannelId(actual.subList(3, 6));
+ }
+
+ @Test
+ public void testMix() {
+ getClient()
+ .readWriteTransaction()
+ .run(
+ transaction -> {
+ try (ResultSet resultSet = transaction.executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+ return transaction.executeUpdate(DML);
+ });
+ try (ReadOnlyTransaction transaction = getClient().readOnlyTransaction()) {
+ for (int i = 0; i < 2; i++) {
+ try (ResultSet resultSet = transaction.executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+ }
+ }
+ try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+ mockSpanner.putStatementResult(
+ StatementResult.query(
+ Statement.of("SELECT my_column FROM my_table WHERE 1=1"), SELECT1_RESULT_SET));
+ try (ResultSet resultSet =
+ getClient().singleUse().read("my_table", KeySet.all(), ImmutableList.of("my_column"))) {
+ while (resultSet.next()) {}
+ }
+
+ assertEquals(
+ ImmutableList.of(
+ ExecuteSqlRequest.class,
+ ExecuteSqlRequest.class,
+ CommitRequest.class,
+ BeginTransactionRequest.class,
+ ExecuteSqlRequest.class,
+ ExecuteSqlRequest.class,
+ ExecuteSqlRequest.class,
+ ReadRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ int requestId = 0;
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1)),
+ actual);
+ verifySameChannelId(actual.subList(0, 3));
+ verifySameChannelId(actual.subList(3, 6));
+ }
+
+ @Test
+ public void testUnaryUnavailable() {
+ mockSpanner.setExecuteSqlExecutionTime(
+ SimulatedExecutionTime.ofException(
+ Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo())));
+
+ getClient().readWriteTransaction().run(transaction -> transaction.executeUpdate(DML));
+
+ assertEquals(
+ ImmutableList.of(ExecuteSqlRequest.class, ExecuteSqlRequest.class, CommitRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 2),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)),
+ actual);
+ verifySameChannelId(actual);
+ }
+
+ @Test
+ public void testStreamingQueryUnavailable() {
+ mockSpanner.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofException(
+ Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo())));
+
+ try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+
+ assertEquals(
+ ImmutableList.of(ExecuteSqlRequest.class, ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 2)),
+ actual);
+ }
+
+ @Test
+ public void testStreamingQueryUnavailableHalfway() {
+ int numRows = 5;
+ Statement statement = Statement.of("select * from random");
+ mockSpanner.putStatementResult(
+ StatementResult.query(statement, new RandomResultSetGenerator(numRows).generate()));
+ mockSpanner.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofStreamException(
+ Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()), 2));
+
+ try (ResultSet resultSet = getClient().singleUse().executeQuery(statement)) {
+ while (resultSet.next()) {}
+ }
+
+ assertEquals(
+ ImmutableList.of(ExecuteSqlRequest.class, ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 2)),
+ actual);
+ List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+ assertEquals(ByteString.empty(), requests.get(0).getResumeToken());
+ assertNotEquals(ByteString.empty(), requests.get(1).getResumeToken());
+ }
+
+ @Test
+ public void testStreamingReadUnavailable() {
+ mockSpanner.setStreamingReadExecutionTime(
+ SimulatedExecutionTime.ofException(
+ Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo())));
+
+ mockSpanner.putStatementResult(
+ StatementResult.query(
+ Statement.of("SELECT my_column FROM my_table WHERE 1=1"), SELECT1_RESULT_SET));
+ try (ResultSet resultSet =
+ getClient().singleUse().read("my_table", KeySet.all(), ImmutableList.of("my_column"))) {
+ while (resultSet.next()) {}
+ }
+
+ assertEquals(
+ ImmutableList.of(ReadRequest.class, ReadRequest.class), mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 2)),
+ actual);
+ }
+
+ @Test
+ public void testStreamingReadUnavailableHalfway() {
+ int numRows = 5;
+ Statement statement = Statement.of("SELECT my_column FROM my_table WHERE 1=1");
+ mockSpanner.putStatementResult(
+ StatementResult.query(statement, new RandomResultSetGenerator(numRows).generate()));
+ mockSpanner.setStreamingReadExecutionTime(
+ SimulatedExecutionTime.ofStreamException(
+ Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()), 2));
+
+ try (ResultSet resultSet =
+ getClient().singleUse().read("my_table", KeySet.all(), ImmutableList.of("my_column"))) {
+ while (resultSet.next()) {}
+ }
+
+ assertEquals(
+ ImmutableList.of(ReadRequest.class, ReadRequest.class), mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 2)),
+ actual);
+ List requests = mockSpanner.getRequestsOfType(ReadRequest.class);
+ assertEquals(ByteString.empty(), requests.get(0).getResumeToken());
+ assertNotEquals(ByteString.empty(), requests.get(1).getResumeToken());
+ }
+
+ @Test
+ public void testPartitionedDml() {
+ getClient().executePartitionedUpdate(DML);
+
+ assertEquals(
+ ImmutableList.of(BeginTransactionRequest.class, ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)),
+ actual);
+ verifySameChannelId(actual);
+ }
+
+ @Test
+ public void testPartitionedDmlError() {
+ Statement invalidDml = Statement.of("update invalid_table set col=true where col=false");
+ mockSpanner.putStatementResult(
+ StatementResult.exception(
+ invalidDml, Status.NOT_FOUND.withDescription("Table not found").asRuntimeException()));
+
+ SpannerException exception =
+ assertThrows(
+ SpannerException.class, () -> getClient().executePartitionedUpdate(invalidDml));
+ assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode());
+ assertNotNull(exception.getRequestId());
+ assertNotEquals("", exception.getRequestId());
+
+ assertEquals(
+ ImmutableList.of(BeginTransactionRequest.class, ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)),
+ actual);
+ verifySameChannelId(actual);
+ assertEquals(XGoogSpannerRequestId.of(exception.getRequestId()), actual.get(1));
+ }
+
+ @Test
+ public void testPartitionedDmlAborted() {
+ mockSpanner.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofException(
+ Status.ABORTED.asRuntimeException(createMinimalRetryInfo())));
+
+ getClient().executePartitionedUpdate(DML);
+
+ assertEquals(
+ ImmutableList.of(
+ BeginTransactionRequest.class,
+ ExecuteSqlRequest.class,
+ BeginTransactionRequest.class,
+ ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 3, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 4, 1)),
+ actual);
+ verifySameChannelId(actual);
+ }
+
+ @Test
+ public void testPartitionedDmlUnavailable() {
+ mockSpanner.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofException(
+ Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo())));
+
+ getClient().executePartitionedUpdate(DML);
+
+ assertEquals(
+ ImmutableList.of(
+ BeginTransactionRequest.class,
+ ExecuteSqlRequest.class,
+ BeginTransactionRequest.class,
+ ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 3, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 4, 1)),
+ actual);
+ verifySameChannelId(actual);
+ }
+
+ @Test
+ public void testPartitionedDmlUnavailableWithResumeToken() {
+ Statement update = Statement.of("UPDATE my_table SET active=true where 1=1");
+ mockSpanner.putStatementResult(
+ StatementResult.query(
+ update,
+ com.google.spanner.v1.ResultSet.newBuilder()
+ .setMetadata(
+ ResultSetMetadata.newBuilder()
+ .setRowType(StructType.newBuilder().build())
+ .build())
+ .addRows(ListValue.newBuilder().build())
+ .addRows(ListValue.newBuilder().build())
+ .addRows(ListValue.newBuilder().build())
+ .setStats(ResultSetStats.newBuilder().setRowCountLowerBound(100L).build())
+ .build()));
+ mockSpanner.setExecuteStreamingSqlExecutionTime(
+ SimulatedExecutionTime.ofStreamException(
+ Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()), 2L));
+
+ getClient().executePartitionedUpdate(update);
+
+ assertEquals(
+ ImmutableList.of(
+ BeginTransactionRequest.class, ExecuteSqlRequest.class, ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 2)),
+ actual);
+ verifySameChannelId(actual);
+ }
+
+ @Test
+ public void testOtherClientId() {
+ // Execute a query with the default client from this test class.
+ try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+ // Create a new client and use that to execute a query. This should use a different client ID.
+ long otherClientId;
+ try (Spanner spanner = createSpanner()) {
+ otherClientId = ((SpannerImpl) spanner).getRequestIdClientId();
+ DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+ try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+ }
+ // Execute another query with the default client. This should use the original client ID.
+ try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) {
+ while (resultSet.next()) {}
+ }
+ assertEquals(
+ ImmutableList.of(
+ ExecuteSqlRequest.class,
+ CreateSessionRequest.class,
+ ExecuteSqlRequest.class,
+ ExecuteSqlRequest.class),
+ mockSpanner.getRequestTypes());
+ List actual = ImmutableList.copyOf(requestIds);
+ verifyRequestIds(
+ ImmutableList.of(
+ XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
+ // The CreateSession RPC from the initialization of the second client is included in
+ // the requests that we see. This request does not include a channel hint, hence the
+ // zero value for the channel number in the request ID.
+ XGoogSpannerRequestId.of(otherClientId, 0, 1, 1),
+ XGoogSpannerRequestId.of(otherClientId, -1, 2, 1),
+ XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)),
+ actual);
+ }
+
+ private void verifyRequestIds(
+ List expectedIds, List actualIds) {
+ assertEquals(message(expectedIds, actualIds), expectedIds.size(), actualIds.size());
+ int i = 0;
+ for (XGoogSpannerRequestId actual : actualIds) {
+ XGoogSpannerRequestId expected = expectedIds.get(i);
+ if (expected.getNthChannelId() > -1) {
+ assertEquals(expected, actual);
+ } else {
+ assertTrue(message(expectedIds, actualIds), equalsIgnoringChannelId(expected, actual));
+ assertTrue(message(expectedIds, actualIds), actual.hasChannelId());
+ }
+ i++;
+ }
+ }
+
+ private void verifySameChannelId(List requestIds) {
+ for (int i = 0; i < requestIds.size() - 1; i++) {
+ XGoogSpannerRequestId requestId = requestIds.get(i);
+ assertTrue(requestId.hasChannelId());
+ assertEquals(requestId.getNthChannelId(), requestIds.get(i + 1).getNthChannelId());
+ }
+ }
+
+ private boolean equalsIgnoringChannelId(
+ XGoogSpannerRequestId expected, XGoogSpannerRequestId actual) {
+ return expected.getNthClientId() == actual.getNthClientId()
+ && expected.getNthRequest() == actual.getNthRequest()
+ && expected.getAttempt() == actual.getAttempt();
+ }
+
+ private String message(List expected, List actual) {
+ return String.format("\n Got: %s\nWant: %s", actual, expected);
+ }
+}
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java
index 14f8df04e97..f13c0bb1237 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java
@@ -25,6 +25,7 @@
import com.google.api.client.util.BackOff;
import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler;
+import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator;
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.ImmutableList;
@@ -168,11 +169,12 @@ private void initWithLimit(int maxBufferSize) {
DefaultErrorHandler.INSTANCE,
SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(),
SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(),
- new XGoogSpannerRequestId.NoopRequestIdCreator()) {
+ NoopRequestIdCreator.INSTANCE) {
@Override
AbstractResultSet.CloseableIterator startStream(
@Nullable ByteString resumeToken,
- AsyncResultSet.StreamMessageListener streamMessageListener) {
+ AsyncResultSet.StreamMessageListener streamMessageListener,
+ XGoogSpannerRequestId requestId) {
return starter.startStream(resumeToken, null);
}
};
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionClientTests.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionClientTests.java
index 5d3ed2bca53..07a76970bfd 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionClientTests.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionClientTests.java
@@ -18,7 +18,7 @@
import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doNothing;
@@ -163,7 +163,6 @@ public void createAndCloseSession() {
assertEquals(
deleteOptionsCaptor.getValue().get(SpannerRpc.Option.CHANNEL_HINT),
options.getValue().get(SpannerRpc.Option.CHANNEL_HINT));
- assertTrue(deleteOptionsCaptor.getValue().containsKey(SpannerRpc.Option.REQUEST_ID));
}
}
@@ -207,9 +206,7 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount
client.createMultiplexedSession(consumer);
}
// for multiplexed session there is no channel hint pass in the RPC options
- assertNotNull(options.getValue());
- assertEquals(options.getValue().get(Option.CHANNEL_HINT), null);
- assertNotNull(options.getValue().get(Option.REQUEST_ID));
+ assertNull(options.getValue());
assertEquals(1, returnedSessionCount.get());
}
@@ -241,9 +238,7 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount
client.createMultiplexedSession(consumer);
}
// for multiplexed session there is no channel hint pass in the RPC options
- assertNotNull(options.getValue());
- assertEquals(options.getValue().get(Option.CHANNEL_HINT), null);
- assertNotNull(options.getValue().get(Option.REQUEST_ID));
+ assertNull(options.getValue());
}
@SuppressWarnings("unchecked")
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java
index df957f75ebf..1ac3b7beaf7 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java
@@ -21,6 +21,7 @@
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyMap;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
@@ -34,6 +35,7 @@
import com.google.cloud.Timestamp;
import com.google.cloud.grpc.GrpcTransportOptions;
import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
+import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
import com.google.protobuf.ByteString;
@@ -144,8 +146,8 @@ public void setUp() {
SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes());
when(rpc.getCommitRetrySettings())
.thenReturn(SpannerStubSettings.newBuilder().commitSettings().getRetrySettings());
+ when(rpc.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
session = spanner.getSessionClient(db).createSession();
- ((SessionImpl) session).setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator());
Span oTspan = mock(Span.class);
ISpan span = new OpenTelemetrySpan(oTspan);
when(oTspan.makeCurrent()).thenReturn(mock(Scope.class));
@@ -446,7 +448,7 @@ public void request(int numMessages) {}
private void mockRead(final PartialResultSet myResultSet) {
final ArgumentCaptor consumer =
ArgumentCaptor.forClass(SpannerRpc.ResultStreamConsumer.class);
- Mockito.when(rpc.read(Mockito.any(), consumer.capture(), anyMap(), eq(false)))
+ Mockito.when(rpc.read(Mockito.any(), consumer.capture(), anyMap(), any(), eq(false)))
.then(
invocation -> {
consumer.getValue().onPartialResultSet(myResultSet);
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
index 8d00f0889b8..b9d6943e523 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
@@ -72,6 +72,7 @@
import com.google.cloud.spanner.SpannerImpl.ClosedException;
import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl;
+import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.cloud.spanner.spi.v1.SpannerRpc.ResultStreamConsumer;
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
@@ -1455,6 +1456,7 @@ public void testSessionNotFoundReadWriteTransaction() {
any(ExecuteSqlRequest.class),
any(ResultStreamConsumer.class),
any(Map.class),
+ any(XGoogSpannerRequestId.class),
eq(true)))
.thenReturn(closedStreamingCall);
when(rpc.executeQuery(any(ExecuteSqlRequest.class), any(Map.class), eq(true)))
@@ -1481,8 +1483,7 @@ public void testSessionNotFoundReadWriteTransaction() {
when(closedSession.getName())
.thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-closed");
when(closedSession.getErrorHandler()).thenReturn(DefaultErrorHandler.INSTANCE);
- when(closedSession.getRequestIdCreator())
- .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator());
+ when(closedSession.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
Span oTspan = mock(Span.class);
ISpan span = new OpenTelemetrySpan(oTspan);
@@ -1523,8 +1524,7 @@ public void testSessionNotFoundReadWriteTransaction() {
TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession);
openTransactionRunner.setSpan(span);
when(openSession.readWriteTransaction()).thenReturn(openTransactionRunner);
- when(openSession.getRequestIdCreator())
- .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator());
+ when(openSession.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
ResultSet openResultSet = mock(ResultSet.class);
when(openResultSet.next()).thenReturn(true, false);
@@ -1648,14 +1648,12 @@ public void testSessionNotFoundWrite() {
SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
List mutations = Collections.singletonList(Mutation.newInsertBuilder("FOO").build());
final SessionImpl closedSession = mockSession();
- closedSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator());
when(closedSession.writeWithOptions(eq(mutations), any())).thenThrow(sessionNotFound);
final SessionImpl openSession = mockSession();
com.google.cloud.spanner.CommitResponse response =
mock(com.google.cloud.spanner.CommitResponse.class);
when(response.getCommitTimestamp()).thenReturn(Timestamp.now());
- openSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator());
when(openSession.writeWithOptions(eq(mutations), any())).thenReturn(response);
doAnswer(
invocation -> {
@@ -1693,7 +1691,6 @@ public void testSessionNotFoundWriteAtLeastOnce() {
SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
List mutations = Collections.singletonList(Mutation.newInsertBuilder("FOO").build());
final SessionImpl closedSession = mockSession();
- closedSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator());
when(closedSession.writeAtLeastOnceWithOptions(eq(mutations), any()))
.thenThrow(sessionNotFound);
@@ -1701,7 +1698,6 @@ public void testSessionNotFoundWriteAtLeastOnce() {
com.google.cloud.spanner.CommitResponse response =
mock(com.google.cloud.spanner.CommitResponse.class);
when(response.getCommitTimestamp()).thenReturn(Timestamp.now());
- openSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator());
when(openSession.writeAtLeastOnceWithOptions(eq(mutations), any())).thenReturn(response);
doAnswer(
invocation -> {
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerCloudMonitoringExporterTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerCloudMonitoringExporterTest.java
index c52cab0546c..590d62db7b5 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerCloudMonitoringExporterTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerCloudMonitoringExporterTest.java
@@ -371,7 +371,7 @@ public void testExportingHistogramDataWithExemplars() {
DoubleExemplarData exemplar =
ImmutableDoubleExemplarData.create(
Attributes.builder()
- .put(XGoogSpannerRequestId.REQUEST_ID, "test")
+ .put(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME, "test")
.put("lang", "java")
.build(),
recordTimeEpoch,
@@ -448,8 +448,10 @@ public void testExportingHistogramDataWithExemplars() {
// Assert only 1 attachment is there with 1 label for request_id.
assertThat(filterAttributes.size()).isEqualTo(1);
assertThat(filterAttributes.get(0).getLabelCount()).isEqualTo(1);
- assertThat(filterAttributes.get(0).containsLabel(XGoogSpannerRequestId.REQUEST_ID)).isTrue();
- assertThat(filterAttributes.get(0).getLabelOrThrow(XGoogSpannerRequestId.REQUEST_ID))
+ assertThat(filterAttributes.get(0).containsLabel(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME))
+ .isTrue();
+ assertThat(
+ filterAttributes.get(0).getLabelOrThrow(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME))
.isEqualTo("test");
}
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java
index 15af6bb8f1b..55b9523d7db 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java
@@ -251,17 +251,4 @@ private Metadata createResourceTypeMetadata(String resourceType, String resource
return trailers;
}
-
- @Test
- public void withRequestId() {
- XGoogSpannerRequestId reqIdIn = XGoogSpannerRequestId.of(1, 2, 3, 4);
- Status status = Status.fromCodeValue(Status.Code.ABORTED.value());
- Exception exc = new StatusRuntimeException(status);
- SpannerException spannerExceptionWithReqId =
- SpannerExceptionFactory.newSpannerException(exc, reqIdIn);
- assertThat(spannerExceptionWithReqId.getRequestId()).isEqualTo(reqIdIn.toString());
- SpannerException spannerExceptionWithoutReqId =
- SpannerExceptionFactory.newSpannerException(exc);
- assertThat(spannerExceptionWithoutReqId.getRequestId()).isEqualTo("");
- }
}
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java
index 4c507b7beae..bbfa7cbd0d5 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java
@@ -27,6 +27,7 @@
import com.google.api.core.ApiFutures;
import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl;
+import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
import com.google.protobuf.ByteString;
@@ -66,9 +67,9 @@ public void setup() {
com.google.spanner.v1.CommitResponse.newBuilder()
.setCommitTimestamp(Timestamp.newBuilder().setSeconds(99L).setNanos(10).build())
.build()));
+ when(rpc.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
when(session.getName()).thenReturn("test");
- when(session.getRequestIdCreator())
- .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator());
+ when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
doNothing().when(span).setStatus(any(Throwable.class));
doNothing().when(span).end();
doNothing().when(span).addAnnotation("Starting Commit");
@@ -212,8 +213,7 @@ public void testReturnCommitStats() {
private void batchDml(int status) {
SessionImpl session = mock(SessionImpl.class);
when(session.getName()).thenReturn("test");
- when(session.getRequestIdCreator())
- .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator());
+ when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
SpannerRpc rpc = mock(SpannerRpc.class);
ExecuteBatchDmlResponse response =
ExecuteBatchDmlResponse.newBuilder()
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java
index 19c03859e66..547f6b70a22 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java
@@ -241,7 +241,7 @@ public void usesPreparedTransaction() {
Mockito.anyString(),
Mockito.anyString(),
Mockito.anyMap(),
- Mockito.anyMap(),
+ Mockito.eq(null),
Mockito.eq(true)))
.thenAnswer(
invocation ->
@@ -324,7 +324,7 @@ public void inlineBegin() {
Mockito.anyString(),
Mockito.anyString(),
Mockito.anyMap(),
- Mockito.anyMap(),
+ Mockito.eq(null),
Mockito.eq(true)))
.thenAnswer(
invocation ->
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java
index 2325f2ac40f..3e3358a53bc 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java
@@ -36,6 +36,7 @@
import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler;
import com.google.cloud.spanner.SessionClient.SessionId;
import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl;
+import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
import com.google.common.base.Preconditions;
@@ -121,8 +122,8 @@ public void setUp() {
when(session.getErrorHandler()).thenReturn(DefaultErrorHandler.INSTANCE);
when(session.newTransaction(eq(Options.fromTransactionOptions()), any())).thenReturn(txn);
when(session.getTracer()).thenReturn(tracer);
- when(session.getRequestIdCreator())
- .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator());
+ when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
+ when(rpc.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE);
when(rpc.executeQuery(Mockito.any(ExecuteSqlRequest.class), Mockito.anyMap(), eq(true)))
.thenAnswer(
invocation -> {
@@ -195,7 +196,7 @@ public void usesPreparedTransaction() {
Mockito.anyString(),
Mockito.anyString(),
Mockito.anyMap(),
- Mockito.anyMap(),
+ Mockito.eq(null),
Mockito.eq(true)))
.thenAnswer(
invocation ->
@@ -336,7 +337,6 @@ public void inlineBegin() {
spanner,
new SessionReference(
"projects/p/instances/i/databases/d/sessions/s", Collections.EMPTY_MAP)) {};
- session.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator());
session.setCurrentSpan(new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)));
TransactionRunnerImpl runner = new TransactionRunnerImpl(session);
runner.setSpan(span);
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java
index 719b94593bf..32d1ac29d2e 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java
@@ -16,6 +16,7 @@
package com.google.cloud.spanner;
+import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
@@ -67,17 +68,14 @@ public void testEnsureHexadecimalFormatForRandProcessID() {
}
public static class ServerHeaderEnforcer implements ServerInterceptor {
- private Map> unaryResults;
- private Map> streamingResults;
- private List gotValues;
- private Set checkMethods;
+ private final Map> unaryResults =
+ new ConcurrentHashMap<>();
+ private final Map> streamingResults =
+ new ConcurrentHashMap<>();
+ private final List gotValues = new CopyOnWriteArrayList<>();
+ private final Set checkMethods;
ServerHeaderEnforcer(Set checkMethods) {
- this.gotValues = new CopyOnWriteArrayList();
- this.unaryResults =
- new ConcurrentHashMap>();
- this.streamingResults =
- new ConcurrentHashMap>();
this.checkMethods = checkMethods;
}
@@ -88,7 +86,7 @@ public ServerCall.Listener interceptCall(
ServerCallHandler next) {
boolean isUnary = call.getMethodDescriptor().getType() == MethodType.UNARY;
String methodName = call.getMethodDescriptor().getFullMethodName();
- String gotReqIdStr = requestHeaders.get(XGoogSpannerRequestId.REQUEST_HEADER_KEY);
+ String gotReqIdStr = requestHeaders.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY);
if (!this.checkMethods.contains(methodName)) {
return next.startCall(call, requestHeaders);
}
@@ -102,7 +100,7 @@ public ServerCall.Listener interceptCall(
Status status =
Status.fromCode(Status.Code.INVALID_ARGUMENT)
.augmentDescription(
- methodName + " lacks " + XGoogSpannerRequestId.REQUEST_HEADER_KEY);
+ methodName + " lacks " + XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY);
call.close(status, requestHeaders);
return next.startCall(call, requestHeaders);
}
@@ -128,14 +126,8 @@ public String[] accumulatedValues() {
}
public void assertIntegrity() {
- this.unaryResults.forEach(
- (String method, CopyOnWriteArrayList values) -> {
- assertMonotonicityOfIds(method, values);
- });
- this.streamingResults.forEach(
- (String method, CopyOnWriteArrayList values) -> {
- assertMonotonicityOfIds(method, values);
- });
+ this.unaryResults.forEach(this::assertMonotonicityOfIds);
+ this.streamingResults.forEach(this::assertMonotonicityOfIds);
}
private void assertMonotonicityOfIds(String prefix, List reqIds) {
@@ -161,22 +153,22 @@ private void assertMonotonicityOfIds(String prefix, List
}
public MethodAndRequestId[] accumulatedUnaryValues() {
- List accumulated = new ArrayList();
+ List accumulated = new ArrayList<>();
this.unaryResults.forEach(
(String method, CopyOnWriteArrayList values) -> {
- for (int i = 0; i < values.size(); i++) {
- accumulated.add(new MethodAndRequestId(method, values.get(i)));
+ for (XGoogSpannerRequestId value : values) {
+ accumulated.add(new MethodAndRequestId(method, value));
}
});
return accumulated.toArray(new MethodAndRequestId[0]);
}
public MethodAndRequestId[] accumulatedStreamingValues() {
- List accumulated = new ArrayList();
+ List accumulated = new ArrayList<>();
this.streamingResults.forEach(
(String method, CopyOnWriteArrayList values) -> {
- for (int i = 0; i < values.size(); i++) {
- accumulated.add(new MethodAndRequestId(method, values.get(i)));
+ for (XGoogSpannerRequestId value : values) {
+ accumulated.add(new MethodAndRequestId(method, value));
}
});
return accumulated.toArray(new MethodAndRequestId[0]);
@@ -188,7 +180,7 @@ public void checkExpectedUnaryXGoogRequestIds(MethodAndRequestId... wantUnaryVal
for (int i = 0; i < gotUnaryValues.length && false; i++) {
System.out.println("\033[33misUnary: #" + i + ":: " + gotUnaryValues[i] + "\033[00m");
}
- assertEquals(wantUnaryValues, gotUnaryValues);
+ assertArrayEquals(wantUnaryValues, gotUnaryValues);
}
public void checkAtLeastHasExpectedUnaryXGoogRequestIds(MethodAndRequestId... wantUnaryValues) {
@@ -200,9 +192,9 @@ public void checkAtLeastHasExpectedUnaryXGoogRequestIds(MethodAndRequestId... wa
if (wantUnaryValues.length < gotUnaryValues.length) {
MethodAndRequestId[] gotSliced =
Arrays.copyOfRange(gotUnaryValues, 0, wantUnaryValues.length);
- assertEquals(wantUnaryValues, gotSliced);
+ assertArrayEquals(wantUnaryValues, gotSliced);
} else {
- assertEquals(wantUnaryValues, gotUnaryValues);
+ assertArrayEquals(wantUnaryValues, gotUnaryValues);
}
}
@@ -218,9 +210,9 @@ public void checkExpectedUnaryXGoogRequestIdsAsSuffixes(MethodAndRequestId... wa
gotUnaryValues,
gotUnaryValues.length - wantUnaryValues.length,
gotUnaryValues.length);
- assertEquals(wantUnaryValues, gotSliced);
+ assertArrayEquals(wantUnaryValues, gotSliced);
} else {
- assertEquals(wantUnaryValues, gotUnaryValues);
+ assertArrayEquals(wantUnaryValues, gotUnaryValues);
}
}
@@ -236,7 +228,7 @@ public void checkExpectedStreamingXGoogRequestIds(MethodAndRequestId... wantStre
"\033[32misStreaming: #" + i + ":: " + gotStreamingValues[i] + "\033[00m");
}
sortValues(gotStreamingValues);
- assertEquals(wantStreamingValues, gotStreamingValues);
+ assertArrayEquals(wantStreamingValues, gotStreamingValues);
}
public void reset() {
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RunTransactionMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RunTransactionMockServerTest.java
index bc0acca3255..d4af5dd4e7e 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RunTransactionMockServerTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RunTransactionMockServerTest.java
@@ -151,7 +151,7 @@ public void testRunTransaction_rollbacksAfterException() {
return null;
}));
assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode());
- assertTrue(exception.getMessage(), exception.getMessage().endsWith("invalid statement"));
+ assertTrue(exception.getMessage(), exception.getMessage().contains("invalid statement"));
}
assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));