feat: add update count verification

Author: olavloite

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

@@ -0,0 +1,45 @@
+/*
+ * 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;
+
+import com.google.cloud.spanner.connection.Connection;
+import java.util.Arrays;
+import java.util.stream.Collectors;
+
+/**
+ * Exception thrown by a {@link Connection} when an automatic DML batch detects that one or more of
+ * the update counts that it returned during the buffering of DML statements does not match with the
+ * actual update counts that were returned after execution.
+ */
+public class DmlBatchUpdateCountVerificationFailedException extends AbortedException {
+  private static final long serialVersionUID = 1L;
+
+  /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
+  DmlBatchUpdateCountVerificationFailedException(
+      DoNotConstructDirectly token, long[] expected, long[] actual) {
+    super(
+        token,
+        String.format(
+            "Actual update counts that were returned during execution do not match the previously returned update counts.\n"
+                + "Expected: %s\n"
+                + "Actual: %s\n"
+                + "Set auto_batch_dml_update_count_verification to false to skip this verification.",
+            Arrays.stream(expected).mapToObj(Long::toString).collect(Collectors.joining()),
+            Arrays.stream(actual).mapToObj(Long::toString).collect(Collectors.joining())),
+        /* cause = */ null);
+  }
+}

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

@@ -116,6 +116,12 @@ public static SpannerBatchUpdateException newSpannerBatchUpdateException(
     return new SpannerBatchUpdateException(token, code, message, updateCounts);
   }
 
+  public static DmlBatchUpdateCountVerificationFailedException
+      newDmlBatchUpdateCountVerificationFailedException(long[] expected, long[] actual) {
+    return new DmlBatchUpdateCountVerificationFailedException(
+        DoNotConstructDirectly.ALLOWED, expected, actual);
+  }
+
   /**
    * Constructs a specific aborted exception that should only be thrown by a connection after an
    * internal retry aborted due to concurrent modifications.

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

@@ -1169,6 +1169,14 @@ default StatementResult execute(Statement statement, Set<ResultType> allowedResu
 
   boolean isAutoBatchDml();
 
+  void setAutoBatchDmlUpdateCount(long updateCount);
+
+  long getAutoBatchDmlUpdateCount();
+
+  void setAutoBatchDmlUpdateCountVerification(boolean verification);
+
+  boolean isAutoBatchDmlUpdateCountVerification();
+
   /**
    * Enable data boost for partitioned queries. See also {@link #partitionQuery(Statement,
    * PartitionOptions, QueryOption...)}

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

@@ -1345,6 +1345,32 @@ public boolean isAutoBatchDml() {
     return Boolean.parseBoolean(System.getProperty("spanner.auto_batch_dml"));
   }
 
+  @Override
+  public void setAutoBatchDmlUpdateCount(long updateCount) {
+    // TODO: Replace with a connection variable
+    System.setProperty("spanner.auto_batch_dml_update_count", String.valueOf(updateCount));
+  }
+
+  @Override
+  public long getAutoBatchDmlUpdateCount() {
+    // TODO: Replace with a connection variable
+    return Long.parseLong(System.getProperty("spanner.auto_batch_dml_update_count", "1"));
+  }
+
+  @Override
+  public void setAutoBatchDmlUpdateCountVerification(boolean verification) {
+    // TODO: Replace with a connection variable
+    System.setProperty(
+        "spanner.auto_batch_dml_update_count_verification", String.valueOf(verification));
+  }
+
+  @Override
+  public boolean isAutoBatchDmlUpdateCountVerification() {
+    // TODO: Replace with a connection variable
+    return Boolean.parseBoolean(
+        System.getProperty("spanner.auto_batch_dml_update_count_verification", "true"));
+  }
+
   @Override
   public void setDataBoostEnabled(boolean dataBoostEnabled) {
     this.dataBoostEnabled = dataBoostEnabled;
@@ -2049,6 +2075,9 @@ UnitOfWork createNewUnitOfWork(
           pushCurrentUnitOfWorkToTransactionStack();
           return DmlBatch.newBuilder()
               .setAutoBatch(autoBatchDml)
+              .setAutoBatchUpdateCountSupplier(this::getAutoBatchDmlUpdateCount)
+              .setAutoBatchUpdateCountVerificationSupplier(
+                  this::isAutoBatchDmlUpdateCountVerification)
               .setTransaction(currentUnitOfWork)
               .setStatementTimeout(statementTimeout)
               .withStatementExecutor(statementExecutor)

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

@@ -32,12 +32,14 @@
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Suppliers;
 import com.google.common.collect.Iterables;
 import com.google.common.util.concurrent.MoreExecutors;
 import io.opentelemetry.context.Scope;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Supplier;
 
 /**
  * {@link UnitOfWork} that is used when a DML batch is started. These batches only accept DML
@@ -46,13 +48,18 @@
  */
 class DmlBatch extends AbstractBaseUnitOfWork {
   private final boolean autoBatch;
+  private final Supplier<Long> autoBatchUpdateCountSupplier;
+  private final Supplier<Boolean> verifyUpdateCountsSupplier;
   private final UnitOfWork transaction;
   private final String statementTag;
   private final List<ParsedStatement> statements = new ArrayList<>();
+  private long[] updateCounts = new long[0];
   private UnitOfWorkState state = UnitOfWorkState.STARTED;
 
   static class Builder extends AbstractBaseUnitOfWork.Builder<Builder, DmlBatch> {
     private boolean autoBatch;
+    private Supplier<Long> autoBatchUpdateCountSupplier = Suppliers.ofInstance(1L);
+    private Supplier<Boolean> verifyUpdateCountsSupplier = Suppliers.ofInstance(Boolean.FALSE);
     private UnitOfWork transaction;
     private String statementTag;
 
@@ -63,6 +70,16 @@ Builder setAutoBatch(boolean autoBatch) {
       return this;
     }
 
+    Builder setAutoBatchUpdateCountSupplier(Supplier<Long> updateCountSupplier) {
+      this.autoBatchUpdateCountSupplier = Preconditions.checkNotNull(updateCountSupplier);
+      return this;
+    }
+
+    Builder setAutoBatchUpdateCountVerificationSupplier(Supplier<Boolean> verificationSupplier) {
+      this.verifyUpdateCountsSupplier = verificationSupplier;
+      return this;
+    }
+
     Builder setTransaction(UnitOfWork transaction) {
       Preconditions.checkNotNull(transaction);
       this.transaction = transaction;
@@ -88,7 +105,8 @@ static Builder newBuilder() {
   private DmlBatch(Builder builder) {
     super(builder);
     this.autoBatch = builder.autoBatch;
-    ;
+    this.autoBatchUpdateCountSupplier = builder.autoBatchUpdateCountSupplier;
+    this.verifyUpdateCountsSupplier = builder.verifyUpdateCountsSupplier;
     this.transaction = Preconditions.checkNotNull(builder.transaction);
     this.statementTag = builder.statementTag;
   }
@@ -172,10 +190,9 @@ public ApiFuture<Void> executeDdlAsync(CallType callType, ParsedStatement ddl) {
   }
 
   long getUpdateCount() {
-    // TODO: Make configurable.
-    // Auto-batching returns update count 1, as this is what ORMs normally expect.
+    // Auto-batching returns update count 1 by default, as this is what ORMs normally expect.
     // Standard batches return -1 by default, to indicate that the update count is unknown.
-    return isAutoBatch() ? 1L : -1L;
+    return isAutoBatch() ? autoBatchUpdateCountSupplier.get() : -1L;
   }
 
   @Override
@@ -189,8 +206,11 @@ public ApiFuture<Long> executeUpdateAsync(
         "Only DML statements are allowed. \""
             + update.getSqlWithoutComments()
             + "\" is not a DML-statement.");
-    statements.add(update);
-    return ApiFutures.immediateFuture(getUpdateCount());
+    long updateCount = getUpdateCount();
+    this.statements.add(update);
+    this.updateCounts = Arrays.copyOf(this.updateCounts, this.updateCounts.length + 1);
+    this.updateCounts[this.updateCounts.length - 1] = updateCount;
+    return ApiFutures.immediateFuture(updateCount);
   }
 
   @Override
@@ -216,10 +236,19 @@ public ApiFuture<long[]> executeBatchUpdateAsync(
               + update.getSqlWithoutComments()
               + "\" is not a DML-statement.");
     }
-    Iterables.addAll(statements, updates);
-    long[] result = new long[Iterables.size(updates)];
-    Arrays.fill(result, getUpdateCount());
-    return ApiFutures.immediateFuture(result);
+    long[] updateCountArray = new long[Iterables.size(updates)];
+    Arrays.fill(updateCountArray, getUpdateCount());
+    Iterables.addAll(this.statements, updates);
+    this.updateCounts =
+        Arrays.copyOf(this.updateCounts, this.updateCounts.length + updateCountArray.length);
+    System.arraycopy(
+        updateCountArray,
+        0,
+        this.updateCounts,
+        this.updateCounts.length - updateCountArray.length,
+        updateCountArray.length);
+
+    return ApiFutures.immediateFuture(updateCountArray);
   }
 
   @Override
@@ -273,7 +302,13 @@ public void onFailure(Throwable t) {
             @Override
             public void onSuccess(long[] result) {
               state = UnitOfWorkState.RAN;
-              res.set(result);
+              if (!verifyUpdateCounts(result)) {
+                res.setException(
+                    SpannerExceptionFactory.newDmlBatchUpdateCountVerificationFailedException(
+                        DmlBatch.this.updateCounts, result));
+              } else {
+                res.set(result);
+              }
             }
           },
           MoreExecutors.directExecutor());
@@ -282,6 +317,15 @@ public void onSuccess(long[] result) {
     }
   }
 
+  private boolean verifyUpdateCounts(long[] actualUpdateCounts) {
+    if (!this.autoBatch || !this.verifyUpdateCountsSupplier.get()) {
+      // We only need to do an actual verification if the batch was an auto-batch and verification
+      // is enabled.
+      return true;
+    }
+    return Arrays.equals(this.updateCounts, actualUpdateCounts);
+  }
+
   @Override
   public void abortBatch() {
     ConnectionPreconditions.checkState(

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 Google LLC
+ * 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.
@@ -26,6 +26,7 @@
 import com.google.cloud.spanner.MockSpannerServiceImpl;
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 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.connection.StatementResult.ResultType;
@@ -550,9 +551,12 @@ public void testDmlWithErrorInBatch() {
 
       // This SELECT statement flushes the batch and is the one that gets the exception, even
       // though the statement itself is valid.
-      SpannerException exception =
-          assertThrows(SpannerException.class, () -> connection.executeQuery(SELECT1_STATEMENT));
+      SpannerBatchUpdateException exception =
+          assertThrows(
+              SpannerBatchUpdateException.class, () -> connection.executeQuery(SELECT1_STATEMENT));
       assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode());
+      // The batch exception contains the update count for the successful DML statements.
+      assertArrayEquals(new long[] {1L}, exception.getUpdateCounts());
 
       connection.commit();
     }
@@ -566,4 +570,102 @@ public void testDmlWithErrorInBatch() {
     assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
     assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
   }
+
+  @Test
+  public void testUpdateCount() {
+    try (Connection connection = createConnection()) {
+
+      // Setting a different update count is reflected in the update count that is returned by
+      // an auto-batch.
+      try {
+        // Disable update count verification to prevent errors.
+        connection.setAutoBatchDmlUpdateCountVerification(false);
+
+        connection.setAutoBatchDmlUpdateCount(2L);
+        assertEquals(2L, connection.executeUpdate(INSERT_STATEMENT));
+        // The update count can be modified during the batch.
+        connection.setAutoBatchDmlUpdateCount(3L);
+        assertEquals(3L, connection.executeUpdate(INSERT_STATEMENT));
+
+        connection.commit();
+
+        // The auto-batch update count setting is not used for explicit batches.
+        connection.startBatchDml();
+        assertEquals(-1L, connection.executeUpdate(INSERT_STATEMENT));
+        connection.runBatch();
+        connection.commit();
+      } finally {
+        // TODO: Remove once a normal connection variable is used for this.
+        System.clearProperty("spanner.auto_batch_dml_update_count");
+        System.clearProperty("spanner.auto_batch_dml_update_count_verification");
+      }
+    }
+
+    assertEquals(2, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(2, mockSpanner.countRequestsOfType(CommitRequest.class));
+  }
+
+  @Test
+  public void testUpdateCountVerification_failsIfDifferent() {
+    try (Connection connection = createConnection()) {
+      try {
+        assertEquals(1L, connection.executeUpdate(INSERT_STATEMENT));
+        // Set a different (expected) update count. This will cause the batch to fail, as the
+        // actual update count will be 1.
+        connection.setAutoBatchDmlUpdateCount(3L);
+        assertEquals(3L, connection.executeUpdate(INSERT_STATEMENT));
+
+        assertThrows(SpannerException.class, connection::commit);
+      } finally {
+        // TODO: Remove once a normal connection variable is used for this.
+        System.clearProperty("spanner.auto_batch_dml_update_count");
+      }
+    }
+
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+  }
+
+  @Test
+  public void testUpdateCountVerification_succeedsIfSame() {
+    Statement statement1 = Statement.of("insert into foo (id, value) values (1, 'One')");
+    Statement statement2 = Statement.of("insert into foo (id, value) values (2, 'Two')");
+    Statement statement3 = Statement.of("insert into foo (id, value) values (3, 'Three')");
+    Statement statement4 = Statement.of("insert into foo (id, value) values (4, 'Four')");
+    Statement statement5 = Statement.of("insert into foo (id, value) values (5, 'Five')");
+    mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement1, 1L));
+    mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement2, 2L));
+    mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement3, 3L));
+    mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement4, 3L));
+    mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement5, 4L));
+
+    try (Connection connection = createConnection()) {
+      try {
+        connection.setAutoBatchDmlUpdateCount(1L);
+        assertEquals(1L, connection.executeUpdate(statement1));
+
+        connection.setAutoBatchDmlUpdateCount(2L);
+        assertEquals(2L, connection.executeUpdate(statement2));
+
+        connection.setAutoBatchDmlUpdateCount(3L);
+        assertArrayEquals(
+            new long[] {3L, 3L},
+            connection.executeBatchUpdate(ImmutableList.of(statement3, statement4)));
+
+        connection.setAutoBatchDmlUpdateCount(4L);
+        assertEquals(4L, connection.executeUpdate(statement5));
+
+        connection.commit();
+      } finally {
+        // TODO: Remove once a normal connection variable is used for this.
+        System.clearProperty("spanner.auto_batch_dml_update_count");
+      }
+    }
+
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    ExecuteBatchDmlRequest request =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
+    assertEquals(5, request.getStatementsCount());
+    assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+  }
 }