fix: retry as PDML dit not retry Resource limit exceeded (googleapis#4258)

Most transactions that exceed the mutation limit for an atomic transaction will fail with the error "The transaction contains too many mutations.". However, it is also possible that the transaction fails with the more generic error message "Transaction resource limits exceeded". This error did not trigger a retry of the statement using a PDML transaction.

Fixes googleapis#4253

Author: olavloite

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

@@ -28,6 +28,9 @@ public class TransactionMutationLimitExceededException extends SpannerException
 
   private static final String ERROR_MESSAGE = "The transaction contains too many mutations.";
 
+  private static final String TRANSACTION_RESOURCE_LIMIT_EXCEEDED_MESSAGE =
+      "Transaction resource limits exceeded";
+
   /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
   TransactionMutationLimitExceededException(
       DoNotConstructDirectly token,
@@ -40,13 +43,17 @@ public class TransactionMutationLimitExceededException extends SpannerException
   }
 
   static boolean isTransactionMutationLimitException(ErrorCode code, String message) {
-    return code == ErrorCode.INVALID_ARGUMENT && message != null && message.contains(ERROR_MESSAGE);
+    return code == ErrorCode.INVALID_ARGUMENT
+        && message != null
+        && (message.contains(ERROR_MESSAGE)
+            || message.contains(TRANSACTION_RESOURCE_LIMIT_EXCEEDED_MESSAGE));
   }
 
   static boolean isTransactionMutationLimitException(Throwable cause, ApiException apiException) {
     if (cause == null
         || cause.getMessage() == null
-        || !cause.getMessage().contains(ERROR_MESSAGE)) {
+        || !(cause.getMessage().contains(ERROR_MESSAGE)
+            || cause.getMessage().contains(TRANSACTION_RESOURCE_LIMIT_EXCEEDED_MESSAGE))) {
       return false;
     }
     // Spanner includes a hint that points to the Spanner limits documentation page when the error
@@ -66,6 +73,9 @@ static boolean isTransactionMutationLimitException(Throwable cause, ApiException
               .getLinks(0)
               .getUrl()
               .equals("https://cloud.google.com/spanner/docs/limits");
+    } else if (cause.getMessage().contains(TRANSACTION_RESOURCE_LIMIT_EXCEEDED_MESSAGE)) {
+      // This more generic error does not contain any additional details.
+      return true;
     }
     return false;
   }

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

@@ -43,10 +43,37 @@
 import io.grpc.StatusRuntimeException;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
 
-@RunWith(JUnit4.class)
+@RunWith(Parameterized.class)
 public class RetryDmlAsPartitionedDmlMockServerTest extends AbstractMockServerTest {
+  private enum ExceptionType {
+    MutationLimitExceeded {
+      @Override
+      StatusRuntimeException createException() {
+        return createTransactionMutationLimitExceededException();
+      }
+    },
+    ResourceLimitExceeded {
+      @Override
+      StatusRuntimeException createException() {
+        return createTransactionResourceLimitExceededException();
+      }
+    };
+
+    abstract StatusRuntimeException createException();
+  }
+
+  @Parameters(name = "exception = {0}")
+  public static Object[] data() {
+    return ExceptionType.values();
+  }
+
+  @SuppressWarnings("ClassEscapesDefinedScope")
+  @Parameter
+  public ExceptionType exceptionType;
 
   static StatusRuntimeException createTransactionMutationLimitExceededException() {
     Metadata.Key<byte[]> key =
@@ -70,10 +97,16 @@ static StatusRuntimeException createTransactionMutationLimitExceededException()
         .asRuntimeException(trailers);
   }
 
+  static StatusRuntimeException createTransactionResourceLimitExceededException() {
+    return Status.INVALID_ARGUMENT
+        .withDescription("Transaction resource limits exceeded")
+        .asRuntimeException();
+  }
+
   @Test
   public void testTransactionMutationLimitExceeded_isNotRetriedByDefault() {
     mockSpanner.setExecuteSqlExecutionTime(
-        SimulatedExecutionTime.ofException(createTransactionMutationLimitExceededException()));
+        SimulatedExecutionTime.ofException(exceptionType.createException()));
 
     try (Connection connection = createConnection()) {
       connection.setAutocommit(true);
@@ -95,7 +128,7 @@ public void testTransactionMutationLimitExceeded_isNotRetriedByDefault() {
   public void testTransactionMutationLimitExceeded_canBeRetriedAsPDML() {
     Statement statement = Statement.of("update test set value=1 where true");
     mockSpanner.setExecuteSqlExecutionTime(
-        SimulatedExecutionTime.ofException(createTransactionMutationLimitExceededException()));
+        SimulatedExecutionTime.ofException(exceptionType.createException()));
     mockSpanner.putStatementResult(
         MockSpannerServiceImpl.StatementResult.update(statement, 100000L));
 
@@ -134,7 +167,7 @@ public void testTransactionMutationLimitExceeded_retryAsPDMLFails() {
     Statement statement = Statement.of("insert into test (id, value) select -id, value from test");
     // The transactional update statement uses ExecuteSql(..).
     mockSpanner.setExecuteSqlExecutionTime(
-        SimulatedExecutionTime.ofException(createTransactionMutationLimitExceededException()));
+        SimulatedExecutionTime.ofException(exceptionType.createException()));
     mockSpanner.putStatementResult(
         MockSpannerServiceImpl.StatementResult.exception(
             statement,
@@ -230,7 +263,7 @@ public void testTransactionMutationLimitExceeded_isWrappedAsCauseOfBatchUpdateEx
     Statement statement = Statement.of(sql);
     mockSpanner.putStatementResult(
         MockSpannerServiceImpl.StatementResult.exception(
-            statement, createTransactionMutationLimitExceededException()));
+            statement, exceptionType.createException()));
 
     try (Connection connection = createConnection()) {
       connection.setAutocommit(true);