From 3f8884a9635a6c1522139343208b8c646c7b072e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Wed, 2 Apr 2025 16:17:08 +0200 Subject: [PATCH] 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. --- .../spanner/SpannerBatchUpdateException.java | 11 ++++++-- .../spanner/SpannerExceptionFactory.java | 6 +++- ...sactionMutationLimitExceededException.java | 8 +++++- .../cloud/spanner/TransactionRunnerImpl.java | 4 +-- ...etryDmlAsPartitionedDmlMockServerTest.java | 28 +++++++++++++++++++ 5 files changed, 50 insertions(+), 7 deletions(-) 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 0e51c5f91f3..b5d0d2a6871 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 @@ -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; } 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 6476b94b144..3e88f06543c 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 @@ -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 */ 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 1b63861bcd1..f04af2cae93 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 @@ -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 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 038fb4b52eb..fefbb383604 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 @@ -1070,7 +1070,7 @@ public long[] batchUpdate(Iterable 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 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(), diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RetryDmlAsPartitionedDmlMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RetryDmlAsPartitionedDmlMockServerTest.java index 610a1a99cbe..c1af1e3f9dc 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RetryDmlAsPartitionedDmlMockServerTest.java +++ b/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)); + } }