feat: auto_batch_dml methods + combination with JDBC batching (#1795)

olavloite · web-flow · commit 98f01bba45ec · 2024-10-25T11:24:07.000+02:00
* feat: auto_batch_dml methods + combination with JDBC batching

Adds methods for enabling/disabling auto_batch_dml.

Also enables the use of this feature in combination with other JDBC
batches. This again allows PreparedStatement#addBatch() and
Statement#addBatch(String) to be used in combination with
auto_batch_dml. This again makes it usable in combination with
frameworks like Hibernate.

* test: remove expected error from test

* chore: add clirr diffs
diff --git a/clirr-ignored-differences.xml b/clirr-ignored-differences.xml
@@ -76,4 +76,36 @@
     <differenceType>8001</differenceType>
     <className>com/google/cloud/spanner/connection/ConnectionHelper</className>
   </difference>
+
+  <!-- DML auto batching -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>boolean isAutoBatchDml()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setAutoBatchDml(boolean)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>long getAutoBatchDmlUpdateCount()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setAutoBatchDmlUpdateCount(long)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>boolean isAutoBatchDmlUpdateCountVerification()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setAutoBatchDmlUpdateCountVerification(boolean)</method>
+  </difference>
 </differences>
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
@@ -406,6 +406,57 @@ default int getMaxPartitionedParallelism() throws SQLException {
     throw new UnsupportedOperationException();
   }
 
+  /**
+   * 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)}.
+   */
+  default void setAutoBatchDml(boolean autoBatchDml) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /** Returns whether automatic DML batching is enabled on this connection. */
+  default boolean isAutoBatchDml() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * 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.
+   */
+  default void setAutoBatchDmlUpdateCount(long updateCount) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Returns the update count that is returned for DML statements that are buffered during an
+   * automatic DML batch.
+   */
+  default long getAutoBatchDmlUpdateCount() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * 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.
+   */
+  default void setAutoBatchDmlUpdateCountVerification(boolean verification) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /** Indicates whether the update counts of automatic DML batches should be verified. */
+  default boolean isAutoBatchDmlUpdateCountVerification() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
   /**
    * @see
    *     com.google.cloud.spanner.connection.Connection#addTransactionRetryListener(com.google.cloud.spanner.connection.TransactionRetryListener)
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
@@ -770,6 +770,36 @@ public int getMaxPartitionedParallelism() throws SQLException {
     return get(Connection::getMaxPartitionedParallelism);
   }
 
+  @Override
+  public void setAutoBatchDml(boolean autoBatchDml) throws SQLException {
+    set(Connection::setAutoBatchDml, autoBatchDml);
+  }
+
+  @Override
+  public boolean isAutoBatchDml() throws SQLException {
+    return get(Connection::isAutoBatchDml);
+  }
+
+  @Override
+  public void setAutoBatchDmlUpdateCount(long updateCount) throws SQLException {
+    set(Connection::setAutoBatchDmlUpdateCount, updateCount);
+  }
+
+  @Override
+  public long getAutoBatchDmlUpdateCount() throws SQLException {
+    return get(Connection::getAutoBatchDmlUpdateCount);
+  }
+
+  @Override
+  public void setAutoBatchDmlUpdateCountVerification(boolean verification) throws SQLException {
+    set(Connection::setAutoBatchDmlUpdateCountVerification, verification);
+  }
+
+  @Override
+  public boolean isAutoBatchDmlUpdateCountVerification() throws SQLException {
+    return get(Connection::isAutoBatchDmlUpdateCountVerification);
+  }
+
   @SuppressWarnings("deprecation")
   private static final class JdbcToSpannerTransactionRetryListener
       implements com.google.cloud.spanner.connection.TransactionRetryListener {
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
@@ -391,7 +391,6 @@ private BatchType determineStatementBatchType(String sql) throws SQLException {
    *     client side statement) or if the connection of this statement has an active batch.
    */
   void checkAndSetBatchType(String sql) throws SQLException {
-    checkConnectionHasNoActiveBatch();
     BatchType type = determineStatementBatchType(sql);
     if (this.currentBatchType == BatchType.NONE) {
       this.currentBatchType = type;
@@ -401,15 +400,6 @@ void checkAndSetBatchType(String sql) throws SQLException {
     }
   }
 
-  private void checkConnectionHasNoActiveBatch() throws SQLException {
-    if (getConnection().getSpannerConnection().isDdlBatchActive()
-        || getConnection().getSpannerConnection().isDmlBatchActive()) {
-      throw JdbcSqlExceptionFactory.of(
-          "Calling addBatch() is not allowed when a DML or DDL batch has been started on the connection.",
-          Code.FAILED_PRECONDITION);
-    }
-  }
-
   @Override
   public void addBatch(String sql) throws SQLException {
     checkClosed();
@@ -420,7 +410,6 @@ public void addBatch(String sql) throws SQLException {
   @Override
   public void clearBatch() throws SQLException {
     checkClosed();
-    checkConnectionHasNoActiveBatch();
     batchedStatements.clear();
     this.currentBatchType = BatchType.NONE;
   }
@@ -436,7 +425,6 @@ public long[] executeLargeBatch() throws SQLException {
 
   private long[] executeBatch(boolean large) throws SQLException {
     checkClosed();
-    checkConnectionHasNoActiveBatch();
     StatementTimeout originalTimeout = setTemporaryStatementTimeout();
     try {
       switch (this.currentBatchType) {
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/AutoBatchDmlMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/AutoBatchDmlMockServerTest.java
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcSimpleStatementsTest.java b/src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcSimpleStatementsTest.java
