Repository: googleapis/java-spanner-jdbc
Pull request #1802 (2024-10-24 16:44:32)

deps: update dependency org.springframework.boot:spring-boot-starter-data-jdbc to v3.3.5
This PR contains the following updates:

| Package | Change | Age | Adoption | Passing | Confidence |
|---|---|---|---|---|---|
| [org.springframework.boot:spring-boot-starter-data-jdbc](https://spring.io/projects/spring-boot) ([source](https://redirect.github.com/spring-projects/spring-boot)) | `3.3.4` -> `3.3.5` | [![age](https://developer.mend.io/api/mc/badges/age/maven/org.springframework.boot:spring-boot-starter-data-jdbc/3.3.5?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![adoption](https://developer.mend.io/api/mc/badges/adoption/maven/org.springframework.boot:spring-boot-starter-data-jdbc/3.3.5?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![passing](https://developer.mend.io/api/mc/badges/compatibility/maven/org.springframework.boot:spring-boot-starter-data-jdbc/3.3.4/3.3.5?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![confidence](https://developer.mend.io/api/mc/badges/confidence/maven/org.springframework.boot:spring-boot-starter-data-jdbc/3.3.4/3.3.5?slim=true)](https://docs.renovatebot.com/merge-confidence/) |

---

### Release Notes

<details>
<summary>spring-projects/spring-boot (org.springframework.boot:spring-boot-starter-data-jdbc)</summary>

### [`v3.3.5`](https://redirect.github.com/spring-projects/spring-boot/compare/v3.3.4...v3.3.5)

</details>

---

### Configuration

📅 **Schedule**: Branch creation - At any time (no schedule defined), Automerge - At any time (no schedule defined).

🚦 **Automerge**: Disabled by config. Please merge this manually once you are satisfied.

♻ **Rebasing**: Whenever PR becomes conflicted, or you tick the rebase/retry checkbox.

🔕 **Ignore**: Close this PR and you won't be reminded about this update again.

---

 - [ ] <!-- rebase-check -->If you want to rebase/retry this PR, check this box

---

This PR was generated by [Mend Renovate](https://mend.io/renovate/). View the [repository job log](https://developer.mend.io/github/googleapis/java-spanner-jdbc).
<!--renovate-debug:eyJjcmVhdGVkSW5WZXIiOiIzOC4xMjAuMSIsInVwZGF0ZWRJblZlciI6IjM4LjEyMC4xIiwidGFyZ2V0QnJhbmNoIjoibWFpbiIsImxhYmVscyI6W119-->


dpebot: /gcbrun

Patch:
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 9defc60..130335d 100644
--- a/pom.xml
+++ b/pom.xml
@@ -61,7 +61,7 @@
       <dependency>
         <groupId>com.google.cloud</groupId>
         <artifactId>google-cloud-spanner-bom</artifactId>
-        <version>6.79.0</version>
+        <version>6.80.0</version>
         <type>pom</type>
         <scope>import</scope>
       </dependency>
diff --git a/samples/quickperf/pom.xml b/samples/quickperf/pom.xml
index 4ab006e..98efb3e 100644
--- a/samples/quickperf/pom.xml
+++ b/samples/quickperf/pom.xml
@@ -79,7 +79,7 @@
     <dependency>
       <groupId>org.springframework.boot</groupId>
       <artifactId>spring-boot</artifactId>
-      <version>3.3.4</version>
+      <version>3.3.5</version>
       <scope>test</scope>
     </dependency>
     <dependency>
diff --git a/samples/spring-data-jdbc/pom.xml b/samples/spring-data-jdbc/pom.xml
index 4ea50f0..bd11b7f 100644
--- a/samples/spring-data-jdbc/pom.xml
+++ b/samples/spring-data-jdbc/pom.xml
@@ -30,7 +30,7 @@
       <dependency>
         <groupId>com.google.cloud</groupId>
         <artifactId>google-cloud-spanner-bom</artifactId>
-        <version>6.79.0</version>
+        <version>6.80.0</version>
         <scope>import</scope>
         <type>pom</type>
       </dependency>
@@ -55,7 +55,7 @@
     <dependency>
       <groupId>org.springframework.boot</groupId>
       <artifactId>spring-boot-starter-data-jdbc</artifactId>
-      <version>3.3.4</version>
+      <version>3.3.5</version>
     </dependency>
 
     <!-- Add both the Cloud Spanner and the PostgreSQL JDBC driver. -->
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.
+      }
     }
   }