diff --git a/README.md b/README.md index fb4045604ac..e7ff525adbe 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ If you are using Maven with [BOM][libraries-bom], add this to your pom.xml file: com.google.cloud libraries-bom - 26.53.0 + 26.54.0 pom import diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index ea21266e19d..145ad67f827 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -698,6 +698,9 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder( if (!isReadOnly()) { builder.setSeqno(getSeqNo()); } + if (options.hasLastStatement()) { + builder.setLastStatement(options.isLastStatement()); + } builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions())); builder.setRequestOptions(buildRequestOptions(options)); return builder; @@ -743,6 +746,9 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder( if (selector != null) { builder.setTransaction(selector); } + if (options.hasLastStatement()) { + builder.setLastStatements(options.isLastStatement()); + } builder.setSeqno(getSeqNo()); builder.setRequestOptions(buildRequestOptions(options)); return builder; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java index c062e89ec2b..c8c588f813a 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java @@ -108,6 +108,9 @@ public interface ReadQueryUpdateTransactionOption /** Marker interface to mark options applicable to Update and Write operations */ public interface UpdateTransactionOption extends UpdateOption, TransactionOption {} + /** Marker interface for options that can be used with both executeQuery and executeUpdate. */ + public interface QueryUpdateOption extends QueryOption, UpdateOption {} + /** * Marker interface to mark options applicable to Create, Update and Delete operations in admin * API. @@ -236,6 +239,20 @@ public static DataBoostQueryOption dataBoostEnabled(Boolean dataBoostEnabled) { return new DataBoostQueryOption(dataBoostEnabled); } + /** + * If set to true, this option marks the end of the transaction. The transaction should be + * committed or aborted after this statement executes, and attempts to execute any other requests + * against this transaction (including reads and queries) will be rejected. Mixing mutations with + * statements that are marked as the last statement is not allowed. + * + *

For DML statements, setting this option may cause some error reporting to be deferred until + * commit time (e.g. validation of unique constraints). Given this, successful execution of a DML + * statement should not be assumed until the transaction commits. + */ + public static QueryUpdateOption lastStatement() { + return new LastStatementUpdateOption(); + } + /** * Specifying this will cause the list operation to start fetching the record from this onwards. */ @@ -494,6 +511,7 @@ void appendToOptions(Options options) { private DecodeMode decodeMode; private RpcOrderBy orderBy; private RpcLockHint lockHint; + private Boolean lastStatement; // Construction is via factory methods below. private Options() {} @@ -630,6 +648,14 @@ OrderBy orderBy() { return orderBy == null ? null : orderBy.proto; } + boolean hasLastStatement() { + return lastStatement != null; + } + + Boolean isLastStatement() { + return lastStatement; + } + boolean hasLockHint() { return lockHint != null; } @@ -694,6 +720,9 @@ public String toString() { if (orderBy != null) { b.append("orderBy: ").append(orderBy).append(' '); } + if (lastStatement != null) { + b.append("lastStatement: ").append(lastStatement).append(' '); + } if (lockHint != null) { b.append("lockHint: ").append(lockHint).append(' '); } @@ -737,6 +766,7 @@ public boolean equals(Object o) { && Objects.equals(dataBoostEnabled(), that.dataBoostEnabled()) && Objects.equals(directedReadOptions(), that.directedReadOptions()) && Objects.equals(orderBy(), that.orderBy()) + && Objects.equals(isLastStatement(), that.isLastStatement()) && Objects.equals(lockHint(), that.lockHint()); } @@ -797,6 +827,9 @@ public int hashCode() { if (orderBy != null) { result = 31 * result + orderBy.hashCode(); } + if (lastStatement != null) { + result = 31 * result + lastStatement.hashCode(); + } if (lockHint != null) { result = 31 * result + lockHint.hashCode(); } @@ -965,4 +998,24 @@ public boolean equals(Object o) { return Objects.equals(filter, ((FilterOption) o).filter); } } + + static final class LastStatementUpdateOption extends InternalOption implements QueryUpdateOption { + + LastStatementUpdateOption() {} + + @Override + void appendToOptions(Options options) { + options.lastStatement = true; + } + + @Override + public int hashCode() { + return LastStatementUpdateOption.class.hashCode(); + } + + @Override + public boolean equals(Object o) { + return o instanceof LastStatementUpdateOption; + } + } } 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 2dd70ce108e..a3f174cda60 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 @@ -265,6 +265,9 @@ static ErrorDetails extractErrorDetails(Throwable cause) { if (cause instanceof ApiException) { return ((ApiException) cause).getErrorDetails(); } + if (cause instanceof SpannerException) { + return ((SpannerException) cause).getErrorDetails(); + } prevCause = cause; cause = cause.getCause(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java index 3c533cb9a7a..a827f82ba36 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java @@ -34,6 +34,7 @@ import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options; import com.google.cloud.spanner.Options.QueryOption; +import com.google.cloud.spanner.Options.QueryUpdateOption; import com.google.cloud.spanner.Options.UpdateOption; import com.google.cloud.spanner.PartitionOptions; import com.google.cloud.spanner.ReadOnlyTransaction; @@ -298,7 +299,8 @@ private ApiFuture executeDmlReturningAsync( writeTransaction.run( transaction -> DirectExecuteResultSet.ofResultSet( - transaction.executeQuery(update.getStatement(), options))); + transaction.executeQuery( + update.getStatement(), appendLastStatement(options)))); state = UnitOfWorkState.COMMITTED; return resultSet; } catch (Throwable t) { @@ -554,11 +556,15 @@ private ApiFuture> executeTransactionalUpdateAsync( transaction -> { if (analyzeMode == AnalyzeMode.NONE) { return Tuple.of( - transaction.executeUpdate(update.getStatement(), options), null); + transaction.executeUpdate( + update.getStatement(), appendLastStatement(options)), + null); } ResultSet resultSet = transaction.analyzeUpdateStatement( - update.getStatement(), analyzeMode.getQueryAnalyzeMode(), options); + update.getStatement(), + analyzeMode.getQueryAnalyzeMode(), + appendLastStatement(options)); return Tuple.of(null, resultSet); }); state = UnitOfWorkState.COMMITTED; @@ -582,6 +588,29 @@ private ApiFuture> executeTransactionalUpdateAsync( return transactionalResult; } + private static final QueryUpdateOption[] LAST_STATEMENT_OPTIONS = + new QueryUpdateOption[] {Options.lastStatement()}; + + private static UpdateOption[] appendLastStatement(UpdateOption[] options) { + if (options.length == 0) { + return LAST_STATEMENT_OPTIONS; + } + UpdateOption[] result = new UpdateOption[options.length + 1]; + System.arraycopy(options, 0, result, 0, options.length); + result[result.length - 1] = LAST_STATEMENT_OPTIONS[0]; + return result; + } + + private static QueryOption[] appendLastStatement(QueryOption[] options) { + if (options.length == 0) { + return LAST_STATEMENT_OPTIONS; + } + QueryOption[] result = new QueryOption[options.length + 1]; + System.arraycopy(options, 0, result, 0, options.length); + result[result.length - 1] = LAST_STATEMENT_OPTIONS[0]; + return result; + } + /** * Adds a callback to the given future that retries the update statement using Partitioned DML if * the original statement fails with a {@link TransactionMutationLimitExceededException}. @@ -719,7 +748,8 @@ private ApiFuture executeTransactionalBatchUpdateAsync( try { long[] res = transaction.batchUpdate( - Iterables.transform(updates, ParsedStatement::getStatement), options); + Iterables.transform(updates, ParsedStatement::getStatement), + appendLastStatement(options)); state = UnitOfWorkState.COMMITTED; return res; } catch (Throwable t) { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java index 8b53bd7efff..eea6658d26d 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -266,6 +267,42 @@ public void testGetExecuteBatchDmlRequestBuilderWithPriority() { assertEquals(Priority.PRIORITY_LOW, request.getRequestOptions().getPriority()); } + @Test + public void testExecuteSqlLastStatement() { + assertFalse( + context + .getExecuteSqlRequestBuilder( + Statement.of("insert into test (id) values (1)"), + QueryMode.NORMAL, + Options.fromUpdateOptions(), + false) + .getLastStatement()); + assertTrue( + context + .getExecuteSqlRequestBuilder( + Statement.of("insert into test (id) values (1)"), + QueryMode.NORMAL, + Options.fromUpdateOptions(Options.lastStatement()), + false) + .getLastStatement()); + } + + @Test + public void testExecuteBatchDmlLastStatement() { + assertFalse( + context + .getExecuteBatchDmlRequestBuilder( + Collections.singleton(Statement.of("insert into test (id) values (1)")), + Options.fromUpdateOptions()) + .getLastStatements()); + assertTrue( + context + .getExecuteBatchDmlRequestBuilder( + Collections.singleton(Statement.of("insert into test (id) values (1)")), + Options.fromUpdateOptions(Options.lastStatement())) + .getLastStatements()); + } + public void executeSqlRequestBuilderWithRequestOptions() { ExecuteSqlRequest request = context diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java index f391088589f..17c25558f3b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java @@ -789,4 +789,22 @@ public void updateOptionsExcludeTxnFromChangeStreams() { assertNull(option3.withExcludeTxnFromChangeStreams()); assertThat(option3.toString()).doesNotContain("withExcludeTxnFromChangeStreams: true"); } + + @Test + public void testLastStatement() { + Options option1 = Options.fromUpdateOptions(Options.lastStatement()); + Options option2 = Options.fromUpdateOptions(Options.lastStatement()); + Options option3 = Options.fromUpdateOptions(); + + assertEquals(option1, option2); + assertEquals(option1.hashCode(), option2.hashCode()); + assertNotEquals(option1, option3); + assertNotEquals(option1.hashCode(), option3.hashCode()); + + assertTrue(option1.isLastStatement()); + assertThat(option1.toString()).contains("lastStatement: true"); + + assertNull(option3.isLastStatement()); + assertThat(option3.toString()).doesNotContain("lastStatement: true"); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AutoCommitMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AutoCommitMockServerTest.java new file mode 100644 index 00000000000..c48c6703353 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AutoCommitMockServerTest.java @@ -0,0 +1,166 @@ +/* + * Copyright 2024 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.connection; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.google.cloud.spanner.ResultSet; +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 org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AutoCommitMockServerTest extends AbstractMockServerTest { + + @Test + public void testQuery() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + //noinspection EmptyTryBlock + try (ResultSet ignore = connection.executeQuery(SELECT1_STATEMENT)) {} + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertTrue(request.getTransaction().hasSingleUse()); + assertTrue(request.getTransaction().getSingleUse().hasReadOnly()); + assertFalse(request.getLastStatement()); + } + + @Test + public void testDml() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + connection.executeUpdate(INSERT_STATEMENT); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertTrue(request.getTransaction().hasBegin()); + assertTrue(request.getTransaction().getBegin().hasReadWrite()); + assertTrue(request.getLastStatement()); + assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class)); + } + + @Test + public void testDmlReturning() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + //noinspection EmptyTryBlock + try (ResultSet ignore = connection.executeQuery(INSERT_RETURNING_STATEMENT)) {} + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertTrue(request.getTransaction().hasBegin()); + assertTrue(request.getTransaction().getBegin().hasReadWrite()); + assertTrue(request.getLastStatement()); + assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class)); + } + + @Test + public void testBatchDml() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + connection.startBatchDml(); + connection.executeUpdate(INSERT_STATEMENT); + connection.executeUpdate(INSERT_STATEMENT); + connection.runBatch(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + ExecuteBatchDmlRequest request = + mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0); + assertTrue(request.getTransaction().hasBegin()); + assertTrue(request.getTransaction().getBegin().hasReadWrite()); + assertTrue(request.getLastStatements()); + assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class)); + } + + @Test + public void testPartitionedDml() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + connection.setAutocommitDmlMode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC); + connection.executeUpdate(INSERT_STATEMENT); + } + assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); + BeginTransactionRequest beginRequest = + mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0); + assertTrue(beginRequest.getOptions().hasPartitionedDml()); + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertTrue(request.getTransaction().hasId()); + assertFalse(request.getLastStatement()); + assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class)); + } + + @Test + public void testDmlAborted() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + mockSpanner.abortNextTransaction(); + connection.executeUpdate(INSERT_STATEMENT); + } + assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + for (ExecuteSqlRequest request : mockSpanner.getRequestsOfType(ExecuteSqlRequest.class)) { + assertTrue(request.getTransaction().hasBegin()); + assertTrue(request.getTransaction().getBegin().hasReadWrite()); + assertTrue(request.getLastStatement()); + } + assertEquals(2, mockSpanner.countRequestsOfType(CommitRequest.class)); + } + + @Test + public void testDmlReturningAborted() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + mockSpanner.abortNextTransaction(); + //noinspection EmptyTryBlock + try (ResultSet ignore = connection.executeQuery(INSERT_RETURNING_STATEMENT)) {} + } + assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + for (ExecuteSqlRequest request : mockSpanner.getRequestsOfType(ExecuteSqlRequest.class)) { + assertTrue(request.getTransaction().hasBegin()); + assertTrue(request.getTransaction().getBegin().hasReadWrite()); + assertTrue(request.getLastStatement()); + } + assertEquals(2, mockSpanner.countRequestsOfType(CommitRequest.class)); + } + + @Test + public void testBatchDmlAborted() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + mockSpanner.abortNextTransaction(); + connection.startBatchDml(); + connection.executeUpdate(INSERT_STATEMENT); + connection.executeUpdate(INSERT_STATEMENT); + connection.runBatch(); + } + assertEquals(2, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + for (ExecuteBatchDmlRequest request : + mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class)) { + assertTrue(request.getTransaction().hasBegin()); + assertTrue(request.getTransaction().getBegin().hasReadWrite()); + assertTrue(request.getLastStatements()); + } + assertEquals(2, mockSpanner.countRequestsOfType(CommitRequest.class)); + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AutocommitDmlModeTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AutocommitDmlModeTest.java index a66d14a8b76..bb845746e13 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AutocommitDmlModeTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AutocommitDmlModeTest.java @@ -18,6 +18,9 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -28,6 +31,7 @@ import com.google.cloud.spanner.BatchClient; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.Options; import com.google.cloud.spanner.Spanner; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.TransactionContext; @@ -82,12 +86,12 @@ public void testAutocommitDmlModeTransactional() { .setCredentials(NoCredentials.getInstance()) .setUri(URI) .build())) { - assertThat(connection.isAutocommit(), is(true)); - assertThat(connection.isReadOnly(), is(false)); - assertThat(connection.getAutocommitDmlMode(), is(AutocommitDmlMode.TRANSACTIONAL)); + assertTrue(connection.isAutocommit()); + assertFalse(connection.isReadOnly()); + assertEquals(AutocommitDmlMode.TRANSACTIONAL, connection.getAutocommitDmlMode()); connection.execute(Statement.of(UPDATE)); - verify(txContext).executeUpdate(Statement.of(UPDATE)); + verify(txContext).executeUpdate(Statement.of(UPDATE), Options.lastStatement()); verify(dbClient, never()).executePartitionedUpdate(Statement.of(UPDATE)); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java index a81005bbb45..d4b7c035658 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java @@ -354,7 +354,8 @@ public TransactionRunner answer(InvocationOnMock invocation) { public T run(TransactionCallable callable) { commitResponse = new CommitResponse(Timestamp.ofTimeSecondsAndNanos(1, 1)); TransactionContext transaction = mock(TransactionContext.class); - when(transaction.executeUpdate(Statement.of(UPDATE))).thenReturn(1L); + when(transaction.executeUpdate(Statement.of(UPDATE), Options.lastStatement())) + .thenReturn(1L); try { return callable.run(transaction); } catch (Exception e) { 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 022c9a92f1f..610a1a99cbe 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 @@ -83,6 +83,8 @@ public void testTransactionMutationLimitExceeded_isNotRetriedByDefault() { assertEquals(0, exception.getSuppressed().length); } assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertTrue(request.getLastStatement()); assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class)); } @@ -108,6 +110,7 @@ public void testTransactionMutationLimitExceeded_canBeRetriedAsPDML() { ExecuteSqlRequest transactionalRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); assertTrue(transactionalRequest.getTransaction().getBegin().hasReadWrite()); + assertTrue(transactionalRequest.getLastStatement()); // Partitioned DML uses an explicit BeginTransaction RPC. assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); @@ -117,6 +120,7 @@ public void testTransactionMutationLimitExceeded_canBeRetriedAsPDML() { ExecuteSqlRequest partitionedDmlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(1); assertTrue(partitionedDmlRequest.getTransaction().hasId()); + assertFalse(partitionedDmlRequest.getLastStatement()); // Partitioned DML transactions are not committed. assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class)); @@ -163,6 +167,7 @@ public void testTransactionMutationLimitExceeded_retryAsPDMLFails() { ExecuteSqlRequest transactionalRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); assertTrue(transactionalRequest.getTransaction().getBegin().hasReadWrite()); + assertTrue(transactionalRequest.getLastStatement()); // Partitioned DML uses an explicit BeginTransaction RPC. assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); @@ -172,6 +177,7 @@ public void testTransactionMutationLimitExceeded_retryAsPDMLFails() { ExecuteSqlRequest partitionedDmlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(1); assertTrue(partitionedDmlRequest.getTransaction().hasId()); + assertFalse(partitionedDmlRequest.getLastStatement()); // Partitioned DML transactions are not committed. assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class)); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java index 6edf46b5623..0e2a322023d 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java @@ -395,6 +395,8 @@ private SingleUseTransaction createSubject( final TransactionContext txContext = mock(TransactionContext.class); when(txContext.executeUpdate(Statement.of(VALID_UPDATE))).thenReturn(VALID_UPDATE_COUNT); + when(txContext.executeUpdate(Statement.of(VALID_UPDATE), Options.lastStatement())) + .thenReturn(VALID_UPDATE_COUNT); when(txContext.executeUpdate(Statement.of(SLOW_UPDATE))) .thenAnswer( invocation -> { @@ -404,6 +406,9 @@ private SingleUseTransaction createSubject( when(txContext.executeUpdate(Statement.of(INVALID_UPDATE))) .thenThrow( SpannerExceptionFactory.newSpannerException(ErrorCode.UNKNOWN, "invalid update")); + when(txContext.executeUpdate(Statement.of(INVALID_UPDATE), Options.lastStatement())) + .thenThrow( + SpannerExceptionFactory.newSpannerException(ErrorCode.UNKNOWN, "invalid update")); SimpleTransactionManager txManager = new SimpleTransactionManager(txContext, commitBehavior); when(dbClient.transactionManager()).thenReturn(txManager);