feat: add TransactionMutationLimitExceededException as cause to SpannerBatchUpdateException (#3723)

If a DML batch fails due to a TransactionMutationLimitExceededException, then add that specific exception as the cause to the SpannerBatchUpdateException. This makes it easier to detect this specific error, and potentially retry the individual statements as PDML.

Author: olavloite

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java

@@ -17,11 +17,16 @@
 package com.google.cloud.spanner;
 
 public class SpannerBatchUpdateException extends SpannerException {
- private long[] updateCounts;
+ private final long[] updateCounts;
+
  /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
  SpannerBatchUpdateException(
- DoNotConstructDirectly token, ErrorCode code, String message, long[] counts) {
- super(token, code, false, message, null);
+ DoNotConstructDirectly token,
+ ErrorCode code,
+ String message,
+ long[] counts,
+ Throwable cause) {
+ super(token, code, false, message, cause);
  updateCounts = counts;
  }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java

@@ -118,7 +118,11 @@ public static SpannerException newSpannerException(Throwable cause) {
  public static SpannerBatchUpdateException newSpannerBatchUpdateException(
  ErrorCode code, String message, long[] updateCounts) {
  DoNotConstructDirectly token = DoNotConstructDirectly.ALLOWED;
- return new SpannerBatchUpdateException(token, code, message, updateCounts);
+ SpannerException cause = null;
+ if (isTransactionMutationLimitException(code, message)) {
+ cause = new TransactionMutationLimitExceededException(token, code, message, null, null);
+ }
+ return new SpannerBatchUpdateException(token, code, message, updateCounts, cause);
  }
 
  /** Constructs a specific error that */

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java

@@ -26,6 +26,8 @@
 public class TransactionMutationLimitExceededException extends SpannerException {
  private static final long serialVersionUID = 1L;
 
+ private static final String ERROR_MESSAGE = "The transaction contains too many mutations.";
+
  /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
  TransactionMutationLimitExceededException(
  DoNotConstructDirectly token,
@@ -36,10 +38,14 @@ public class TransactionMutationLimitExceededException extends SpannerException
  super(token, errorCode, /*retryable = */ false, message, cause, apiException);
  }
 
+ static boolean isTransactionMutationLimitException(ErrorCode code, String message) {
+ return code == ErrorCode.INVALID_ARGUMENT && message != null && message.contains(ERROR_MESSAGE);
+ }
+
  static boolean isTransactionMutationLimitException(Throwable cause) {
  if (cause == null
  || cause.getMessage() == null
- || !cause.getMessage().contains("The transaction contains too many mutations.")) {
+ || !cause.getMessage().contains(ERROR_MESSAGE)) {
  return false;
  }
  // Spanner includes a hint that points to the Spanner limits documentation page when the error

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java

@@ -1070,7 +1070,7 @@ public long[] batchUpdate(Iterable<Statement> statements, UpdateOption... update
  // In all other cases, we should throw a BatchUpdateException.
  if (response.getStatus().getCode() == Code.ABORTED_VALUE) {
  throw createAbortedExceptionForBatchDml(response);
- } else if (response.getStatus().getCode() != 0) {
+ } else if (response.getStatus().getCode() != Code.OK_VALUE) {
  throw newSpannerBatchUpdateException(
  ErrorCode.fromRpcStatus(response.getStatus()),
  response.getStatus().getMessage(),
@@ -1137,7 +1137,7 @@ public ApiFuture<long[]> batchUpdateAsync(
  // In all other cases, we should throw a BatchUpdateException.
  if (batchDmlResponse.getStatus().getCode() == Code.ABORTED_VALUE) {
  throw createAbortedExceptionForBatchDml(batchDmlResponse);
- } else if (batchDmlResponse.getStatus().getCode() != 0) {
+ } else if (batchDmlResponse.getStatus().getCode() != Code.OK_VALUE) {
  throw newSpannerBatchUpdateException(
  ErrorCode.fromRpcStatus(batchDmlResponse.getStatus()),
  batchDmlResponse.getStatus().getMessage(),

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RetryDmlAsPartitionedDmlMockServerTest.java

@@ -18,6 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
@@ -26,6 +27,7 @@
 import com.google.cloud.spanner.MockSpannerServiceImpl;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.SpannerBatchUpdateException;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.TransactionMutationLimitExceededException;
@@ -34,6 +36,7 @@
 import com.google.rpc.Help.Link;
 import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CommitRequest;
+import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import io.grpc.Metadata;
 import io.grpc.Status;
@@ -219,4 +222,29 @@ public void testSqlStatements() {
  }
  }
  }
+
+ @Test
+ public void testTransactionMutationLimitExceeded_isWrappedAsCauseOfBatchUpdateException() {
+ String sql = "update test set value=1 where true";
+ Statement statement = Statement.of(sql);
+ mockSpanner.putStatementResult(
+ MockSpannerServiceImpl.StatementResult.exception(
+ statement, createTransactionMutationLimitExceededException()));
+
+ try (Connection connection = createConnection()) {
+ connection.setAutocommit(true);
+ assertEquals(AutocommitDmlMode.TRANSACTIONAL, connection.getAutocommitDmlMode());
+
+ connection.startBatchDml();
+ connection.execute(statement);
+ SpannerBatchUpdateException batchUpdateException =
+ assertThrows(SpannerBatchUpdateException.class, connection::runBatch);
+ assertNotNull(batchUpdateException.getCause());
+ assertEquals(
+ TransactionMutationLimitExceededException.class,
+ batchUpdateException.getCause().getClass());
+ }
+ assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+ assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+ }
 }