Repository: googleapis/java-spanner-jdbc
Pull request #1300 (2023-07-30 13:15:51)

feat: support partitioned queries
Adds support for partitioned queries to the JDBC driver. Partitioned queries can be executed either using SQL statements or specific Cloud Spanner methods that are exposed by the specific Cloud Spanner JDBC interfaces.


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/snippets/pom.xml b/samples/snippets/pom.xml
index 0012684..a00b7a5 100644
--- a/samples/snippets/pom.xml
+++ b/samples/snippets/pom.xml
@@ -30,7 +30,7 @@
       <dependency>
         <groupId>com.google.cloud</groupId>
         <artifactId>libraries-bom</artifactId>
-        <version>26.22.0</version>
+        <version>26.23.0</version>
         <type>pom</type>
         <scope>import</scope>
       </dependency>
diff --git a/samples/spring-data-jdbc/pom.xml b/samples/spring-data-jdbc/pom.xml
index d0862d2..ee5a443 100644
--- a/samples/spring-data-jdbc/pom.xml
+++ b/samples/spring-data-jdbc/pom.xml
@@ -30,7 +30,7 @@
       <dependency>
         <groupId>com.google.cloud</groupId>
         <artifactId>libraries-bom</artifactId>
-        <version>26.22.0</version>
+        <version>26.23.0</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..c2e0c1d 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.jdbc;
 
 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,7 +34,8 @@ 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;
@@ -128,4 +130,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..b62449e 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
@@ -17,6 +17,9 @@
 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;
 
@@ -24,18 +27,28 @@ 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.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;
 
@@ -46,6 +59,14 @@ import org.junit.runners.JUnit4;
 @RunWith(JUnit4.class)
 public class ExecuteMockServerTest extends AbstractMockServerTest {
   private static final String QUERY = "select * from my_table";
+  private static final String DML = "insert into my_table (id, value) values (1, 'One')";
+  private static final String DML_THEN_RETURN_ID = DML + "\nTHEN RETURN `id`";
+  private static final String LARGE_DML = "update my_table set value='new value' where true";
+  private static final String LARGE_DML_THEN_RETURN_ID = LARGE_DML + "\nTHEN RETURN `id`";
+  private static final String DML_RETURNING =
+      "insert into my_table (id, value) values (1, 'One') THEN RETURN *";
+  private static final String DDL = "create table my_table";
+  private static final long LARGE_UPDATE_COUNT = 2L * Integer.MAX_VALUE;
 
   @Before
   public void setupResults() {
@@ -74,8 +95,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.update(
+            com.google.cloud.spanner.Statement.of(LARGE_DML), LARGE_UPDATE_COUNT));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            com.google.cloud.spanner.Statement.of(DML_RETURNING),
+            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 +155,503 @@ 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(DML_RETURNING)) {
+        verifyResultSet(resultSet);
+      }
+      verifyException(() -> statement.executeQuery(DML));
+      verifyException(() -> statement.executeQuery(LARGE_DML));
+      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(LARGE_DML));
+      verifyException(() -> statement.executeUpdate(QUERY));
+      verifyException(() -> statement.executeUpdate(DML_RETURNING));
+    }
+  }
+
+  @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(LARGE_DML, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeUpdate(QUERY, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeUpdate(DML_RETURNING, 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(LARGE_DML, new String[] {"id"}));
+      verifyException(
+          () -> statement.executeUpdate(QUERY, new String[] {"id"}), Code.FAILED_PRECONDITION);
+      assertEquals(1, statement.executeUpdate(DML_RETURNING, 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(LARGE_DML, new int[] {1}));
+      verifyException(() -> statement.executeUpdate(QUERY, new int[] {1}));
+      verifyException(() -> statement.executeUpdate(DML_RETURNING, 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(LARGE_DML));
+      verifyException(() -> statement.executeLargeUpdate(QUERY));
+      verifyException(() -> statement.executeLargeUpdate(DML_RETURNING));
+    }
+  }
+
+  @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(LARGE_DML, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeLargeUpdate(QUERY, Statement.NO_GENERATED_KEYS));
+      verifyException(
+          () -> statement.executeLargeUpdate(DML_RETURNING, 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(LARGE_DML, new String[] {"id"}));
+      verifyException(
+          () -> statement.executeLargeUpdate(QUERY, new String[] {"id"}), Code.FAILED_PRECONDITION);
+      assertEquals(1L, statement.executeLargeUpdate(DML_RETURNING, 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(LARGE_DML, new int[] {1}));
+      verifyException(() -> statement.executeLargeUpdate(QUERY, new int[] {1}));
+      verifyException(() -> statement.executeLargeUpdate(DML_RETURNING, 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(LARGE_DML), LARGE_UPDATE_COUNT);
+      verifyUpdateCount(statement, () -> statement.execute(DDL), Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(QUERY));
+      verifyResultSet(statement, () -> statement.execute(DML_RETURNING));
+    }
+  }
+
+  @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(LARGE_DML, 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(DML_RETURNING, 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(LARGE_DML, 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(DML_RETURNING, 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(LARGE_DML, 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(DML_RETURNING, 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(DML_RETURNING).executeQuery()) {
+        verifyResultSet(resultSet);
+      }
+      verifyException(() -> connection.prepareStatement(DML).executeQuery());
+      verifyException(() -> connection.prepareStatement(LARGE_DML).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(LARGE_DML).executeUpdate());
+      verifyException(() -> connection.prepareStatement(QUERY).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(DML_RETURNING).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(LARGE_DML, Statement.NO_GENERATED_KEYS).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(QUERY, Statement.NO_GENERATED_KEYS).executeUpdate());
+      verifyException(
+          () ->
+              connection
+                  .prepareStatement(DML_RETURNING, 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(LARGE_DML, new String[] {"id"}).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(QUERY, new String[] {"id"}).executeUpdate(),
+          Code.FAILED_PRECONDITION);
+      assertEquals(
+          1, connection.prepareStatement(DML_RETURNING, 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(LARGE_DML, new int[] {1}).executeUpdate());
+      verifyException(() -> connection.prepareStatement(QUERY, new int[] {1}).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(DML_RETURNING, 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(LARGE_DML).executeLargeUpdate());
+      verifyException(() -> connection.prepareStatement(QUERY).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(DML_RETURNING).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(LARGE_DML, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      verifyException(
+          () ->
+              connection.prepareStatement(QUERY, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      verifyException(
+          () ->
+              connection
+                  .prepareStatement(DML_RETURNING, 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(LARGE_DML, new String[] {"id"}).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(QUERY, new String[] {"id"}).executeLargeUpdate(),
+          Code.FAILED_PRECONDITION);
+      assertEquals(
+          1L, connection.prepareStatement(DML_RETURNING, 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(LARGE_DML, new int[] {1}).executeLargeUpdate());
+      verifyException(() -> connection.prepareStatement(QUERY, new int[] {1}).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(DML_RETURNING, 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(LARGE_DML), PreparedStatement::execute, LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL), PreparedStatement::execute, Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(connection.prepareStatement(QUERY), PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(DML_RETURNING), 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(LARGE_DML, 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(DML_RETURNING, 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(LARGE_DML, 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(DML_RETURNING, 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(LARGE_DML, 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(DML_RETURNING, 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;
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());
+      }
+    }
+  }
+}