Repository: googleapis/java-spanner-jdbc
Pull request #1347 (2023-09-15 10:20:09)

deps: update dependency org.springframework.data:spring-data-bom to v2023.0.4
[![Mend Renovate](https://app.renovatebot.com/images/banner.svg)](https://renovatebot.com)

This PR contains the following updates:

| Package | Change | Age | Adoption | Passing | Confidence |
|---|---|---|---|---|---|
| [org.springframework.data:spring-data-bom](https://togithub.com/spring-projects/spring-data-bom) | `2023.0.3` -> `2023.0.4` | [![age](https://developer.mend.io/api/mc/badges/age/maven/org.springframework.data:spring-data-bom/2023.0.4?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![adoption](https://developer.mend.io/api/mc/badges/adoption/maven/org.springframework.data:spring-data-bom/2023.0.4?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![passing](https://developer.mend.io/api/mc/badges/compatibility/maven/org.springframework.data:spring-data-bom/2023.0.3/2023.0.4?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![confidence](https://developer.mend.io/api/mc/badges/confidence/maven/org.springframework.data:spring-data-bom/2023.0.3/2023.0.4?slim=true)](https://docs.renovatebot.com/merge-confidence/) |

---

### Release Notes

<details>
<summary>spring-projects/spring-data-bom (org.springframework.data:spring-data-bom)</summary>

### [`v2023.0.4`](https://togithub.com/spring-projects/spring-data-bom/releases/tag/2023.0.4)

[Compare Source](https://togithub.com/spring-projects/spring-data-bom/compare/2023.0.3...2023.0.4)

#### :shipit: Participating Modules

-   [Spring Data BOM 2023.0.4](https://togithub.com/spring-projects/spring-data-bom/releases/tag/2023.0.4)
-   [Spring Data Build 3.1.4](https://togithub.com/spring-projects/spring-data-build/releases/tag/3.1.4)
-   [Spring Data Cassandra 4.1.4](https://togithub.com/spring-projects/spring-data-cassandra/releases/tag/4.1.4)
-   [Spring Data Commons 3.1.4](https://togithub.com/spring-projects/spring-data-commons/releases/tag/3.1.4)
-   [Spring Data Couchbase 5.1.4](https://togithub.com/spring-projects/spring-data-couchbase/releases/tag/5.1.4)
-   [Spring Data Elasticsearch 5.1.4](https://togithub.com/spring-projects/spring-data-elasticsearch/releases/tag/5.1.4)
-   [Spring Data JPA 3.1.4](https://togithub.com/spring-projects/spring-data-jpa/releases/tag/3.1.4)
-   [Spring Data KeyValue 3.1.4](https://togithub.com/spring-projects/spring-data-keyvalue/releases/tag/3.1.4)
-   [Spring Data LDAP 3.1.4](https://togithub.com/spring-projects/spring-data-ldap/releases/tag/3.1.4)
-   [Spring Data MongoDB 4.1.4](https://togithub.com/spring-projects/spring-data-mongodb/releases/tag/4.1.4)
-   [Spring Data Neo4j 7.1.4](https://togithub.com/spring-projects/spring-data-neo4j/releases/tag/7.1.4)
-   [Spring Data REST 4.1.4](https://togithub.com/spring-projects/spring-data-rest/releases/tag/4.1.4)
-   [Spring Data Redis 3.1.4](https://togithub.com/spring-projects/spring-data-redis/releases/tag/3.1.4)
-   [Spring Data Relational 3.1.4](https://togithub.com/spring-projects/spring-data-relational/releases/tag/3.1.4)

</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 has been generated by [Mend Renovate](https://www.mend.io/free-developer-tools/renovate/). View repository job log [here](https://developer.mend.io/github/googleapis/java-spanner-jdbc).
<!--renovate-debug:eyJjcmVhdGVkSW5WZXIiOiIzNi44My4wIiwidXBkYXRlZEluVmVyIjoiMzYuODMuMCIsInRhcmdldEJyYW5jaCI6Im1haW4ifQ==-->



Patch:
diff --git a/clirr-ignored-differences.xml b/clirr-ignored-differences.xml
index a92ad53..acf11c5 100644
--- a/clirr-ignored-differences.xml
+++ b/clirr-ignored-differences.xml
@@ -16,4 +16,46 @@
     <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
     <method>void setSavepointSupport(com.google.cloud.spanner.connection.SavepointSupport)</method>
   </difference>
+
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>int getMaxPartitionedParallelism()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>int getMaxPartitions()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>boolean isAutoPartitionMode()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>boolean isDataBoostEnabled()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setAutoPartitionMode(boolean)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setDataBoostEnabled(boolean)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setMaxPartitionedParallelism(int)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setMaxPartitions(int)</method>
+  </difference>
+
 </differences>
diff --git a/samples/spring-data-jdbc/pom.xml b/samples/spring-data-jdbc/pom.xml
index ee5a443..ba78496 100644
--- a/samples/spring-data-jdbc/pom.xml
+++ b/samples/spring-data-jdbc/pom.xml
@@ -23,7 +23,7 @@
       <dependency>
         <groupId>org.springframework.data</groupId>
         <artifactId>spring-data-bom</artifactId>
-        <version>2023.0.3</version>
+        <version>2023.0.4</version>
         <scope>import</scope>
         <type>pom</type>
       </dependency>
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
index e7b1a47..9713172 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
@@ -24,12 +24,15 @@ import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
+import com.google.rpc.Code;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 /** Base class for Cloud Spanner JDBC {@link Statement}s */
 abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Statement {
@@ -117,6 +120,25 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
     }
   }
 
+  /** Functional interface that throws {@link SQLException}. */
+  interface JdbcFunction<T, R> {
+    R apply(T t) throws SQLException;
+  }
+
+  /** Runs the given function with the timeout that has been set on this statement. */
+  protected <T> T runWithStatementTimeout(JdbcFunction<Connection, T> function)
+      throws SQLException {
+    checkClosed();
+    StatementTimeout originalTimeout = setTemporaryStatementTimeout();
+    try {
+      return function.apply(getConnection().getSpannerConnection());
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    } finally {
+      resetStatementTimeout(originalTimeout);
+    }
+  }
+
   /**
    * Sets the statement timeout of the Spanner {@link Connection} to the query timeout of this JDBC
    * {@link Statement} and returns the original timeout of the Spanner {@link Connection} so it can
@@ -183,21 +205,74 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
       QueryAnalyzeMode analyzeMode,
       QueryOption... options)
       throws SQLException {
+    Options.QueryOption[] queryOptions = getQueryOptions(options);
+    return doWithStatementTimeout(
+        () -> {
+          com.google.cloud.spanner.ResultSet resultSet;
+          if (analyzeMode == null) {
+            resultSet = connection.getSpannerConnection().executeQuery(statement, queryOptions);
+          } else {
+            resultSet = connection.getSpannerConnection().analyzeQuery(statement, analyzeMode);
+          }
+          return JdbcResultSet.of(this, resultSet);
+        });
+  }
+
+  private <T> T doWithStatementTimeout(Supplier<T> runnable) throws SQLException {
+    return doWithStatementTimeout(runnable, ignore -> Boolean.TRUE);
+  }
+
+  private <T> T doWithStatementTimeout(
+      Supplier<T> runnable, Function<T, Boolean> shouldResetTimeout) throws SQLException {
     StatementTimeout originalTimeout = setTemporaryStatementTimeout();
+    T result = null;
     try {
-      com.google.cloud.spanner.ResultSet resultSet;
-      if (analyzeMode == null) {
-        resultSet =
-            connection.getSpannerConnection().executeQuery(statement, getQueryOptions(options));
-      } else {
-        resultSet = connection.getSpannerConnection().analyzeQuery(statement, analyzeMode);
-      }
-      return JdbcResultSet.of(this, resultSet);
-    } catch (SpannerException e) {
-      throw JdbcSqlExceptionFactory.of(e);
+      result = runnable.get();
+      return result;
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
     } finally {
-      resetStatementTimeout(originalTimeout);
+      if (shouldResetTimeout.apply(result)) {
+        resetStatementTimeout(originalTimeout);
+      }
+    }
+  }
+
+  /**
+   * Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
+   * statement.
+   *
+   * @param statement The SQL statement to execute
+   * @return the number of rows that was inserted/updated/deleted
+   * @throws SQLException if a database error occurs, or if the number of rows affected is larger
+   *     than {@link Integer#MAX_VALUE}
+   */
+  int executeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
+    return checkedCast(executeLargeUpdate(statement));
+  }
+
+  /**
+   * Do a checked cast from long to int. Throws a {@link SQLException} with code {@link
+   * Code#OUT_OF_RANGE} if the update count is too big to fit in an int.
+   */
+  int checkedCast(long updateCount) throws SQLException {
+    if (updateCount > Integer.MAX_VALUE) {
+      throw JdbcSqlExceptionFactory.of(
+          "update count too large for executeUpdate: " + updateCount, Code.OUT_OF_RANGE);
     }
+    return (int) updateCount;
+  }
+
+  /**
+   * Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
+   * statement.
+   *
+   * @param statement The SQL statement to execute
+   * @return the number of rows that was inserted/updated/deleted
+   * @throws SQLException if a database error occurs
+   */
+  long executeLargeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
+    return doWithStatementTimeout(() -> connection.getSpannerConnection().executeUpdate(statement));
   }
 
   /**
@@ -210,24 +285,16 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
    * @throws SQLException if a database error occurs.
    */
   StatementResult execute(com.google.cloud.spanner.Statement statement) throws SQLException {
-    StatementTimeout originalTimeout = setTemporaryStatementTimeout();
-    boolean mustResetTimeout = false;
-    try {
-      StatementResult result = connection.getSpannerConnection().execute(statement);
-      mustResetTimeout = !resultIsSetStatementTimeout(result);
-      if (mustResetTimeout && resultIsShowStatementTimeout(result)) {
-        // it was a 'SHOW STATEMENT_TIMEOUT statement, we need to re-run to get the correct value
-        mustResetTimeout = false;
-        result = rerunShowStatementTimeout(statement, originalTimeout);
-      }
-      return result;
-    } catch (SpannerException e) {
-      throw JdbcSqlExceptionFactory.of(e);
-    } finally {
-      if (mustResetTimeout) {
-        resetStatementTimeout(originalTimeout);
-      }
+    StatementResult statementResult =
+        doWithStatementTimeout(
+            () -> connection.getSpannerConnection().execute(statement),
+            result -> !resultIsSetStatementTimeout(result));
+    if (resultIsShowStatementTimeout(statementResult)) {
+      // We can safely re-run it without first resetting the timeout to the original value, as that
+      // has already been done by the 'doWithStatementTimeout' function.
+      return rerunShowStatementTimeout(statement);
     }
+    return statementResult;
   }
 
   /**
@@ -250,18 +317,22 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
    *     executed was a SET STATEMENT_TIMEOUT statement.
    */
   private boolean resultIsSetStatementTimeout(StatementResult result) {
-    return result.getClientSideStatementType() == ClientSideStatementType.SET_STATEMENT_TIMEOUT;
+    return result != null
+        && result.getClientSideStatementType() == ClientSideStatementType.SET_STATEMENT_TIMEOUT;
   }
 
   private boolean resultIsShowStatementTimeout(StatementResult result) {
-    return result.getClientSideStatementType() == ClientSideStatementType.SHOW_STATEMENT_TIMEOUT;
+    return result != null
+        && result.getClientSideStatementType() == ClientSideStatementType.SHOW_STATEMENT_TIMEOUT;
   }
 
-  private StatementResult rerunShowStatementTimeout(
-      com.google.cloud.spanner.Statement statement, StatementTimeout originalTimeout)
+  private StatementResult rerunShowStatementTimeout(com.google.cloud.spanner.Statement statement)
       throws SQLException {
-    resetStatementTimeout(originalTimeout);
-    return connection.getSpannerConnection().execute(statement);
+    try {
+      return connection.getSpannerConnection().execute(statement);
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    }
   }
 
   @Override
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 cf368f6..c161ee5 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
@@ -22,12 +22,14 @@ import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.CommitStats;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.Mutation;
-import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.AutocommitDmlMode;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.connection.TransactionMode;
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Timestamp;
 import java.util.Iterator;
@@ -334,6 +336,73 @@ public interface CloudSpannerJdbcConnection extends Connection {
     return Dialect.GOOGLE_STANDARD_SQL;
   }
 
+  /**
+   * Enable data boost for partitioned queries. See also {@link
+   * CloudSpannerJdbcStatement#partitionQuery(String, PartitionOptions, QueryOption...)} and {@link
+   * CloudSpannerJdbcPreparedStatement#partitionQuery(PartitionOptions, QueryOption...)}.
+   */
+  default void setDataBoostEnabled(boolean dataBoostEnabled) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Returns whether data boost is enabled for partitioned queries. See also {@link
+   * CloudSpannerJdbcStatement#partitionQuery(String, PartitionOptions, QueryOption...)} and {@link
+   * CloudSpannerJdbcPreparedStatement#partitionQuery(PartitionOptions, QueryOption...)}.
+   */
+  default boolean isDataBoostEnabled() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Sets whether this connection should always use partitioned queries when a query is executed on
+   * this connection. Setting this flag to <code>true</code> and then executing a query that cannot
+   * be partitioned, or executing a query in a read/write transaction, will cause an error. Use this
+   * flag in combination with {@link #setDataBoostEnabled(boolean)} to force all queries on this
+   * connection to use data boost.
+   */
+  default void setAutoPartitionMode(boolean alwaysUsePartitionedQueries) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /** Returns whether this connection will execute all queries as partitioned queries. */
+  default boolean isAutoPartitionMode() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Sets the maximum number of partitions that should be included as a hint to Cloud Spanner when
+   * partitioning a query on this connection. Note that this is only a hint and Cloud Spanner might
+   * choose to ignore the hint.
+   */
+  default void setMaxPartitions(int maxPartitions) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Gets the maximum number of partitions that should be included as a hint to Cloud Spanner when
+   * partitioning a query on this connection. Note that this is only a hint and Cloud Spanner might
+   * choose to ignore the hint.
+   */
+  default int getMaxPartitions() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Sets the maximum degree of parallelism that is used when executing a partitioned query. A
+   * partitioned query will use up to <code>maxThreads</code> to execute and retrieve the results
+   * from Cloud Spanner. Set this value to <code>0</code>> to use the number of available processors
+   * as returned by {@link Runtime#availableProcessors()}.
+   */
+  default void setMaxPartitionedParallelism(int maxThreads) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /** Returns the maximum degree of parallelism that is used for partitioned queries. */
+  default int getMaxPartitionedParallelism() 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/CloudSpannerJdbcPartitionedQueryResultSet.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPartitionedQueryResultSet.java
new file mode 100644
index 0000000..3b6118f
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPartitionedQueryResultSet.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright 2023 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 com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
+import java.sql.ResultSet;
+
+/**
+ * Result set that is returned for partitioned queries, e.g. for 'run partitioned query select ...'
+ * or for {@link CloudSpannerJdbcPreparedStatement#runPartitionedQuery(PartitionOptions,
+ * QueryOption...)}.
+ */
+public interface CloudSpannerJdbcPartitionedQueryResultSet extends ResultSet {
+  /** Returns the number of partitions that this result set contains. */
+  int getNumPartitions();
+
+  /** Returns the degree of parallelism that this result set uses. */
+  int getParallelism();
+}
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPreparedStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPreparedStatement.java
new file mode 100644
index 0000000..8e72d95
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPreparedStatement.java
@@ -0,0 +1,61 @@
+/*
+ * Copyright 2023 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 com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+
+/**
+ * This interface is implemented by {@link PreparedStatement}s that are created on Cloud Spanner
+ * JDBC connections.
+ */
+public interface CloudSpannerJdbcPreparedStatement extends PreparedStatement {
+
+  /**
+   * Partitions this query, so it can be executed in parallel. This method returns a {@link
+   * ResultSet} with a string-representation of the partitions that were created. These strings can
+   * be used to execute a partition either on this connection or an any other connection (on this
+   * host or an any other host) by calling the method {@link #runPartition()}. This method will
+   * automatically enable data boost for the query if {@link
+   * CloudSpannerJdbcConnection#isDataBoostEnabled()} returns true.
+   */
+  ResultSet partitionQuery(PartitionOptions partitionOptions, QueryOption... options)
+      throws SQLException;
+
+  /**
+   * Executes the given partition of a query. The partition that should be executed must be set as a
+   * string parameter on this {@link PreparedStatement} using {@link #setString(int, String)}. The
+   * value should be a string that was returned by {@link #partitionQuery(PartitionOptions,
+   * QueryOption...)}.
+   */
+  ResultSet runPartition() throws SQLException;
+
+  /**
+   * Executes the given query as a partitioned query. The query will first be partitioned using the
+   * {@link #partitionQuery(PartitionOptions, QueryOption...)} method. Each of the partitions will
+   * then be executed in the background, and the results will be merged into a single result set.
+   *
+   * <p>This method will use {@link CloudSpannerJdbcConnection#getMaxPartitionedParallelism()}
+   * threads to execute the partitioned query. Set this variable to a higher/lower value to
+   * increase/decrease the degree of parallelism used for execution.
+   */
+  CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
+      PartitionOptions partitionOptions, QueryOption... options) throws SQLException;
+}
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcStatement.java
new file mode 100644
index 0000000..6115c0b
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcStatement.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright 2023 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 com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+/**
+ * This interface is implemented by {@link Statement}s that are created on Cloud Spanner JDBC
+ * connections.
+ */
+public interface CloudSpannerJdbcStatement extends Statement {
+
+  /**
+   * Partitions the given query, so it can be executed in parallel. This method returns a {@link
+   * ResultSet} with a string-representation of the partitions that were created. These strings can
+   * be used to execute a partition either on this connection or an any other connection (on this
+   * host or an any other host) by calling the method {@link #runPartition(String)}. This method
+   * will automatically enable data boost for the query if {@link
+   * CloudSpannerJdbcConnection#isDataBoostEnabled()} returns true.
+   */
+  ResultSet partitionQuery(String query, PartitionOptions partitionOptions, QueryOption... options)
+      throws SQLException;
+
+  /**
+   * Executes the given partition of a query. The encodedPartitionId should be a string that was
+   * returned by {@link #partitionQuery(String, PartitionOptions, QueryOption...)}.
+   */
+  ResultSet runPartition(String encodedPartitionId) throws SQLException;
+
+  /**
+   * Executes the given query as a partitioned query. The query will first be partitioned using the
+   * {@link #partitionQuery(String, PartitionOptions, QueryOption...)} method. Each of the
+   * partitions will then be executed in the background, and the results will be merged into a
+   * single result set.
+   *
+   * <p>This method will use {@link CloudSpannerJdbcConnection#getMaxPartitionedParallelism()}
+   * threads to execute the partitioned query. Set this variable to a higher/lower value to
+   * increase/decrease the degree of parallelism used for execution.
+   */
+  CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
+      String query, PartitionOptions partitionOptions, QueryOption... options) throws SQLException;
+}
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 6d6bd80..6031afe 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
@@ -25,6 +25,7 @@ import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.AutocommitDmlMode;
+import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.connection.TransactionMode;
@@ -44,6 +45,8 @@ import java.sql.Timestamp;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
 
 /** Jdbc Connection class for Google Cloud Spanner */
 class JdbcConnection extends AbstractJdbcConnection {
@@ -592,6 +595,72 @@ class JdbcConnection extends AbstractJdbcConnection {
     }
   }
 
+  /**
+   * Convenience method for calling a setter and translating any {@link SpannerException} to a
+   * {@link SQLException}.
+   */
+  private <T> void set(BiConsumer<Connection, T> setter, T value) throws SQLException {
+    checkClosed();
+    try {
+      setter.accept(getSpannerConnection(), value);
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    }
+  }
+
+  /**
+   * Convenience method for calling a getter and translating any {@link SpannerException} to a
+   * {@link SQLException}.
+   */
+  private <R> R get(Function<Connection, R> getter) throws SQLException {
+    checkClosed();
+    try {
+      return getter.apply(getSpannerConnection());
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    }
+  }
+
+  @Override
+  public void setDataBoostEnabled(boolean dataBoostEnabled) throws SQLException {
+    set(Connection::setDataBoostEnabled, dataBoostEnabled);
+  }
+
+  @Override
+  public boolean isDataBoostEnabled() throws SQLException {
+    return get(Connection::isDataBoostEnabled);
+  }
+
+  @Override
+  public void setAutoPartitionMode(boolean autoPartitionMode) throws SQLException {
+    set(Connection::setAutoPartitionMode, autoPartitionMode);
+  }
+
+  @Override
+  public boolean isAutoPartitionMode() throws SQLException {
+    return get(Connection::isAutoPartitionMode);
+  }
+
+  @Override
+  public void setMaxPartitions(int maxPartitions) throws SQLException {
+    set(Connection::setMaxPartitions, maxPartitions);
+  }
+
+  @Override
+  public int getMaxPartitions() throws SQLException {
+    return get(Connection::getMaxPartitions);
+  }
+
+  @Override
+  public void setMaxPartitionedParallelism(int maxThreads) throws SQLException {
+    set(Connection::setMaxPartitionedParallelism, maxThreads);
+  }
+
+  @Override
+  public int getMaxPartitionedParallelism() throws SQLException {
+    return get(Connection::getMaxPartitionedParallelism);
+  }
+
   @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/JdbcPartitionedQueryResultSet.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPartitionedQueryResultSet.java
new file mode 100644
index 0000000..1d5ef66
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPartitionedQueryResultSet.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2023 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 com.google.cloud.spanner.connection.PartitionedQueryResultSet;
+import com.google.common.base.Preconditions;
+import java.sql.Statement;
+
+/**
+ * {@link java.sql.ResultSet} implementation that is returned for queries that are executed with
+ * `RUN PARTITIONED QUERY ...`.
+ */
+class JdbcPartitionedQueryResultSet extends JdbcResultSet
+    implements CloudSpannerJdbcPartitionedQueryResultSet {
+  static JdbcPartitionedQueryResultSet of(
+      Statement statement, PartitionedQueryResultSet resultSet) {
+    return new JdbcPartitionedQueryResultSet(
+        Preconditions.checkNotNull(statement), Preconditions.checkNotNull(resultSet));
+  }
+
+  private final PartitionedQueryResultSet partitionedQueryResultSet;
+
+  private JdbcPartitionedQueryResultSet(Statement statement, PartitionedQueryResultSet resultSet) {
+    super(statement, resultSet);
+    this.partitionedQueryResultSet = resultSet;
+  }
+
+  @Override
+  public int getNumPartitions() {
+    return partitionedQueryResultSet.getNumPartitions();
+  }
+
+  @Override
+  public int getParallelism() {
+    return partitionedQueryResultSet.getParallelism();
+  }
+}
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
index 3ab302c..518807d 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
@@ -16,7 +16,9 @@
 
 package com.google.cloud.spanner.jdbc;
 
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 import com.google.cloud.spanner.ResultSets;
 import com.google.cloud.spanner.SpannerException;
@@ -33,10 +35,10 @@ import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 
 /** Implementation of {@link PreparedStatement} for Cloud Spanner. */
-class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
+class JdbcPreparedStatement extends AbstractJdbcPreparedStatement
+    implements CloudSpannerJdbcPreparedStatement {
   private static final char POS_PARAM_CHAR = '?';
   private final String sql;
-  private final String sqlWithoutComments;
   private final ParametersInfo parameters;
   private final ImmutableList<String> generatedKeysColumns;
 
@@ -46,9 +48,15 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
     super(connection);
     this.sql = sql;
     try {
-      this.sqlWithoutComments = parser.removeCommentsAndTrim(this.sql);
+      // The PostgreSQL parser allows comments to be present in the SQL string that is used to parse
+      // the query parameters.
+      String sqlForParameterExtraction =
+          getConnection().getDialect() == Dialect.POSTGRESQL
+              ? this.sql
+              : parser.removeCommentsAndTrim(this.sql);
       this.parameters =
-          parser.convertPositionalParametersToNamedParameters(POS_PARAM_CHAR, sqlWithoutComments);
+          parser.convertPositionalParametersToNamedParameters(
+              POS_PARAM_CHAR, sqlForParameterExtraction);
     } catch (SpannerException e) {
       throw JdbcSqlExceptionFactory.of(e);
     }
@@ -128,4 +136,37 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
       return rs.getMetaData();
     }
   }
+
+  @Override
+  public ResultSet partitionQuery(PartitionOptions partitionOptions, QueryOption... options)
+      throws SQLException {
+    return runWithStatementTimeout(
+        connection ->
+            JdbcResultSet.of(
+                this, connection.partitionQuery(createStatement(), partitionOptions, options)));
+  }
+
+  @Override
+  public ResultSet runPartition() throws SQLException {
+    return runWithStatementTimeout(
+        connection -> {
+          if (getParameters().getHighestIndex() < 1 || getParameters().getParameter(1) == null) {
+            throw JdbcSqlExceptionFactory.of(
+                "No query parameter has been set. runPartition() requires the partition ID to be set as a query parameter with index 1. Call PreparedStatement#setString(1, \"some-partition-id\") before calling runPartition().",
+                Code.FAILED_PRECONDITION);
+          }
+          String partitionId = getParameters().getParameter(1).toString();
+          return JdbcResultSet.of(this, connection.runPartition(partitionId));
+        });
+  }
+
+  @Override
+  public CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
+      PartitionOptions partitionOptions, QueryOption... options) throws SQLException {
+    return runWithStatementTimeout(
+        connection ->
+            JdbcPartitionedQueryResultSet.of(
+                this,
+                connection.runPartitionedQuery(createStatement(), partitionOptions, options)));
+  }
 }
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
index 94b4604..4152fa5 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
@@ -21,6 +21,7 @@ import com.google.cloud.spanner.Struct;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.Type.Code;
 import com.google.cloud.spanner.Value;
+import com.google.cloud.spanner.connection.PartitionedQueryResultSet;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import java.io.ByteArrayInputStream;
@@ -59,9 +60,13 @@ class JdbcResultSet extends AbstractJdbcResultSet {
   }
 
   static JdbcResultSet of(Statement statement, com.google.cloud.spanner.ResultSet resultSet) {
-    Preconditions.checkNotNull(statement);
-    Preconditions.checkNotNull(resultSet);
-    return new JdbcResultSet(statement, resultSet);
+    // Return a JDBC version of a PartitionedQueryResultSet if the Cloud Spanner Java client
+    // returned a PartitionedQueryResultSet.
+    if (resultSet instanceof PartitionedQueryResultSet) {
+      return JdbcPartitionedQueryResultSet.of(statement, (PartitionedQueryResultSet) resultSet);
+    }
+    return new JdbcResultSet(
+        Preconditions.checkNotNull(statement), Preconditions.checkNotNull(resultSet));
   }
 
   /**
@@ -123,7 +128,7 @@ class JdbcResultSet extends AbstractJdbcResultSet {
   private boolean nextCalledForMetaDataResult = false;
   private long currentRow = 0L;
 
-  private JdbcResultSet(Statement statement, com.google.cloud.spanner.ResultSet spanner) {
+  JdbcResultSet(Statement statement, com.google.cloud.spanner.ResultSet spanner) {
     super(spanner);
     this.statement = statement;
   }
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 6c6a45a..60cc3e5 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
@@ -19,6 +19,8 @@ package com.google.cloud.spanner.jdbc;
 import static com.google.cloud.spanner.jdbc.JdbcConnection.NO_GENERATED_KEY_COLUMNS;
 
 import com.google.cloud.spanner.Options;
+import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
 import com.google.cloud.spanner.ResultSets;
 import com.google.cloud.spanner.SpannerBatchUpdateException;
 import com.google.cloud.spanner.SpannerException;
@@ -40,7 +42,7 @@ import java.util.List;
 import java.util.stream.Collectors;
 
 /** Implementation of {@link java.sql.Statement} for Google Cloud Spanner. */
-class JdbcStatement extends AbstractJdbcStatement {
+class JdbcStatement extends AbstractJdbcStatement implements CloudSpannerJdbcStatement {
   static final ImmutableList<String> ALL_COLUMNS = ImmutableList.of("*");
 
   enum BatchType {
@@ -79,11 +81,7 @@ class JdbcStatement extends AbstractJdbcStatement {
 
   private int executeUpdate(String sql, ImmutableList<String> generatedKeysColumns)
       throws SQLException {
-    long result = executeLargeUpdate(sql, generatedKeysColumns);
-    if (result > Integer.MAX_VALUE) {
-      throw JdbcSqlExceptionFactory.of("update count too large: " + result, Code.OUT_OF_RANGE);
-    }
-    return (int) result;
+    return checkedCast(executeLargeUpdate(sql, generatedKeysColumns));
   }
 
   /**
@@ -300,11 +298,7 @@ class JdbcStatement extends AbstractJdbcStatement {
   @Override
   public int getUpdateCount() throws SQLException {
     checkClosed();
-    if (currentUpdateCount > Integer.MAX_VALUE) {
-      throw JdbcSqlExceptionFactory.of(
-          "update count too large: " + currentUpdateCount, Code.OUT_OF_RANGE);
-    }
-    return (int) currentUpdateCount;
+    return checkedCast(currentUpdateCount);
   }
 
   /**
@@ -481,12 +475,7 @@ class JdbcStatement extends AbstractJdbcStatement {
   int[] convertUpdateCounts(long[] updateCounts) throws SQLException {
     int[] res = new int[updateCounts.length];
     for (int index = 0; index < updateCounts.length; index++) {
-      if (updateCounts[index] > Integer.MAX_VALUE) {
-        throw JdbcSqlExceptionFactory.of(
-            String.format("Update count too large for int: %d", updateCounts[index]),
-            Code.OUT_OF_RANGE);
-      }
-      res[index] = (int) updateCounts[index];
+      res[index] = checkedCast(updateCounts[index]);
     }
     return res;
   }
@@ -598,4 +587,29 @@ class JdbcStatement extends AbstractJdbcStatement {
   static boolean isNullOrEmpty(String[] columnNames) {
     return columnNames == null || columnNames.length == 0;
   }
+
+  @Override
+  public ResultSet partitionQuery(
+      String query, PartitionOptions partitionOptions, QueryOption... options) throws SQLException {
+    return runWithStatementTimeout(
+        connection ->
+            JdbcResultSet.of(
+                this, connection.partitionQuery(Statement.of(query), partitionOptions, options)));
+  }
+
+  @Override
+  public ResultSet runPartition(String encodedPartitionId) throws SQLException {
+    return runWithStatementTimeout(
+        connection -> JdbcResultSet.of(this, connection.runPartition(encodedPartitionId)));
+  }
+
+  @Override
+  public CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
+      String query, PartitionOptions partitionOptions, QueryOption... options) throws SQLException {
+    return runWithStatementTimeout(
+        connection ->
+            JdbcPartitionedQueryResultSet.of(
+                this,
+                connection.runPartitionedQuery(Statement.of(query), partitionOptions, options)));
+  }
 }


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
index f74de66..65a54a8 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
@@ -17,39 +17,97 @@
 package com.google.cloud.spanner.jdbc;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.cloud.spanner.SessionPoolOptions;
 import com.google.cloud.spanner.connection.AbstractMockServerTest;
 import com.google.cloud.spanner.connection.ConnectionOptions;
+import com.google.cloud.spanner.connection.SpannerPool;
+import com.google.longrunning.Operation;
+import com.google.protobuf.Any;
+import com.google.protobuf.Empty;
 import com.google.protobuf.ListValue;
 import com.google.protobuf.Value;
 import com.google.rpc.Code;
+import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
 import com.google.spanner.v1.ResultSetMetadata;
+import com.google.spanner.v1.ResultSetStats;
 import com.google.spanner.v1.StructType;
 import com.google.spanner.v1.StructType.Field;
 import com.google.spanner.v1.Type;
 import com.google.spanner.v1.TypeCode;
 import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
 
 /**
  * Test class for verifying that the methods execute, executeQuery, and executeUpdate work as
- * intended.
+ * intended. It also verifies that they always also include any comments in the statement for the
+ * PostgreSQL dialect, as these may contain hints.
  */
-@RunWith(JUnit4.class)
+@RunWith(Parameterized.class)
 public class ExecuteMockServerTest extends AbstractMockServerTest {
-  private static final String QUERY = "select * from my_table";
+
+  @Parameters(name = "dialect = {0}")
+  public static Object[] parameters() {
+    return Dialect.values();
+  }
+
+  @Parameter public Dialect dialect;
+
+  private static final String DDL = "create table my_table";
+  private static final long LARGE_UPDATE_COUNT = 2L * Integer.MAX_VALUE;
+
+  private String query;
+  private String dml;
+  private String largeDml;
+  private String dmlReturning;
 
   @Before
   public void setupResults() {
+    query =
+        dialect == Dialect.POSTGRESQL
+            ? "/*@ lock_scanned_ranges = exclusive */ select * from my_table"
+            : "select * from my_table";
+    dml =
+        dialect == Dialect.POSTGRESQL
+            ? "/*@ lock_scanned_ranges = exclusive */ insert into my_table (id, value) values (1, 'One')"
+            : "insert into my_table (id, value) values (1, 'One')";
+    String DML_THEN_RETURN_ID =
+        dml + (dialect == Dialect.POSTGRESQL ? "\nRETURNING \"id\"" : "\nTHEN RETURN `id`");
+    largeDml =
+        dialect == Dialect.POSTGRESQL
+            ? "/*@ lock_scanned_ranges = exclusive */ update my_table set value='new value' where true"
+            : "update my_table set value='new value' where true";
+    String LARGE_DML_THEN_RETURN_ID =
+        largeDml + (dialect == Dialect.POSTGRESQL ? "\nRETURNING \"id\"" : "\nTHEN RETURN `id`");
+    dmlReturning =
+        dialect == Dialect.POSTGRESQL
+            ? "/*@ lock_scanned_ranges = exclusive */ insert into my_table (id, value) values (1, 'One') RETURNING *"
+            : "insert into my_table (id, value) values (1, 'One') THEN RETURN *";
+
+    // This forces a refresh of the Spanner instance that is used for a connection, which again is
+    // needed in order to refresh the dialect of the database.
+    SpannerPool.closeSpannerPool();
+    mockSpanner.putStatementResult(StatementResult.detectDialectResult(dialect));
     super.setupResults();
+
     com.google.spanner.v1.ResultSet resultSet =
         com.google.spanner.v1.ResultSet.newBuilder()
             .setMetadata(
@@ -74,8 +132,58 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
                     .addValues(Value.newBuilder().setStringValue("One").build())
                     .build())
             .build();
+    com.google.spanner.v1.ResultSet returnIdResultSet =
+        com.google.spanner.v1.ResultSet.newBuilder()
+            .setMetadata(
+                ResultSetMetadata.newBuilder()
+                    .setRowType(
+                        StructType.newBuilder()
+                            .addFields(
+                                Field.newBuilder()
+                                    .setType(Type.newBuilder().setCode(TypeCode.INT64).build())
+                                    .setName("id")
+                                    .build())
+                            .build())
+                    .build())
+            .addRows(
+                ListValue.newBuilder()
+                    .addValues(Value.newBuilder().setStringValue("1").build())
+                    .build())
+            .build();
+    mockSpanner.putStatementResult(
+        StatementResult.query(com.google.cloud.spanner.Statement.of(query), resultSet));
+    mockSpanner.putStatementResult(
+        StatementResult.update(com.google.cloud.spanner.Statement.of(dml), 1L));
     mockSpanner.putStatementResult(
-        StatementResult.query(com.google.cloud.spanner.Statement.of(QUERY), resultSet));
+        StatementResult.update(
+            com.google.cloud.spanner.Statement.of(largeDml), LARGE_UPDATE_COUNT));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            com.google.cloud.spanner.Statement.of(dmlReturning),
+            resultSet
+                .toBuilder()
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            com.google.cloud.spanner.Statement.of(DML_THEN_RETURN_ID),
+            returnIdResultSet
+                .toBuilder()
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            com.google.cloud.spanner.Statement.of(LARGE_DML_THEN_RETURN_ID),
+            returnIdResultSet
+                .toBuilder()
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(LARGE_UPDATE_COUNT).build())
+                .build()));
+    mockDatabaseAdmin.addResponse(
+        Operation.newBuilder()
+            .setDone(true)
+            .setResponse(Any.pack(Empty.getDefaultInstance()))
+            .setMetadata(Any.pack(UpdateDatabaseDdlMetadata.getDefaultInstance()))
+            .build());
   }
 
   private String createUrl() {
@@ -84,6 +192,501 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
         getPort(), "proj", "inst", "db");
   }
 
+  private Connection createConnection() throws SQLException {
+    return DriverManager.getConnection(createUrl());
+  }
+
+  @Test
+  public void testStatementExecuteQuery() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      try (ResultSet resultSet = statement.executeQuery(query)) {
+        verifyResultSet(resultSet);
+      }
+      try (ResultSet resultSet = statement.executeQuery(dmlReturning)) {
+        verifyResultSet(resultSet);
+      }
+      verifyException(() -> statement.executeQuery(dml));
+      verifyException(() -> statement.executeQuery(largeDml));
+      verifyException(() -> statement.executeQuery(DDL));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteUpdate() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeUpdate(dml));
+      assertEquals(0, statement.executeUpdate(DDL));
+      verifyOverflow(() -> statement.executeUpdate(largeDml));
+      verifyException(() -> statement.executeUpdate(query));
+      verifyException(() -> statement.executeUpdate(dmlReturning));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteUpdateReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      assertEquals(1, statement.executeUpdate(dml, Statement.NO_GENERATED_KEYS));
+      assertEquals(0, statement.executeUpdate(DDL, Statement.NO_GENERATED_KEYS));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeUpdate(query, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeUpdate(dmlReturning, Statement.NO_GENERATED_KEYS));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteUpdateReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeUpdate(dml, new String[] {"id"}));
+      assertEquals(0, statement.executeUpdate(DDL, new String[] {"id"}));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, new String[] {"id"}));
+      verifyException(
+          () -> statement.executeUpdate(query, new String[] {"id"}), Code.FAILED_PRECONDITION);
+      assertEquals(1, statement.executeUpdate(dmlReturning, new String[] {"id"}));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteUpdateReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeUpdate(dml, new int[] {1}));
+      assertEquals(0, statement.executeUpdate(DDL, new int[] {1}));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, new int[] {1}));
+      verifyException(() -> statement.executeUpdate(query, new int[] {1}));
+      verifyException(() -> statement.executeUpdate(dmlReturning, new int[] {1}));
+    }
+  }
+
+  @Test
+  public void testStatementLargeExecuteUpdate() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1L, statement.executeLargeUpdate(dml));
+      assertEquals(0L, statement.executeLargeUpdate(DDL));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml));
+      verifyException(() -> statement.executeLargeUpdate(query));
+      verifyException(() -> statement.executeLargeUpdate(dmlReturning));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteLargeUpdateReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      assertEquals(1, statement.executeLargeUpdate(dml, Statement.NO_GENERATED_KEYS));
+      assertEquals(0, statement.executeLargeUpdate(DDL, Statement.NO_GENERATED_KEYS));
+      assertEquals(
+          LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeLargeUpdate(query, Statement.NO_GENERATED_KEYS));
+      verifyException(
+          () -> statement.executeLargeUpdate(dmlReturning, Statement.NO_GENERATED_KEYS));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteLargeUpdateReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeLargeUpdate(dml, new String[] {"id"}));
+      assertEquals(0, statement.executeLargeUpdate(DDL, new String[] {"id"}));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, new String[] {"id"}));
+      verifyException(
+          () -> statement.executeLargeUpdate(query, new String[] {"id"}), Code.FAILED_PRECONDITION);
+      assertEquals(1L, statement.executeLargeUpdate(dmlReturning, new String[] {"id"}));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteLargeUpdateReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeLargeUpdate(dml, new int[] {1}));
+      assertEquals(0, statement.executeLargeUpdate(DDL, new int[] {1}));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, new int[] {1}));
+      verifyException(() -> statement.executeLargeUpdate(query, new int[] {1}));
+      verifyException(() -> statement.executeLargeUpdate(dmlReturning, new int[] {1}));
+    }
+  }
+
+  @Test
+  public void testStatementExecute() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      verifyUpdateCount(statement, () -> statement.execute(dml), 1L);
+      verifyUpdateCount(statement, () -> statement.execute(largeDml), LARGE_UPDATE_COUNT);
+      verifyUpdateCount(statement, () -> statement.execute(DDL), Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(query));
+      verifyResultSet(statement, () -> statement.execute(dmlReturning));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      verifyUpdateCount(statement, () -> statement.execute(dml, Statement.NO_GENERATED_KEYS), 1L);
+      verifyUpdateCount(
+          statement,
+          () -> statement.execute(largeDml, Statement.NO_GENERATED_KEYS),
+          LARGE_UPDATE_COUNT);
+      verifyUpdateCount(
+          statement,
+          () -> statement.execute(DDL, Statement.NO_GENERATED_KEYS),
+          Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(query, Statement.NO_GENERATED_KEYS));
+      verifyResultSet(
+          statement, () -> statement.execute(dmlReturning, Statement.NO_GENERATED_KEYS));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      verifyUpdateCount(statement, () -> statement.execute(dml, new String[] {"id"}), 1L);
+      verifyUpdateCount(
+          statement, () -> statement.execute(largeDml, new String[] {"id"}), LARGE_UPDATE_COUNT);
+      verifyUpdateCount(
+          statement, () -> statement.execute(DDL, new String[] {"id"}), Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(query, new String[] {"id"}));
+      verifyResultSet(statement, () -> statement.execute(dmlReturning, new String[] {"id"}));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      verifyUpdateCount(statement, () -> statement.execute(dml, new int[] {1}), 1L);
+      verifyUpdateCount(
+          statement, () -> statement.execute(largeDml, new int[] {1}), LARGE_UPDATE_COUNT);
+      verifyUpdateCount(
+          statement, () -> statement.execute(DDL, new int[] {1}), Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(query, new int[] {1}));
+      verifyResultSet(statement, () -> statement.execute(dmlReturning, new int[] {1}));
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteQuery() throws SQLException {
+    try (Connection connection = createConnection()) {
+      try (ResultSet resultSet = connection.prepareStatement(query).executeQuery()) {
+        verifyResultSet(resultSet);
+      }
+      try (ResultSet resultSet = connection.prepareStatement(dmlReturning).executeQuery()) {
+        verifyResultSet(resultSet);
+      }
+      verifyException(() -> connection.prepareStatement(dml).executeQuery());
+      verifyException(() -> connection.prepareStatement(largeDml).executeQuery());
+      verifyException(() -> connection.prepareStatement(DDL).executeQuery());
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdate() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml).executeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL).executeUpdate());
+      verifyOverflow(() -> connection.prepareStatement(largeDml).executeUpdate());
+      verifyException(() -> connection.prepareStatement(query).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(dmlReturning).executeUpdate(), Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdateReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      assertEquals(
+          1, connection.prepareStatement(dml, Statement.NO_GENERATED_KEYS).executeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL,
+      // Statement.NO_GENERATED_KEYS).executeUpdate());
+      verifyOverflow(
+          () -> connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(query, Statement.NO_GENERATED_KEYS).executeUpdate());
+      verifyException(
+          () ->
+              connection
+                  .prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS)
+                  .executeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdateReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml, new String[] {"id"}).executeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL, new String[] {"id"}).executeUpdate());
+      verifyOverflow(
+          () -> connection.prepareStatement(largeDml, new String[] {"id"}).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(query, new String[] {"id"}).executeUpdate(),
+          Code.FAILED_PRECONDITION);
+      assertEquals(
+          1, connection.prepareStatement(dmlReturning, new String[] {"id"}).executeUpdate());
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdateReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml, new int[] {1}).executeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL, new int[] {1}).executeUpdate());
+      verifyOverflow(() -> connection.prepareStatement(largeDml, new int[] {1}).executeUpdate());
+      verifyException(() -> connection.prepareStatement(query, new int[] {1}).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(dmlReturning, new int[] {1}).executeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementLargeExecuteUpdate() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1L, connection.prepareStatement(dml).executeLargeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0L, connection.prepareStatement(DDL).executeLargeUpdate());
+      assertEquals(LARGE_UPDATE_COUNT, connection.prepareStatement(largeDml).executeLargeUpdate());
+      verifyException(() -> connection.prepareStatement(query).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(dmlReturning).executeLargeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteLargeUpdateReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      assertEquals(
+          1, connection.prepareStatement(dml, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL,
+      // Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      assertEquals(
+          LARGE_UPDATE_COUNT,
+          connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      verifyException(
+          () ->
+              connection.prepareStatement(query, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      verifyException(
+          () ->
+              connection
+                  .prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS)
+                  .executeLargeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteLargeUpdateReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml, new String[] {"id"}).executeLargeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL, new String[]
+      // {"id"}).executeLargeUpdate());
+      assertEquals(
+          LARGE_UPDATE_COUNT,
+          connection.prepareStatement(largeDml, new String[] {"id"}).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(query, new String[] {"id"}).executeLargeUpdate(),
+          Code.FAILED_PRECONDITION);
+      assertEquals(
+          1L, connection.prepareStatement(dmlReturning, new String[] {"id"}).executeLargeUpdate());
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteLargeUpdateReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml, new int[] {1}).executeLargeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL, new int[] {1}).executeLargeUpdate());
+      assertEquals(
+          LARGE_UPDATE_COUNT,
+          connection.prepareStatement(largeDml, new int[] {1}).executeLargeUpdate());
+      verifyException(() -> connection.prepareStatement(query, new int[] {1}).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(dmlReturning, new int[] {1}).executeLargeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecute() throws SQLException {
+    try (Connection connection = createConnection()) {
+      verifyPreparedUpdateCount(connection.prepareStatement(dml), PreparedStatement::execute, 1L);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(largeDml), PreparedStatement::execute, LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL), PreparedStatement::execute, Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(connection.prepareStatement(query), PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(dmlReturning), PreparedStatement::execute);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(dml, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute,
+          1L);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute,
+          LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute,
+          Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(
+          connection.prepareStatement(query, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection()) {
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(dml, new String[] {"id"}), PreparedStatement::execute, 1L);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(largeDml, new String[] {"id"}),
+          PreparedStatement::execute,
+          LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL, new String[] {"id"}),
+          PreparedStatement::execute,
+          Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(
+          connection.prepareStatement(query, new String[] {"id"}), PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(dmlReturning, new String[] {"id"}),
+          PreparedStatement::execute);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection()) {
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(dml, new int[] {1}), PreparedStatement::execute, 1L);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(largeDml, new int[] {1}),
+          PreparedStatement::execute,
+          LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL, new int[] {1}),
+          PreparedStatement::execute,
+          Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(
+          connection.prepareStatement(query, new int[] {1}), PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(dmlReturning, new int[] {1}), PreparedStatement::execute);
+    }
+  }
+
+  private void verifyResultSet(ResultSet resultSet) throws SQLException {
+    assertNotNull(resultSet.getMetaData());
+    assertEquals(2, resultSet.getMetaData().getColumnCount());
+    assertTrue(resultSet.next());
+    assertFalse(resultSet.next());
+  }
+
+  private void verifyException(ThrowingRunnable runnable) {
+    verifyException(runnable, Code.INVALID_ARGUMENT);
+  }
+
+  private void verifyOverflow(ThrowingRunnable runnable) {
+    verifyException(runnable, Code.OUT_OF_RANGE);
+  }
+
+  private void verifyException(ThrowingRunnable runnable, Code code) {
+    SQLException exception = assertThrows(SQLException.class, runnable);
+    assertTrue(exception instanceof JdbcSqlException);
+    JdbcSqlException sqlException = (JdbcSqlException) exception;
+    assertEquals(code, sqlException.getCode());
+  }
+
+  interface ThrowingFunction<T> {
+    T apply() throws SQLException;
+  }
+
+  interface ThrowingPreparedFunction<T> {
+    T apply(PreparedStatement statement) throws SQLException;
+  }
+
+  private void verifyPreparedUpdateCount(
+      PreparedStatement statement, ThrowingPreparedFunction<Boolean> function, long updateCount)
+      throws SQLException {
+    verifyUpdateCount(statement, () -> function.apply(statement), updateCount);
+  }
+
+  private void verifyUpdateCount(
+      Statement statement, ThrowingFunction<Boolean> function, long updateCount)
+      throws SQLException {
+    assertFalse(function.apply());
+    if (updateCount > Integer.MAX_VALUE) {
+      verifyOverflow(statement::getUpdateCount);
+    } else {
+      assertEquals((int) updateCount, statement.getUpdateCount());
+    }
+    assertEquals(updateCount, statement.getLargeUpdateCount());
+    assertNull(statement.getResultSet());
+    assertFalse(statement.getMoreResults());
+    assertEquals(-1, statement.getUpdateCount());
+    assertEquals(-1L, statement.getLargeUpdateCount());
+  }
+
+  private void verifyPreparedResultSet(
+      PreparedStatement statement, ThrowingPreparedFunction<Boolean> function) throws SQLException {
+    verifyResultSet(statement, () -> function.apply(statement));
+  }
+
+  private void verifyResultSet(Statement statement, ThrowingFunction<Boolean> function)
+      throws SQLException {
+    assertTrue(function.apply());
+    assertEquals(-1, statement.getUpdateCount());
+    assertEquals(-1L, statement.getLargeUpdateCount());
+    assertNotNull(statement.getResultSet());
+    try (ResultSet resultSet = statement.getResultSet()) {
+      assertTrue(resultSet.next());
+      assertFalse(resultSet.next());
+    }
+
+    assertFalse(statement.getMoreResults());
+    assertEquals(-1, statement.getUpdateCount());
+    assertEquals(-1L, statement.getLargeUpdateCount());
+  }
+
   @Test
   public void testInvalidExecuteUpdate_shouldNotLeakSession() throws SQLException {
     int maxSessions = 1;
@@ -103,7 +706,7 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
       for (int i = 0; i < (maxSessions + 1); i++) {
         SQLException exception =
             assertThrows(
-                SQLException.class, () -> connection.createStatement().executeUpdate(QUERY));
+                SQLException.class, () -> connection.createStatement().executeUpdate(query));
         assertTrue(exception instanceof JdbcSqlException);
         JdbcSqlException jdbcSqlException = (JdbcSqlException) exception;
         // This would be RESOURCE_EXHAUSTED if the query leaked a session.
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/PartitionedQueryMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/PartitionedQueryMockServerTest.java
new file mode 100644
index 0000000..d6edf99
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/jdbc/PartitionedQueryMockServerTest.java
@@ -0,0 +1,486 @@
+/*
+ * Copyright 2023 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.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.PartitionOptions;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.cloud.spanner.connection.SpannerPool;
+import com.google.protobuf.ByteString;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.PartitionQueryRequest;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class PartitionedQueryMockServerTest extends AbstractMockServerTest {
+  @Parameter public Dialect dialect;
+
+  @Parameters(name = "dialect = {0}")
+  public static Object[] data() {
+    return Dialect.values();
+  }
+
+  @Before
+  public void setupDialect() {
+    mockSpanner.putStatementResult(StatementResult.detectDialectResult(this.dialect));
+  }
+
+  @After
+  public void clearRequests() {
+    mockSpanner.clearRequests();
+    SpannerPool.closeSpannerPool();
+  }
+
+  private String createUrl() {
+    return String.format(
+        "jdbc:cloudspanner://localhost:%d/projects/%s/instances/%s/databases/%s?usePlainText=true",
+        getPort(), "proj", "inst", "db");
+  }
+
+  private Connection createConnection() throws SQLException {
+    return DriverManager.getConnection(createUrl());
+  }
+
+  @Test
+  public void testPartitionedQueryUsingSql() throws SQLException {
+    int numRows = 5;
+    int maxPartitions = 4;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows);
+    Statement statement =
+        Statement.newBuilder(
+                String.format(
+                    "select * from my_table where active=%s",
+                    dialect == Dialect.POSTGRESQL ? "$1" : "@p1"))
+            .bind("p1")
+            .to(true)
+            .build();
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+    String prefix = dialect == Dialect.POSTGRESQL ? "spanner." : "";
+
+    try (Connection connection = createConnection()) {
+      // This will automatically enable Data Boost for any partitioned query that is executed on
+      // this connection. The property is ignored for any query that is not a partitioned query.
+      connection.createStatement().execute(String.format("set %sdata_boost_enabled=true", prefix));
+      // Sets the maximum number of partitions that should be used by Cloud Spanner.
+      // This is just a hint that can be ignored by Cloud Spanner, but the mock server that is used
+      // for testing respects this hint.
+      connection
+          .createStatement()
+          .execute(String.format("set %smax_partitions=%d", prefix, maxPartitions));
+
+      List<String> partitionIds = new ArrayList<>();
+      try (PreparedStatement partitionStatement =
+          connection.prepareStatement("partition select * from my_table where active=?")) {
+        partitionStatement.setBoolean(1, true);
+        try (ResultSet partitions = partitionStatement.executeQuery()) {
+          assertNotNull(partitions.getMetaData());
+          while (partitions.next()) {
+            partitionIds.add(partitions.getString(1));
+          }
+        }
+      }
+
+      for (String partitionId : partitionIds) {
+        // Execute `run partition '<partition-id>'` to get the results of a single partition.
+        // This can be done from any JDBC connection, including from other hosts.
+        try (ResultSet resultSet =
+            connection
+                .createStatement()
+                .executeQuery(String.format("run partition '%s'", partitionId))) {
+          assertNotNull(resultSet.getMetaData());
+          int rowCount = 0;
+          while (resultSet.next()) {
+            rowCount++;
+          }
+          assertEquals(numRows, rowCount);
+        }
+      }
+
+      for (String partitionId : partitionIds) {
+        // The partition ID may also be set as a query parameter.
+        try (PreparedStatement runStatement = connection.prepareStatement("run partition")) {
+          runStatement.setString(1, partitionId);
+          int rowCount = 0;
+          try (ResultSet resultSet = runStatement.executeQuery()) {
+            assertNotNull(resultSet.getMetaData());
+            while (resultSet.next()) {
+              rowCount++;
+            }
+            assertEquals(numRows, rowCount);
+          }
+        }
+      }
+
+      // You can also run a query directly as a partitioned query on a single JDBC connection. This
+      // will not give you the full benefits of a partitioned query, as the entire query is handled
+      // by a single host and is returned as a single result set. The result set uses multiple
+      // threads internally to execute the separate partitions.
+      // This gives users a simple way to access the Data Boost feature that should be accessible
+      // from most generic frameworks and tools that work with JDBC.
+      try (PreparedStatement partitionStatement =
+          connection.prepareStatement(
+              "run partitioned query select * from my_table where active=?")) {
+        partitionStatement.setBoolean(1, true);
+        try (ResultSet results = partitionStatement.executeQuery()) {
+          assertNotNull(results.getMetaData());
+          assertEquals(18, results.getMetaData().getColumnCount());
+          int rowCount = 0;
+          while (results.next()) {
+            rowCount++;
+          }
+          // The mock server is not smart enough to actually partition the query and only return
+          // a fraction of the rows per partition. The total row count will therefore be equal to
+          // the number of partitions multiplied by the number of rows.
+          assertEquals(numRows * maxPartitions, rowCount);
+
+          // Partitioned queries return a result set with some additional metadata that can be
+          // inspected to determine the number of partitions and the degree of parallelism that the
+          // query used.
+          CloudSpannerJdbcPartitionedQueryResultSet partitionedQueryResultSet =
+              results.unwrap(CloudSpannerJdbcPartitionedQueryResultSet.class);
+          assertEquals(maxPartitions, partitionedQueryResultSet.getNumPartitions());
+          assertEquals(1, partitionedQueryResultSet.getParallelism());
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testPartitionedQueryStatement() throws SQLException {
+    int numRows = 5;
+    int maxPartitions = 4;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows);
+    Statement statement = Statement.of("select * from my_table where active=true");
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+
+    try (Connection connection = createConnection()) {
+      CloudSpannerJdbcConnection cloudSpannerJdbcConnection =
+          connection.unwrap(CloudSpannerJdbcConnection.class);
+      // This will automatically enable Data Boost for any partitioned query that is executed on
+      // this connection. The property is ignored for any query that is not a partitioned query.
+      cloudSpannerJdbcConnection.setDataBoostEnabled(true);
+      // Sets the maximum number of partitions that should be used by Cloud Spanner.
+      // This is just a hint that can be ignored by Cloud Spanner, but the mock server that is used
+      // for testing respects this hint.
+      cloudSpannerJdbcConnection.setMaxPartitions(maxPartitions);
+
+      List<String> partitionIds = new ArrayList<>();
+      try (ResultSet partitions =
+          connection
+              .createStatement()
+              .unwrap(CloudSpannerJdbcStatement.class)
+              .partitionQuery(
+                  "select * from my_table where active=true",
+                  PartitionOptions.getDefaultInstance())) {
+        assertNotNull(partitions.getMetaData());
+        while (partitions.next()) {
+          partitionIds.add(partitions.getString(1));
+        }
+      }
+
+      for (String partitionId : partitionIds) {
+        // Execute `run partition '<partition-id>'` to get the results of a single partition.
+        // This can be done from any JDBC connection, including from other hosts.
+        try (ResultSet resultSet =
+            connection
+                .createStatement()
+                .unwrap(CloudSpannerJdbcStatement.class)
+                .runPartition(partitionId)) {
+          assertNotNull(resultSet.getMetaData());
+          int rowCount = 0;
+          while (resultSet.next()) {
+            rowCount++;
+          }
+          assertEquals(numRows, rowCount);
+        }
+      }
+
+      // You can also run a query directly as a partitioned query on a single JDBC connection. This
+      // will not give you the full benefits of a partitioned query, as the entire query is handled
+      // by a single host and is returned as a single result set. The result set uses multiple
+      // threads internally to execute the separate partitions.
+      // This gives users a simple way to access the Data Boost feature that should be accessible
+      // from most generic frameworks and tools that work with JDBC.
+      try (CloudSpannerJdbcPartitionedQueryResultSet results =
+          connection
+              .createStatement()
+              .unwrap(CloudSpannerJdbcStatement.class)
+              .runPartitionedQuery(
+                  "select * from my_table where active=true",
+                  PartitionOptions.getDefaultInstance())) {
+        assertNotNull(results.getMetaData());
+        int rowCount = 0;
+        while (results.next()) {
+          rowCount++;
+        }
+        // The mock server is not smart enough to actually partition the query and only return
+        // a fraction of the rows per partition. The total row count will therefore be equal to
+        // the number of partitions multiplied by the number of rows.
+        assertEquals(numRows * maxPartitions, rowCount);
+
+        // Partitioned queries return a result set with some additional metadata that can be
+        // inspected to determine the number of partitions and the degree of parallelism that the
+        // query used.
+        assertEquals(maxPartitions, results.getNumPartitions());
+        assertEquals(1, results.getParallelism());
+      }
+    }
+  }
+
+  @Test
+  public void testPartitionedQueryPreparedStatement() throws SQLException {
+    int numRows = 5;
+    int maxPartitions = 4;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows);
+    Statement statement =
+        Statement.newBuilder(
+                String.format(
+                    "select * from my_table where active=%s",
+                    dialect == Dialect.POSTGRESQL ? "$1" : "@p1"))
+            .bind("p1")
+            .to(true)
+            .build();
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+
+    try (Connection connection = createConnection()) {
+      CloudSpannerJdbcConnection cloudSpannerJdbcConnection =
+          connection.unwrap(CloudSpannerJdbcConnection.class);
+      // This will automatically enable Data Boost for any partitioned query that is executed on
+      // this connection. The property is ignored for any query that is not a partitioned query.
+      cloudSpannerJdbcConnection.setDataBoostEnabled(true);
+      // Sets the maximum number of partitions that should be used by Cloud Spanner.
+      // This is just a hint that can be ignored by Cloud Spanner, but the mock server that is used
+      // for testing respects this hint.
+      cloudSpannerJdbcConnection.setMaxPartitions(maxPartitions);
+
+      List<String> partitionIds = new ArrayList<>();
+      try (PreparedStatement partitionStatement =
+          connection.prepareStatement("select * from my_table where active=?")) {
+        partitionStatement.setBoolean(1, true);
+        try (ResultSet partitions =
+            partitionStatement
+                .unwrap(CloudSpannerJdbcPreparedStatement.class)
+                .partitionQuery(PartitionOptions.getDefaultInstance())) {
+          assertNotNull(partitions.getMetaData());
+          while (partitions.next()) {
+            partitionIds.add(partitions.getString(1));
+          }
+        }
+      }
+
+      for (String partitionId : partitionIds) {
+        // The partition ID can also be set as a query parameter.
+        // The statement text in this case is ignored and can be an empty string.
+        try (PreparedStatement runStatement = connection.prepareStatement("")) {
+          runStatement.setString(1, partitionId);
+          int rowCount = 0;
+          try (ResultSet resultSet =
+              runStatement.unwrap(CloudSpannerJdbcPreparedStatement.class).runPartition()) {
+            assertNotNull(resultSet.getMetaData());
+            while (resultSet.next()) {
+              rowCount++;
+            }
+            assertEquals(numRows, rowCount);
+          }
+        }
+      }
+
+      // You can also run a query directly as a partitioned query on a single JDBC connection. This
+      // will not give you the full benefits of a partitioned query, as the entire query is handled
+      // by a single host and is returned as a single result set. The result set uses multiple
+      // threads internally to execute the separate partitions.
+      // This gives users a simple way to access the Data Boost feature that should be accessible
+      // from most generic frameworks and tools that work with JDBC.
+      try (PreparedStatement preparedStatement =
+          connection.prepareStatement("select * from my_table where active=?")) {
+        preparedStatement.setBoolean(1, true);
+        try (CloudSpannerJdbcPartitionedQueryResultSet results =
+            preparedStatement
+                .unwrap(CloudSpannerJdbcPreparedStatement.class)
+                .runPartitionedQuery(PartitionOptions.getDefaultInstance())) {
+          assertNotNull(results.getMetaData());
+          int rowCount = 0;
+          while (results.next()) {
+            rowCount++;
+          }
+          // The mock server is not smart enough to actually partition the query and only return
+          // a fraction of the rows per partition. The total row count will therefore be equal to
+          // the number of partitions multiplied by the number of rows.
+          assertEquals(numRows * maxPartitions, rowCount);
+
+          // Partitioned queries return a result set with some additional metadata that can be
+          // inspected to determine the number of partitions and the degree of parallelism that the
+          // query used.
+          assertEquals(maxPartitions, results.getNumPartitions());
+          assertEquals(1, results.getParallelism());
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testAutoPartitionMode() throws SQLException {
+    int numRows = 5;
+    int maxPartitions = 4;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows);
+    Statement statement = Statement.of("select * from my_table where active=true");
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+
+    try (Connection connection = createConnection()) {
+      CloudSpannerJdbcConnection cloudSpannerJdbcConnection =
+          connection.unwrap(CloudSpannerJdbcConnection.class);
+      // This will automatically enable Data Boost for any partitioned query that is executed on
+      // this connection. The property is ignored for any query that is not a partitioned query.
+      cloudSpannerJdbcConnection.setDataBoostEnabled(true);
+      // Sets the maximum number of partitions that should be used by Cloud Spanner.
+      // This is just a hint that can be ignored by Cloud Spanner, but the mock server that is used
+      // for testing respects this hint.
+      cloudSpannerJdbcConnection.setMaxPartitions(maxPartitions);
+      cloudSpannerJdbcConnection.setAutoPartitionMode(true);
+
+      try (ResultSet results =
+          connection.createStatement().executeQuery("select * from my_table where active=true")) {
+        assertNotNull(results.getMetaData());
+        assertEquals(18, results.getMetaData().getColumnCount());
+        int rowCount = 0;
+        while (results.next()) {
+          rowCount++;
+        }
+        // The mock server is not smart enough to actually partition the query and only return
+        // a fraction of the rows per partition. The total row count will therefore be equal to
+        // the number of partitions multiplied by the number of rows.
+        assertEquals(numRows * maxPartitions, rowCount);
+        assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
+
+        // Partitioned queries return a result set with some additional metadata that can be
+        // inspected to determine the number of partitions and the degree of parallelism that the
+        // query used.
+        assertEquals(
+            maxPartitions, results.unwrap(JdbcPartitionedQueryResultSet.class).getNumPartitions());
+        assertEquals(1, results.unwrap(JdbcPartitionedQueryResultSet.class).getParallelism());
+
+        // Verify that we can run metadata queries in auto_partition_mode.
+        // Just add a random result for the table metadata query. We don't care about the result,
+        // only about the fact that it should be allowed, and that it is executed in normal mode.
+        if (dialect == Dialect.GOOGLE_STANDARD_SQL) {
+          mockSpanner.putPartialStatementResult(
+              StatementResult.query(
+                  Statement.of(
+                      "SELECT TABLE_CATALOG AS TABLE_CAT, TABLE_SCHEMA AS TABLE_SCHEM, TABLE_NAME,\n"
+                          + "       CASE WHEN TABLE_TYPE = 'BASE TABLE' THEN 'TABLE' ELSE TABLE_TYPE END AS TABLE_TYPE,\n"
+                          + "       NULL AS REMARKS, NULL AS TYPE_CAT, NULL AS TYPE_SCHEM, NULL AS TYPE_NAME,\n"
+                          + "       NULL AS SELF_REFERENCING_COL_NAME, NULL AS REF_GENERATION\n"
+                          + "FROM INFORMATION_SCHEMA.TABLES AS T"),
+                  SELECT_COUNT_RESULTSET_BEFORE_INSERT));
+        } else {
+          mockSpanner.putPartialStatementResult(
+              StatementResult.query(
+                  Statement.of(
+                      "SELECT TABLE_CATALOG AS \"TABLE_CAT\", TABLE_SCHEMA AS \"TABLE_SCHEM\", TABLE_NAME AS \"TABLE_NAME\",\n"
+                          + "       CASE WHEN TABLE_TYPE = 'BASE TABLE' THEN 'TABLE' ELSE TABLE_TYPE END AS \"TABLE_TYPE\",\n"
+                          + "       NULL AS \"REMARKS\", NULL AS \"TYPE_CAT\", NULL AS \"TYPE_SCHEM\", NULL AS \"TYPE_NAME\",\n"
+                          + "       NULL AS \"SELF_REFERENCING_COL_NAME\", NULL AS \"REF_GENERATION\"\n"
+                          + "FROM INFORMATION_SCHEMA.TABLES AS T"),
+                  SELECT_COUNT_RESULTSET_BEFORE_INSERT));
+        }
+        try (ResultSet tables = connection.getMetaData().getTables(null, null, null, null)) {
+          assertTrue(tables.next());
+          assertEquals(0, tables.getInt(1));
+          assertFalse(tables.next());
+        }
+        assertEquals(
+            1,
+            mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).stream()
+                .filter(req -> req.getSql().contains("FROM INFORMATION_SCHEMA.TABLES AS T"))
+                .count());
+        ExecuteSqlRequest request =
+            mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).stream()
+                .filter(req -> req.getSql().contains("FROM INFORMATION_SCHEMA.TABLES AS T"))
+                .findFirst()
+                .orElse(ExecuteSqlRequest.getDefaultInstance());
+        assertTrue(request.hasTransaction());
+        assertTrue(request.getTransaction().hasSingleUse());
+        assertTrue(request.getTransaction().getSingleUse().hasReadOnly());
+        assertTrue(request.getTransaction().getSingleUse().getReadOnly().hasStrong());
+        assertEquals(ByteString.EMPTY, request.getPartitionToken());
+        assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
+      }
+    }
+  }
+
+  @Test
+  public void testAutoPartitionModeEmptyResult() throws SQLException {
+    int numRows = 0;
+    int maxPartitions = 1;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows);
+    Statement statement = Statement.of("select * from my_table where active=true");
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+
+    try (Connection connection = createConnection()) {
+      CloudSpannerJdbcConnection cloudSpannerJdbcConnection =
+          connection.unwrap(CloudSpannerJdbcConnection.class);
+      // This will automatically enable Data Boost for any partitioned query that is executed on
+      // this connection. The property is ignored for any query that is not a partitioned query.
+      cloudSpannerJdbcConnection.setDataBoostEnabled(true);
+      // Sets the maximum number of partitions that should be used by Cloud Spanner.
+      // This is just a hint that can be ignored by Cloud Spanner, but the mock server that is used
+      // for testing respects this hint.
+      cloudSpannerJdbcConnection.setMaxPartitions(maxPartitions);
+      cloudSpannerJdbcConnection.setAutoPartitionMode(true);
+
+      try (ResultSet results =
+          connection.createStatement().executeQuery("select * from my_table where active=true")) {
+        assertNotNull(results.getMetaData());
+        assertEquals(18, results.getMetaData().getColumnCount());
+        int rowCount = 0;
+        while (results.next()) {
+          rowCount++;
+        }
+        assertEquals(0, rowCount);
+        assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
+
+        // Partitioned queries return a result set with some additional metadata that can be
+        // inspected to determine the number of partitions and the degree of parallelism that the
+        // query used.
+        assertEquals(
+            maxPartitions, results.unwrap(JdbcPartitionedQueryResultSet.class).getNumPartitions());
+        assertEquals(1, results.unwrap(JdbcPartitionedQueryResultSet.class).getParallelism());
+      }
+    }
+  }
+}