Skip to content

Commit bf01bb2

Browse files
committed
feat: add update count verification
1 parent ce47b75 commit bf01bb2

File tree

6 files changed

+248
-14
lines changed

6 files changed

+248
-14
lines changed
Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner;
18+
19+
import com.google.cloud.spanner.connection.Connection;
20+
import java.util.Arrays;
21+
import java.util.stream.Collectors;
22+
23+
/**
24+
* Exception thrown by a {@link Connection} when an automatic DML batch detects that one or more of
25+
* the update counts that it returned during the buffering of DML statements does not match with the
26+
* actual update counts that were returned after execution.
27+
*/
28+
public class DmlBatchUpdateCountVerificationFailedException extends AbortedException {
29+
private static final long serialVersionUID = 1L;
30+
31+
/** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
32+
DmlBatchUpdateCountVerificationFailedException(
33+
DoNotConstructDirectly token, long[] expected, long[] actual) {
34+
super(
35+
token,
36+
String.format(
37+
"Actual update counts that were returned during execution do not match the previously returned update counts.\n"
38+
+ "Expected: %s\n"
39+
+ "Actual: %s\n"
40+
+ "Set auto_batch_dml_update_count_verification to false to skip this verification.",
41+
Arrays.stream(expected).mapToObj(Long::toString).collect(Collectors.joining()),
42+
Arrays.stream(actual).mapToObj(Long::toString).collect(Collectors.joining())),
43+
/* cause = */ null);
44+
}
45+
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,12 @@ public static SpannerBatchUpdateException newSpannerBatchUpdateException(
116116
return new SpannerBatchUpdateException(token, code, message, updateCounts);
117117
}
118118

119+
public static DmlBatchUpdateCountVerificationFailedException
120+
newDmlBatchUpdateCountVerificationFailedException(long[] expected, long[] actual) {
121+
return new DmlBatchUpdateCountVerificationFailedException(
122+
DoNotConstructDirectly.ALLOWED, expected, actual);
123+
}
124+
119125
/**
120126
* Constructs a specific aborted exception that should only be thrown by a connection after an
121127
* internal retry aborted due to concurrent modifications.

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1169,6 +1169,14 @@ default StatementResult execute(Statement statement, Set<ResultType> allowedResu
11691169

11701170
boolean isAutoBatchDml();
11711171

1172+
void setAutoBatchDmlUpdateCount(long updateCount);
1173+
1174+
long getAutoBatchDmlUpdateCount();
1175+
1176+
void setAutoBatchDmlUpdateCountVerification(boolean verification);
1177+
1178+
boolean isAutoBatchDmlUpdateCountVerification();
1179+
11721180
/**
11731181
* Enable data boost for partitioned queries. See also {@link #partitionQuery(Statement,
11741182
* PartitionOptions, QueryOption...)}

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1345,6 +1345,32 @@ public boolean isAutoBatchDml() {
13451345
return Boolean.parseBoolean(System.getProperty("spanner.auto_batch_dml"));
13461346
}
13471347

1348+
@Override
1349+
public void setAutoBatchDmlUpdateCount(long updateCount) {
1350+
// TODO: Replace with a connection variable
1351+
System.setProperty("spanner.auto_batch_dml_update_count", String.valueOf(updateCount));
1352+
}
1353+
1354+
@Override
1355+
public long getAutoBatchDmlUpdateCount() {
1356+
// TODO: Replace with a connection variable
1357+
return Long.parseLong(System.getProperty("spanner.auto_batch_dml_update_count", "1"));
1358+
}
1359+
1360+
@Override
1361+
public void setAutoBatchDmlUpdateCountVerification(boolean verification) {
1362+
// TODO: Replace with a connection variable
1363+
System.setProperty(
1364+
"spanner.auto_batch_dml_update_count_verification", String.valueOf(verification));
1365+
}
1366+
1367+
@Override
1368+
public boolean isAutoBatchDmlUpdateCountVerification() {
1369+
// TODO: Replace with a connection variable
1370+
return Boolean.parseBoolean(
1371+
System.getProperty("spanner.auto_batch_dml_update_count_verification", "true"));
1372+
}
1373+
13481374
@Override
13491375
public void setDataBoostEnabled(boolean dataBoostEnabled) {
13501376
this.dataBoostEnabled = dataBoostEnabled;
@@ -2049,6 +2075,9 @@ UnitOfWork createNewUnitOfWork(
20492075
pushCurrentUnitOfWorkToTransactionStack();
20502076
return DmlBatch.newBuilder()
20512077
.setAutoBatch(autoBatchDml)
2078+
.setAutoBatchUpdateCountSupplier(this::getAutoBatchDmlUpdateCount)
2079+
.setAutoBatchUpdateCountVerificationSupplier(
2080+
this::isAutoBatchDmlUpdateCountVerification)
20522081
.setTransaction(currentUnitOfWork)
20532082
.setStatementTimeout(statementTimeout)
20542083
.withStatementExecutor(statementExecutor)

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java

Lines changed: 55 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,12 +32,14 @@
3232
import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
3333
import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
3434
import com.google.common.base.Preconditions;
35+
import com.google.common.base.Suppliers;
3536
import com.google.common.collect.Iterables;
3637
import com.google.common.util.concurrent.MoreExecutors;
3738
import io.opentelemetry.context.Scope;
3839
import java.util.ArrayList;
3940
import java.util.Arrays;
4041
import java.util.List;
42+
import java.util.function.Supplier;
4143

4244
/**
4345
* {@link UnitOfWork} that is used when a DML batch is started. These batches only accept DML
@@ -46,13 +48,18 @@
4648
*/
4749
class DmlBatch extends AbstractBaseUnitOfWork {
4850
private final boolean autoBatch;
51+
private final Supplier<Long> autoBatchUpdateCountSupplier;
52+
private final Supplier<Boolean> verifyUpdateCountsSupplier;
4953
private final UnitOfWork transaction;
5054
private final String statementTag;
5155
private final List<ParsedStatement> statements = new ArrayList<>();
56+
private long[] updateCounts = new long[0];
5257
private UnitOfWorkState state = UnitOfWorkState.STARTED;
5358

5459
static class Builder extends AbstractBaseUnitOfWork.Builder<Builder, DmlBatch> {
5560
private boolean autoBatch;
61+
private Supplier<Long> autoBatchUpdateCountSupplier = Suppliers.ofInstance(1L);
62+
private Supplier<Boolean> verifyUpdateCountsSupplier = Suppliers.ofInstance(Boolean.FALSE);
5663
private UnitOfWork transaction;
5764
private String statementTag;
5865

@@ -63,6 +70,16 @@ Builder setAutoBatch(boolean autoBatch) {
6370
return this;
6471
}
6572

73+
Builder setAutoBatchUpdateCountSupplier(Supplier<Long> updateCountSupplier) {
74+
this.autoBatchUpdateCountSupplier = Preconditions.checkNotNull(updateCountSupplier);
75+
return this;
76+
}
77+
78+
Builder setAutoBatchUpdateCountVerificationSupplier(Supplier<Boolean> verificationSupplier) {
79+
this.verifyUpdateCountsSupplier = verificationSupplier;
80+
return this;
81+
}
82+
6683
Builder setTransaction(UnitOfWork transaction) {
6784
Preconditions.checkNotNull(transaction);
6885
this.transaction = transaction;
@@ -88,7 +105,8 @@ static Builder newBuilder() {
88105
private DmlBatch(Builder builder) {
89106
super(builder);
90107
this.autoBatch = builder.autoBatch;
91-
;
108+
this.autoBatchUpdateCountSupplier = builder.autoBatchUpdateCountSupplier;
109+
this.verifyUpdateCountsSupplier = builder.verifyUpdateCountsSupplier;
92110
this.transaction = Preconditions.checkNotNull(builder.transaction);
93111
this.statementTag = builder.statementTag;
94112
}
@@ -172,10 +190,9 @@ public ApiFuture<Void> executeDdlAsync(CallType callType, ParsedStatement ddl) {
172190
}
173191

174192
long getUpdateCount() {
175-
// TODO: Make configurable.
176-
// Auto-batching returns update count 1, as this is what ORMs normally expect.
193+
// Auto-batching returns update count 1 by default, as this is what ORMs normally expect.
177194
// Standard batches return -1 by default, to indicate that the update count is unknown.
178-
return isAutoBatch() ? 1L : -1L;
195+
return isAutoBatch() ? autoBatchUpdateCountSupplier.get() : -1L;
179196
}
180197

181198
@Override
@@ -189,8 +206,11 @@ public ApiFuture<Long> executeUpdateAsync(
189206
"Only DML statements are allowed. \""
190207
+ update.getSqlWithoutComments()
191208
+ "\" is not a DML-statement.");
192-
statements.add(update);
193-
return ApiFutures.immediateFuture(getUpdateCount());
209+
long updateCount = getUpdateCount();
210+
this.statements.add(update);
211+
this.updateCounts = Arrays.copyOf(this.updateCounts, this.updateCounts.length + 1);
212+
this.updateCounts[this.updateCounts.length - 1] = updateCount;
213+
return ApiFutures.immediateFuture(updateCount);
194214
}
195215

196216
@Override
@@ -216,10 +236,19 @@ public ApiFuture<long[]> executeBatchUpdateAsync(
216236
+ update.getSqlWithoutComments()
217237
+ "\" is not a DML-statement.");
218238
}
219-
Iterables.addAll(statements, updates);
220-
long[] result = new long[Iterables.size(updates)];
221-
Arrays.fill(result, getUpdateCount());
222-
return ApiFutures.immediateFuture(result);
239+
long[] updateCountArray = new long[Iterables.size(updates)];
240+
Arrays.fill(updateCountArray, getUpdateCount());
241+
Iterables.addAll(this.statements, updates);
242+
this.updateCounts =
243+
Arrays.copyOf(this.updateCounts, this.updateCounts.length + updateCountArray.length);
244+
System.arraycopy(
245+
updateCountArray,
246+
0,
247+
this.updateCounts,
248+
this.updateCounts.length - updateCountArray.length,
249+
updateCountArray.length);
250+
251+
return ApiFutures.immediateFuture(updateCountArray);
223252
}
224253

225254
@Override
@@ -273,7 +302,13 @@ public void onFailure(Throwable t) {
273302
@Override
274303
public void onSuccess(long[] result) {
275304
state = UnitOfWorkState.RAN;
276-
res.set(result);
305+
if (!verifyUpdateCounts(result)) {
306+
res.setException(
307+
SpannerExceptionFactory.newDmlBatchUpdateCountVerificationFailedException(
308+
DmlBatch.this.updateCounts, result));
309+
} else {
310+
res.set(result);
311+
}
277312
}
278313
},
279314
MoreExecutors.directExecutor());
@@ -282,6 +317,15 @@ public void onSuccess(long[] result) {
282317
}
283318
}
284319

320+
private boolean verifyUpdateCounts(long[] actualUpdateCounts) {
321+
if (!this.autoBatch || !this.verifyUpdateCountsSupplier.get()) {
322+
// We only need to do an actual verification if the batch was an auto-batch and verification
323+
// is enabled.
324+
return true;
325+
}
326+
return Arrays.equals(this.updateCounts, actualUpdateCounts);
327+
}
328+
285329
@Override
286330
public void abortBatch() {
287331
ConnectionPreconditions.checkState(

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

Lines changed: 105 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 Google LLC
2+
* Copyright 2024 Google LLC
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@
2626
import com.google.cloud.spanner.MockSpannerServiceImpl;
2727
import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
2828
import com.google.cloud.spanner.ResultSet;
29+
import com.google.cloud.spanner.SpannerBatchUpdateException;
2930
import com.google.cloud.spanner.SpannerException;
3031
import com.google.cloud.spanner.Statement;
3132
import com.google.cloud.spanner.connection.StatementResult.ResultType;
@@ -550,9 +551,12 @@ public void testDmlWithErrorInBatch() {
550551

551552
// This SELECT statement flushes the batch and is the one that gets the exception, even
552553
// though the statement itself is valid.
553-
SpannerException exception =
554-
assertThrows(SpannerException.class, () -> connection.executeQuery(SELECT1_STATEMENT));
554+
SpannerBatchUpdateException exception =
555+
assertThrows(
556+
SpannerBatchUpdateException.class, () -> connection.executeQuery(SELECT1_STATEMENT));
555557
assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode());
558+
// The batch exception contains the update count for the successful DML statements.
559+
assertArrayEquals(new long[] {1L}, exception.getUpdateCounts());
556560

557561
connection.commit();
558562
}
@@ -566,4 +570,102 @@ public void testDmlWithErrorInBatch() {
566570
assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
567571
assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
568572
}
573+
574+
@Test
575+
public void testUpdateCount() {
576+
try (Connection connection = createConnection()) {
577+
578+
// Setting a different update count is reflected in the update count that is returned by
579+
// an auto-batch.
580+
try {
581+
// Disable update count verification to prevent errors.
582+
connection.setAutoBatchDmlUpdateCountVerification(false);
583+
584+
connection.setAutoBatchDmlUpdateCount(2L);
585+
assertEquals(2L, connection.executeUpdate(INSERT_STATEMENT));
586+
// The update count can be modified during the batch.
587+
connection.setAutoBatchDmlUpdateCount(3L);
588+
assertEquals(3L, connection.executeUpdate(INSERT_STATEMENT));
589+
590+
connection.commit();
591+
592+
// The auto-batch update count setting is not used for explicit batches.
593+
connection.startBatchDml();
594+
assertEquals(-1L, connection.executeUpdate(INSERT_STATEMENT));
595+
connection.runBatch();
596+
connection.commit();
597+
} finally {
598+
// TODO: Remove once a normal connection variable is used for this.
599+
System.clearProperty("spanner.auto_batch_dml_update_count");
600+
System.clearProperty("spanner.auto_batch_dml_update_count_verification");
601+
}
602+
}
603+
604+
assertEquals(2, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
605+
assertEquals(2, mockSpanner.countRequestsOfType(CommitRequest.class));
606+
}
607+
608+
@Test
609+
public void testUpdateCountVerification_failsIfDifferent() {
610+
try (Connection connection = createConnection()) {
611+
try {
612+
assertEquals(1L, connection.executeUpdate(INSERT_STATEMENT));
613+
// Set a different (expected) update count. This will cause the batch to fail, as the
614+
// actual update count will be 1.
615+
connection.setAutoBatchDmlUpdateCount(3L);
616+
assertEquals(3L, connection.executeUpdate(INSERT_STATEMENT));
617+
618+
assertThrows(SpannerException.class, connection::commit);
619+
} finally {
620+
// TODO: Remove once a normal connection variable is used for this.
621+
System.clearProperty("spanner.auto_batch_dml_update_count");
622+
}
623+
}
624+
625+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
626+
assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
627+
}
628+
629+
@Test
630+
public void testUpdateCountVerification_succeedsIfSame() {
631+
Statement statement1 = Statement.of("insert into foo (id, value) values (1, 'One')");
632+
Statement statement2 = Statement.of("insert into foo (id, value) values (2, 'Two')");
633+
Statement statement3 = Statement.of("insert into foo (id, value) values (3, 'Three')");
634+
Statement statement4 = Statement.of("insert into foo (id, value) values (4, 'Four')");
635+
Statement statement5 = Statement.of("insert into foo (id, value) values (5, 'Five')");
636+
mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement1, 1L));
637+
mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement2, 2L));
638+
mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement3, 3L));
639+
mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement4, 3L));
640+
mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement5, 4L));
641+
642+
try (Connection connection = createConnection()) {
643+
try {
644+
connection.setAutoBatchDmlUpdateCount(1L);
645+
assertEquals(1L, connection.executeUpdate(statement1));
646+
647+
connection.setAutoBatchDmlUpdateCount(2L);
648+
assertEquals(2L, connection.executeUpdate(statement2));
649+
650+
connection.setAutoBatchDmlUpdateCount(3L);
651+
assertArrayEquals(
652+
new long[] {3L, 3L},
653+
connection.executeBatchUpdate(ImmutableList.of(statement3, statement4)));
654+
655+
connection.setAutoBatchDmlUpdateCount(4L);
656+
assertEquals(4L, connection.executeUpdate(statement5));
657+
658+
connection.commit();
659+
} finally {
660+
// TODO: Remove once a normal connection variable is used for this.
661+
System.clearProperty("spanner.auto_batch_dml_update_count");
662+
}
663+
}
664+
665+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
666+
ExecuteBatchDmlRequest request =
667+
mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
668+
assertEquals(5, request.getStatementsCount());
669+
assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
670+
}
569671
}

0 commit comments

Comments
 (0)