feat: support DML auto-batching in Connection API (#3386)

* feat: support DML auto-batching in Connection API

Adds support for automatically batching DML statements that are
executed on a connection. This can be used in combination with
frameworks like Hibernate to improve performance, especially on
multi-region instances with high round-trip latencies.

* chore: generate libraries at Sun Oct  6 06:00:17 UTC 2024

* feat: add update count verification

* fix: add ignored diffs

* test: add transaction retry test

* chore: cleanup and add comments

* chore: generate libraries at Fri Oct 11 15:21:53 UTC 2024

* feat: use connection variables for auto_batch_dml

---------

Co-authored-by: cloud-java-bot <[email protected]>

Author: olavloite

google-cloud-spanner/clirr-ignored-differences.xml

@@ -758,4 +758,37 @@
     <className>com/google/cloud/spanner/connection/Connection</className>
     <method>boolean isKeepTransactionAlive()</method>
   </difference>
+  
+  <!-- Automatic DML batching -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void setAutoBatchDml(boolean)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>boolean isAutoBatchDml()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void setAutoBatchDmlUpdateCount(long)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>long getAutoBatchDmlUpdateCount()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>void setAutoBatchDmlUpdateCountVerification(boolean)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>boolean isAutoBatchDmlUpdateCountVerification()</method>
+  </difference>
+  
 </differences>

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

@@ -0,0 +1,68 @@
+/*
+ * 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 final long[] expected;
+
+  private final long[] actual;
+
+  /** 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);
+    this.expected = expected;
+    this.actual = actual;
+  }
+
+  /**
+   * The expected update counts. These were returned to the client application when the DML
+   * statements were buffered.
+   */
+  public long[] getExpected() {
+    return Arrays.copyOf(this.expected, this.expected.length);
+  }
+
+  /**
+   * The actual update counts. These were returned by Spanner to the client when the DML statements
+   * were actually executed, and are the update counts that the client application would have
+   * received if auto-batching had not been enabled.
+   */
+  public long[] getActual() {
+    return Arrays.copyOf(this.actual, this.actual.length);
+  }
+}

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

@@ -116,6 +116,13 @@ public static SpannerBatchUpdateException newSpannerBatchUpdateException(
     return new SpannerBatchUpdateException(token, code, message, updateCounts);
   }
 
+  /** Constructs a specific error that */
+  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/ClientSideStatementValueConverters.java

@@ -177,6 +177,35 @@ public Integer convert(String value) {
     }
   }
 
+  /** Converter from string to a long. */
+  static class LongConverter implements ClientSideStatementValueConverter<Long> {
+    static final LongConverter INSTANCE = new LongConverter();
+
+    private LongConverter() {}
+
+    /** Constructor needed for reflection. */
+    public LongConverter(String allowedValues) {}
+
+    @Override
+    public Class<Long> getParameterClass() {
+      return Long.class;
+    }
+
+    @Override
+    public Long convert(String value) {
+      try {
+        long res = Long.parseLong(value);
+        if (res < 0) {
+          // The conventions for these converters is to return null if the value is invalid.
+          return null;
+        }
+        return res;
+      } catch (Exception ignore) {
+        return null;
+      }
+    }
+  }
+
   /** Converter from string to {@link Duration}. */
   static class DurationConverter implements ClientSideStatementValueConverter<Duration> {
     static final DurationConverter INSTANCE =

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

@@ -1165,6 +1165,45 @@ default StatementResult execute(Statement statement, Set<ResultType> allowedResu
    */
   ResultSet analyzeQuery(Statement query, QueryAnalyzeMode queryMode);
 
+  /**
+   * Enables or disables automatic batching of DML statements. When enabled, DML statements that are
+   * executed on this connection will be buffered in memory instead of actually being executed. The
+   * buffered DML statements are flushed to Spanner when a statement that cannot be part of a DML
+   * batch is executed on the connection. This can be a query, a DDL statement with a THEN RETURN
+   * clause, or a Commit call. The update count that is returned for DML statements that are
+   * buffered is determined by the value that has been set with {@link
+   * #setAutoBatchDmlUpdateCount(long)}. The default is 1. The connection verifies that the update
+   * counts that were returned while buffering DML statements match the actual update counts that
+   * are returned by Spanner when the batch is executed. This verification can be disabled by
+   * calling {@link #setAutoBatchDmlUpdateCountVerification(boolean)}.
+   */
+  void setAutoBatchDml(boolean autoBatchDml);
+
+  /** Returns whether automatic DML batching is enabled on this connection. */
+  boolean isAutoBatchDml();
+
+  /**
+   * Sets the update count that is returned for DML statements that are buffered during an automatic
+   * DML batch. This value is only used if {@link #isAutoBatchDml()} is enabled.
+   */
+  void setAutoBatchDmlUpdateCount(long updateCount);
+
+  /**
+   * Returns the update count that is returned for DML statements that are buffered during an
+   * automatic DML batch.
+   */
+  long getAutoBatchDmlUpdateCount();
+
+  /**
+   * Sets whether the update count that is returned by Spanner after executing an automatic DML
+   * batch should be verified against the update counts that were returned during the buffering of
+   * those statements.
+   */
+  void setAutoBatchDmlUpdateCountVerification(boolean verification);
+
+  /** Indicates whether the update counts of automatic DML batches should be verified. */
+  boolean isAutoBatchDmlUpdateCountVerification();
+
   /**
    * Enable data boost for partitioned queries. See also {@link #partitionQuery(Statement,
    * PartitionOptions, QueryOption...)}