feat: add TransactionMutationLimitExceededException as cause to SpannerBatchUpdateException

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));
+  }
 }