Repository: googleapis/java-spanner-jdbc
Pull request #1795 (2024-10-21 16:13:12)

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.



Patch:
diff --git a/.github/workflows/unmanaged_dependency_check.yaml b/.github/workflows/unmanaged_dependency_check.yaml
index 514cb3a..df0f894 100644
--- a/.github/workflows/unmanaged_dependency_check.yaml
+++ b/.github/workflows/unmanaged_dependency_check.yaml
@@ -14,6 +14,6 @@ jobs:
       shell: bash
       run: .kokoro/build.sh
     - name: Unmanaged dependency check
-      uses: googleapis/sdk-platform-java/java-shared-dependencies/unmanaged-dependency-check@google-cloud-shared-dependencies/v3.37.0
+      uses: googleapis/sdk-platform-java/java-shared-dependencies/unmanaged-dependency-check@google-cloud-shared-dependencies/v3.38.0
       with:
         bom-path: pom.xml
diff --git a/.kokoro/presubmit/graalvm-native-17.cfg b/.kokoro/presubmit/graalvm-native-17.cfg
index 0f8b919..9a49f7d 100644
--- a/.kokoro/presubmit/graalvm-native-17.cfg
+++ b/.kokoro/presubmit/graalvm-native-17.cfg
@@ -3,7 +3,7 @@
 # Configure the docker image for kokoro-trampoline.
 env_vars: {
   key: "TRAMPOLINE_IMAGE"
-  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_b:3.37.0"
+  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_b:3.38.0"
 }
 
 env_vars: {
diff --git a/.kokoro/presubmit/graalvm-native.cfg b/.kokoro/presubmit/graalvm-native.cfg
index e6553bd..e434ee6 100644
--- a/.kokoro/presubmit/graalvm-native.cfg
+++ b/.kokoro/presubmit/graalvm-native.cfg
@@ -3,7 +3,7 @@
 # Configure the docker image for kokoro-trampoline.
 env_vars: {
   key: "TRAMPOLINE_IMAGE"
-  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_a:3.37.0"
+  value: "gcr.io/cloud-devrel-public-resources/graalvm_sdk_platform_a:3.38.0"
 }
 
 env_vars: {
diff --git a/clirr-ignored-differences.xml b/clirr-ignored-differences.xml
index 431f540..e82c992 100644
--- 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/pom.xml b/pom.xml
index 094ed7c..9defc60 100644
--- a/pom.xml
+++ b/pom.xml
@@ -14,7 +14,7 @@
   <parent>
     <groupId>com.google.cloud</groupId>
     <artifactId>sdk-platform-java-config</artifactId>
-    <version>3.37.0</version>
+    <version>3.38.0</version>
   </parent>
   <developers>
     <developer>
@@ -166,7 +166,7 @@
     <dependency>
       <groupId>org.testcontainers</groupId>
       <artifactId>testcontainers</artifactId>
-      <version>1.20.2</version>
+      <version>1.20.3</version>
       <scope>test</scope>
     </dependency>
     <dependency>
@@ -459,7 +459,7 @@
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-project-info-reports-plugin</artifactId>
-        <version>3.7.0</version>
+        <version>3.8.0</version>
         <reportSets>
           <reportSet>
             <reports>
diff --git a/samples/quickperf/pom.xml b/samples/quickperf/pom.xml
index b1b7d05..4ab006e 100644
--- a/samples/quickperf/pom.xml
+++ b/samples/quickperf/pom.xml
@@ -12,7 +12,7 @@
   <parent>
     <groupId>com.google.cloud</groupId>
     <artifactId>sdk-platform-java-config</artifactId>
-    <version>3.37.0</version>
+    <version>3.38.0</version>
     <relativePath/>
   </parent>
 
@@ -73,7 +73,7 @@
     <dependency>
       <groupId>org.testcontainers</groupId>
       <artifactId>testcontainers</artifactId>
-      <version>1.20.2</version>
+      <version>1.20.3</version>
       <scope>test</scope>
     </dependency>
     <dependency>
@@ -95,7 +95,7 @@
       <plugin>
         <groupId>org.codehaus.mojo</groupId>
         <artifactId>exec-maven-plugin</artifactId>
-        <version>3.4.1</version>
+        <version>3.5.0</version>
         <configuration>
           <mainClass>com.google.cloud.jdbc.quickperf.QuickPerf</mainClass>
         </configuration>
diff --git a/samples/snippets/pom.xml b/samples/snippets/pom.xml
index 5efd8b8..16d2dce 100644
--- a/samples/snippets/pom.xml
+++ b/samples/snippets/pom.xml
@@ -9,7 +9,7 @@
   <parent>
     <groupId>com.google.cloud</groupId>
     <artifactId>sdk-platform-java-config</artifactId>
-    <version>3.37.0</version>
+    <version>3.38.0</version>
     <relativePath/>
   </parent>
 
@@ -54,7 +54,7 @@
     <dependency>
       <groupId>org.testcontainers</groupId>
       <artifactId>testcontainers</artifactId>
-      <version>1.20.2</version>
+      <version>1.20.3</version>
       <scope>test</scope>
     </dependency>
     <dependency>
diff --git a/samples/spring-data-jdbc/pom.xml b/samples/spring-data-jdbc/pom.xml
index bc66cd3..4ea50f0 100644
--- a/samples/spring-data-jdbc/pom.xml
+++ b/samples/spring-data-jdbc/pom.xml
@@ -114,7 +114,7 @@
     <dependency>
       <groupId>org.testcontainers</groupId>
       <artifactId>testcontainers</artifactId>
-      <version>1.20.2</version>
+      <version>1.20.3</version>
       <scope>test</scope>
     </dependency>
   </dependencies>
diff --git a/samples/spring-data-mybatis/pom.xml b/samples/spring-data-mybatis/pom.xml
index 7aece5e..80ebd18 100644
--- a/samples/spring-data-mybatis/pom.xml
+++ b/samples/spring-data-mybatis/pom.xml
@@ -42,7 +42,7 @@
       <dependency>
         <groupId>org.testcontainers</groupId>
         <artifactId>testcontainers-bom</artifactId>
-        <version>1.20.2</version>
+        <version>1.20.3</version>
         <scope>import</scope>
         <type>pom</type>
       </dependency>
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
index 6795054..b8491d1 100644
--- 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 @@ public interface CloudSpannerJdbcConnection extends Connection {
     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
index 26b381d..e603c29 100644
--- 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 @@ class JdbcConnection extends AbstractJdbcConnection {
     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
index 19b3256..5d8dfd9 100644
--- 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 @@ class JdbcStatement extends AbstractJdbcStatement implements CloudSpannerJdbcSta
    *     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 @@ class JdbcStatement extends AbstractJdbcStatement implements CloudSpannerJdbcSta
     }
   }
 
-  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 @@ class JdbcStatement extends AbstractJdbcStatement implements CloudSpannerJdbcSta
   @Override
   public void clearBatch() throws SQLException {
     checkClosed();
-    checkConnectionHasNoActiveBatch();
     batchedStatements.clear();
     this.currentBatchType = BatchType.NONE;
   }
@@ -436,7 +425,6 @@ class JdbcStatement extends AbstractJdbcStatement implements CloudSpannerJdbcSta
 
   private long[] executeBatch(boolean large) throws SQLException {
     checkClosed();
-    checkConnectionHasNoActiveBatch();
     StatementTimeout originalTimeout = setTemporaryStatementTimeout();
     try {
       switch (this.currentBatchType) {


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/AutoBatchDmlMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/AutoBatchDmlMockServerTest.java
new file mode 100644
index 0000000..8ba1a3c
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/jdbc/AutoBatchDmlMockServerTest.java
@@ -0,0 +1,332 @@
+/*
+ * 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.jdbc;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.spanner.v1.ExecuteBatchDmlRequest;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.concurrent.ThreadLocalRandom;
+import org.junit.After;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class AutoBatchDmlMockServerTest extends AbstractMockServerTest {
+  private static final String NON_PARAMETERIZED_INSERT =
+      "insert into foo (id, value) values (1, 'One')";
+  private static final String NON_PARAMETERIZED_UPDATE = "update foo set value='Zero' where id=0";
+  private static final String PARAMETERIZED_INSERT =
+      "insert into foo (id, value) values (@p1, @p2)";
+  private static final String PARAMETERIZED_UPDATE = "update foo set value=@p1 where id=@p2";
+
+  @BeforeClass
+  public static void setup() {
+    mockSpanner.putStatementResult(
+        StatementResult.update(
+            com.google.cloud.spanner.Statement.of(NON_PARAMETERIZED_INSERT), 1L));
+    mockSpanner.putStatementResult(
+        StatementResult.update(
+            com.google.cloud.spanner.Statement.of(NON_PARAMETERIZED_UPDATE), 1L));
+
+    mockSpanner.putStatementResult(
+        StatementResult.update(
+            com.google.cloud.spanner.Statement.newBuilder(PARAMETERIZED_INSERT)
+                .bind("p1")
+                .to(1L)
+                .bind("p2")
+                .to("One")
+                .build(),
+            1L));
+    mockSpanner.putStatementResult(
+        StatementResult.update(
+            com.google.cloud.spanner.Statement.newBuilder(PARAMETERIZED_INSERT)
+                .bind("p1")
+                .to(2L)
+                .bind("p2")
+                .to("Two")
+                .build(),
+            1L));
+    mockSpanner.putStatementResult(
+        StatementResult.update(
+            com.google.cloud.spanner.Statement.newBuilder(PARAMETERIZED_UPDATE)
+                .bind("p2")
+                .to(1L)
+                .bind("p1")
+                .to("One")
+                .build(),
+            1L));
+    mockSpanner.putStatementResult(
+        StatementResult.update(
+            com.google.cloud.spanner.Statement.newBuilder(PARAMETERIZED_UPDATE)
+                .bind("p2")
+                .to(2L)
+                .bind("p1")
+                .to("Two")
+                .build(),
+            1L));
+  }
+
+  @After
+  public void clearRequests() {
+    mockSpanner.clearRequests();
+  }
+
+  @Test
+  public void testStatementExecute() throws SQLException {
+    try (Connection connection = createJdbcConnection()) {
+      connection.setAutoCommit(false);
+      connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true);
+
+      try (Statement statement = connection.createStatement()) {
+        assertFalse(statement.execute(NON_PARAMETERIZED_INSERT));
+        assertFalse(statement.execute(NON_PARAMETERIZED_UPDATE));
+      }
+      connection.commit();
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  @Test
+  public void testStatementExecuteUpdate() throws SQLException {
+    try (Connection connection = createJdbcConnection()) {
+      connection.setAutoCommit(false);
+      connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true);
+
+      try (Statement statement = connection.createStatement()) {
+        assertEquals(1, statement.executeUpdate(NON_PARAMETERIZED_INSERT));
+        assertEquals(1, statement.executeUpdate(NON_PARAMETERIZED_UPDATE));
+      }
+      connection.commit();
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  @Test
+  public void testStatementBatch() throws SQLException {
+    try (Connection connection = createJdbcConnection()) {
+      connection.setAutoCommit(false);
+      connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true);
+
+      try (Statement statement = connection.createStatement()) {
+        repeat(
+            () -> {
+              statement.addBatch(NON_PARAMETERIZED_INSERT);
+              statement.addBatch(NON_PARAMETERIZED_UPDATE);
+              assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
+            },
+            2);
+      }
+      connection.commit();
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  @Test
+  public void testStatementCombination() throws SQLException {
+    try (Connection connection = createJdbcConnection()) {
+      connection.setAutoCommit(false);
+      connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true);
+
+      try (Statement statement = connection.createStatement()) {
+        statement.executeUpdate(NON_PARAMETERIZED_UPDATE);
+        repeat(
+            () -> {
+              statement.addBatch(NON_PARAMETERIZED_INSERT);
+              statement.addBatch(NON_PARAMETERIZED_UPDATE);
+              assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
+            },
+            ThreadLocalRandom.current().nextInt(1, 5));
+        repeat(
+            () -> {
+              statement.execute(NON_PARAMETERIZED_INSERT);
+              statement.executeUpdate(NON_PARAMETERIZED_UPDATE);
+            },
+            ThreadLocalRandom.current().nextInt(1, 5));
+      }
+      connection.commit();
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  @Test
+  public void testPreparedStatementExecute() throws SQLException {
+    try (Connection connection = createJdbcConnection()) {
+      connection.setAutoCommit(false);
+      connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true);
+
+      try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) {
+        statement.setLong(1, 1L);
+        statement.setString(2, "One");
+        assertFalse(statement.execute());
+        statement.setLong(1, 2L);
+        statement.setString(2, "Two");
+        assertFalse(statement.execute());
+      }
+      try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) {
+        statement.setLong(2, 1L);
+        statement.setString(1, "One");
+        assertFalse(statement.execute());
+        statement.setLong(2, 2L);
+        statement.setString(1, "Two");
+        assertFalse(statement.execute());
+      }
+      connection.commit();
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdate() throws SQLException {
+    try (Connection connection = createJdbcConnection()) {
+      connection.setAutoCommit(false);
+      connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true);
+
+      try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) {
+        statement.setLong(1, 1L);
+        statement.setString(2, "One");
+        assertEquals(1, statement.executeUpdate());
+        statement.setLong(1, 2L);
+        statement.setString(2, "Two");
+        assertEquals(1, statement.executeUpdate());
+      }
+      try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) {
+        statement.setLong(2, 1L);
+        statement.setString(1, "One");
+        assertEquals(1, statement.executeUpdate());
+        statement.setLong(2, 2L);
+        statement.setString(1, "Two");
+        assertEquals(1, statement.executeUpdate());
+      }
+      connection.commit();
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  @Test
+  public void testPreparedStatementBatch() throws SQLException {
+    try (Connection connection = createJdbcConnection()) {
+      connection.setAutoCommit(false);
+      connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true);
+
+      repeat(
+          () -> {
+            try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) {
+              statement.setLong(1, 1L);
+              statement.setString(2, "One");
+              statement.addBatch();
+              statement.setLong(1, 2L);
+              statement.setString(2, "Two");
+              statement.addBatch();
+              statement.executeBatch();
+            }
+            try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) {
+              statement.setLong(2, 1L);
+              statement.setString(1, "One");
+              statement.addBatch();
+              statement.setLong(2, 2L);
+              statement.setString(1, "Two");
+              statement.addBatch();
+              statement.executeBatch();
+            }
+          },
+          2);
+      connection.commit();
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  @Test
+  public void testPreparedStatementCombination() throws SQLException {
+    try (Connection connection = createJdbcConnection()) {
+      connection.setAutoCommit(false);
+      connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true);
+
+      try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) {
+        statement.setLong(1, 1L);
+        statement.setString(2, "One");
+        assertFalse(statement.execute());
+      }
+      repeat(
+          () -> {
+            try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) {
+              statement.setLong(1, 1L);
+              statement.setString(2, "One");
+              statement.addBatch();
+              statement.setLong(1, 2L);
+              statement.setString(2, "Two");
+              statement.addBatch();
+              assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
+            }
+            try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) {
+              statement.setLong(2, 1L);
+              statement.setString(1, "One");
+              statement.addBatch();
+              statement.setLong(2, 2L);
+              statement.setString(1, "Two");
+              statement.addBatch();
+              assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
+            }
+          },
+          ThreadLocalRandom.current().nextInt(1, 5));
+      repeat(
+          () -> {
+            try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) {
+              statement.setLong(1, 1L);
+              statement.setString(2, "One");
+              assertEquals(1, statement.executeUpdate());
+            }
+            try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) {
+              statement.setLong(2, 2L);
+              statement.setString(1, "Two");
+              assertFalse(statement.execute());
+            }
+          },
+          ThreadLocalRandom.current().nextInt(1, 5));
+      connection.commit();
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  interface SQLRunnable {
+    void run() throws SQLException;
+  }
+
+  static void repeat(SQLRunnable runnable, int count) throws SQLException {
+    for (int i = 0; i < count; i++) {
+      runnable.run();
+    }
+  }
+}
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
index f6948a9..7dd7153 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcSimpleStatementsTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcSimpleStatementsTest.java
@@ -17,11 +17,11 @@
 package com.google.cloud.spanner.jdbc.it;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeFalse;
 
 import com.google.cloud.spanner.Database;
@@ -102,12 +102,6 @@ public class ITJdbcSimpleStatementsTest extends ITAbstractJdbcTest {
 
   @Test
   public void testPreparedStatement() throws SQLException {
-    // skipping the test when dialect is POSTGRESQL because of exception below
-    // INVALID_ARGUMENT: io.grpc.StatusRuntimeException: INVALID_ARGUMENT: Statements with set
-    // operations in subqueries are not supported
-    assumeFalse(
-        "select array of structs is not supported on POSTGRESQL",
-        dialect.dialect == Dialect.POSTGRESQL);
     String sql =
         "select * from (select 1 as number union all select 2 union all select 3) numbers where number=?";
     try (Connection connection = createConnection(env, database)) {
@@ -188,15 +182,16 @@ public class ITJdbcSimpleStatementsTest extends ITAbstractJdbcTest {
   }
 
   @Test
-  public void testAddBatchWhenAlreadyInBatch() {
+  public void testAddBatchWhenAlreadyInBatch() throws SQLException {
     try (Connection connection = createConnection(env, database)) {
-      connection.createStatement().execute("START BATCH DML");
-      connection.createStatement().addBatch("INSERT INTO Singers (SingerId) VALUES (-1)");
-      fail("missing expected exception");
-    } catch (SQLException e) {
-      assertThat(e.getMessage())
-          .contains(
-              "Calling addBatch() is not allowed when a DML or DDL batch has been started on the connection.");
+      try (Statement statement = connection.createStatement()) {
+        statement.execute("START BATCH DML");
+        statement.addBatch("INSERT INTO Singers (SingerId) VALUES (-1)");
+        statement.addBatch("INSERT INTO Singers (SingerId) VALUES (-2)");
+        // The returned update count for DML statements in a batch is -1.
+        assertArrayEquals(new int[] {-1, -1}, statement.executeBatch());
+        // Note: The 'Singers' table does not actually exist, so we're not executing the batch.
+      }
     }
   }