Repository: googleapis/java-spanner-jdbc
Pull request #1331 (2023-08-25 13:28:42)

fix: comments should be sent to Spanner for PostgreSQL databases
Comments for PostgreSQL-dialect databases should be included in the request that is sent to Cloud Spanner. This is needed, because for PostgreSQL comments can contain query hints.


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/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
index 99bee40..9713172 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
@@ -120,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
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 fcc0959..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 {
@@ -585,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 b62449e..65a54a8 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
@@ -23,10 +23,12 @@ 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;
@@ -50,27 +52,62 @@ 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";
-  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 *";
+
+  @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(
@@ -114,15 +151,15 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
                     .build())
             .build();
     mockSpanner.putStatementResult(
-        StatementResult.query(com.google.cloud.spanner.Statement.of(QUERY), resultSet));
+        StatementResult.query(com.google.cloud.spanner.Statement.of(query), resultSet));
     mockSpanner.putStatementResult(
-        StatementResult.update(com.google.cloud.spanner.Statement.of(DML), 1L));
+        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));
+            com.google.cloud.spanner.Statement.of(largeDml), LARGE_UPDATE_COUNT));
     mockSpanner.putStatementResult(
         StatementResult.query(
-            com.google.cloud.spanner.Statement.of(DML_RETURNING),
+            com.google.cloud.spanner.Statement.of(dmlReturning),
             resultSet
                 .toBuilder()
                 .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
@@ -163,14 +200,14 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testStatementExecuteQuery() throws SQLException {
     try (Connection connection = createConnection();
         Statement statement = connection.createStatement()) {
-      try (ResultSet resultSet = statement.executeQuery(QUERY)) {
+      try (ResultSet resultSet = statement.executeQuery(query)) {
         verifyResultSet(resultSet);
       }
-      try (ResultSet resultSet = statement.executeQuery(DML_RETURNING)) {
+      try (ResultSet resultSet = statement.executeQuery(dmlReturning)) {
         verifyResultSet(resultSet);
       }
-      verifyException(() -> statement.executeQuery(DML));
-      verifyException(() -> statement.executeQuery(LARGE_DML));
+      verifyException(() -> statement.executeQuery(dml));
+      verifyException(() -> statement.executeQuery(largeDml));
       verifyException(() -> statement.executeQuery(DDL));
     }
   }
@@ -179,11 +216,11 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testStatementExecuteUpdate() throws SQLException {
     try (Connection connection = createConnection();
         Statement statement = connection.createStatement()) {
-      assertEquals(1, statement.executeUpdate(DML));
+      assertEquals(1, statement.executeUpdate(dml));
       assertEquals(0, statement.executeUpdate(DDL));
-      verifyOverflow(() -> statement.executeUpdate(LARGE_DML));
-      verifyException(() -> statement.executeUpdate(QUERY));
-      verifyException(() -> statement.executeUpdate(DML_RETURNING));
+      verifyOverflow(() -> statement.executeUpdate(largeDml));
+      verifyException(() -> statement.executeUpdate(query));
+      verifyException(() -> statement.executeUpdate(dmlReturning));
     }
   }
 
@@ -192,11 +229,11 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
     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(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));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeUpdate(query, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeUpdate(dmlReturning, Statement.NO_GENERATED_KEYS));
     }
   }
 
@@ -204,12 +241,12 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testStatementExecuteUpdateReturnColumnNames() throws SQLException {
     try (Connection connection = createConnection();
         Statement statement = connection.createStatement()) {
-      assertEquals(1, statement.executeUpdate(DML, new String[] {"id"}));
+      assertEquals(1, statement.executeUpdate(dml, new String[] {"id"}));
       assertEquals(0, statement.executeUpdate(DDL, new String[] {"id"}));
-      verifyOverflow(() -> statement.executeUpdate(LARGE_DML, new String[] {"id"}));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, new String[] {"id"}));
       verifyException(
-          () -> statement.executeUpdate(QUERY, new String[] {"id"}), Code.FAILED_PRECONDITION);
-      assertEquals(1, statement.executeUpdate(DML_RETURNING, new String[] {"id"}));
+          () -> statement.executeUpdate(query, new String[] {"id"}), Code.FAILED_PRECONDITION);
+      assertEquals(1, statement.executeUpdate(dmlReturning, new String[] {"id"}));
     }
   }
 
@@ -217,11 +254,11 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testStatementExecuteUpdateReturnColumnIndexes() throws SQLException {
     try (Connection connection = createConnection();
         Statement statement = connection.createStatement()) {
-      assertEquals(1, statement.executeUpdate(DML, new int[] {1}));
+      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}));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, new int[] {1}));
+      verifyException(() -> statement.executeUpdate(query, new int[] {1}));
+      verifyException(() -> statement.executeUpdate(dmlReturning, new int[] {1}));
     }
   }
 
@@ -229,11 +266,11 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testStatementLargeExecuteUpdate() throws SQLException {
     try (Connection connection = createConnection();
         Statement statement = connection.createStatement()) {
-      assertEquals(1L, statement.executeLargeUpdate(DML));
+      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));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml));
+      verifyException(() -> statement.executeLargeUpdate(query));
+      verifyException(() -> statement.executeLargeUpdate(dmlReturning));
     }
   }
 
@@ -242,13 +279,13 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
     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(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));
+          LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeLargeUpdate(query, Statement.NO_GENERATED_KEYS));
       verifyException(
-          () -> statement.executeLargeUpdate(DML_RETURNING, Statement.NO_GENERATED_KEYS));
+          () -> statement.executeLargeUpdate(dmlReturning, Statement.NO_GENERATED_KEYS));
     }
   }
 
@@ -256,13 +293,12 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testStatementExecuteLargeUpdateReturnColumnNames() throws SQLException {
     try (Connection connection = createConnection();
         Statement statement = connection.createStatement()) {
-      assertEquals(1, statement.executeLargeUpdate(DML, new String[] {"id"}));
+      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"}));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, new String[] {"id"}));
       verifyException(
-          () -> statement.executeLargeUpdate(QUERY, new String[] {"id"}), Code.FAILED_PRECONDITION);
-      assertEquals(1L, statement.executeLargeUpdate(DML_RETURNING, new String[] {"id"}));
+          () -> statement.executeLargeUpdate(query, new String[] {"id"}), Code.FAILED_PRECONDITION);
+      assertEquals(1L, statement.executeLargeUpdate(dmlReturning, new String[] {"id"}));
     }
   }
 
@@ -270,11 +306,11 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testStatementExecuteLargeUpdateReturnColumnIndexes() throws SQLException {
     try (Connection connection = createConnection();
         Statement statement = connection.createStatement()) {
-      assertEquals(1, statement.executeLargeUpdate(DML, new int[] {1}));
+      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}));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, new int[] {1}));
+      verifyException(() -> statement.executeLargeUpdate(query, new int[] {1}));
+      verifyException(() -> statement.executeLargeUpdate(dmlReturning, new int[] {1}));
     }
   }
 
@@ -282,11 +318,11 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   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(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(DML_RETURNING));
+      verifyResultSet(statement, () -> statement.execute(query));
+      verifyResultSet(statement, () -> statement.execute(dmlReturning));
     }
   }
 
@@ -295,18 +331,18 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
     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(dml, Statement.NO_GENERATED_KEYS), 1L);
       verifyUpdateCount(
           statement,
-          () -> statement.execute(LARGE_DML, Statement.NO_GENERATED_KEYS),
+          () -> 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(query, Statement.NO_GENERATED_KEYS));
       verifyResultSet(
-          statement, () -> statement.execute(DML_RETURNING, Statement.NO_GENERATED_KEYS));
+          statement, () -> statement.execute(dmlReturning, Statement.NO_GENERATED_KEYS));
     }
   }
 
@@ -314,13 +350,13 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   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(dml, new String[] {"id"}), 1L);
       verifyUpdateCount(
-          statement, () -> statement.execute(LARGE_DML, new String[] {"id"}), LARGE_UPDATE_COUNT);
+          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(DML_RETURNING, new String[] {"id"}));
+      verifyResultSet(statement, () -> statement.execute(query, new String[] {"id"}));
+      verifyResultSet(statement, () -> statement.execute(dmlReturning, new String[] {"id"}));
     }
   }
 
@@ -328,27 +364,27 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   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(dml, new int[] {1}), 1L);
       verifyUpdateCount(
-          statement, () -> statement.execute(LARGE_DML, new int[] {1}), LARGE_UPDATE_COUNT);
+          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(DML_RETURNING, new int[] {1}));
+      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()) {
+      try (ResultSet resultSet = connection.prepareStatement(query).executeQuery()) {
         verifyResultSet(resultSet);
       }
-      try (ResultSet resultSet = connection.prepareStatement(DML_RETURNING).executeQuery()) {
+      try (ResultSet resultSet = connection.prepareStatement(dmlReturning).executeQuery()) {
         verifyResultSet(resultSet);
       }
-      verifyException(() -> connection.prepareStatement(DML).executeQuery());
-      verifyException(() -> connection.prepareStatement(LARGE_DML).executeQuery());
+      verifyException(() -> connection.prepareStatement(dml).executeQuery());
+      verifyException(() -> connection.prepareStatement(largeDml).executeQuery());
       verifyException(() -> connection.prepareStatement(DDL).executeQuery());
     }
   }
@@ -356,14 +392,14 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   @Test
   public void testPreparedStatementExecuteUpdate() throws SQLException {
     try (Connection connection = createConnection()) {
-      assertEquals(1, connection.prepareStatement(DML).executeUpdate());
+      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());
+      verifyOverflow(() -> connection.prepareStatement(largeDml).executeUpdate());
+      verifyException(() -> connection.prepareStatement(query).executeUpdate());
       verifyException(
-          () -> connection.prepareStatement(DML_RETURNING).executeUpdate(), Code.INVALID_ARGUMENT);
+          () -> connection.prepareStatement(dmlReturning).executeUpdate(), Code.INVALID_ARGUMENT);
     }
   }
 
@@ -372,20 +408,19 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
     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());
+          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());
+          () -> connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS).executeUpdate());
       verifyException(
-          () -> connection.prepareStatement(QUERY, Statement.NO_GENERATED_KEYS).executeUpdate());
+          () -> connection.prepareStatement(query, Statement.NO_GENERATED_KEYS).executeUpdate());
       verifyException(
           () ->
               connection
-                  .prepareStatement(DML_RETURNING, Statement.NO_GENERATED_KEYS)
+                  .prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS)
                   .executeUpdate(),
           Code.INVALID_ARGUMENT);
     }
@@ -394,31 +429,31 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   @Test
   public void testPreparedStatementExecuteUpdateReturnColumnNames() throws SQLException {
     try (Connection connection = createConnection()) {
-      assertEquals(1, connection.prepareStatement(DML, new String[] {"id"}).executeUpdate());
+      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());
+          () -> connection.prepareStatement(largeDml, new String[] {"id"}).executeUpdate());
       verifyException(
-          () -> connection.prepareStatement(QUERY, new String[] {"id"}).executeUpdate(),
+          () -> connection.prepareStatement(query, new String[] {"id"}).executeUpdate(),
           Code.FAILED_PRECONDITION);
       assertEquals(
-          1, connection.prepareStatement(DML_RETURNING, new String[] {"id"}).executeUpdate());
+          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());
+      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());
+      verifyOverflow(() -> connection.prepareStatement(largeDml, new int[] {1}).executeUpdate());
+      verifyException(() -> connection.prepareStatement(query, new int[] {1}).executeUpdate());
       verifyException(
-          () -> connection.prepareStatement(DML_RETURNING, new int[] {1}).executeUpdate(),
+          () -> connection.prepareStatement(dmlReturning, new int[] {1}).executeUpdate(),
           Code.INVALID_ARGUMENT);
     }
   }
@@ -426,14 +461,14 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   @Test
   public void testPreparedStatementLargeExecuteUpdate() throws SQLException {
     try (Connection connection = createConnection()) {
-      assertEquals(1L, connection.prepareStatement(DML).executeLargeUpdate());
+      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());
+      assertEquals(LARGE_UPDATE_COUNT, connection.prepareStatement(largeDml).executeLargeUpdate());
+      verifyException(() -> connection.prepareStatement(query).executeLargeUpdate());
       verifyException(
-          () -> connection.prepareStatement(DML_RETURNING).executeLargeUpdate(),
+          () -> connection.prepareStatement(dmlReturning).executeLargeUpdate(),
           Code.INVALID_ARGUMENT);
     }
   }
@@ -443,21 +478,21 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
     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());
+          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());
+          connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
       verifyException(
           () ->
-              connection.prepareStatement(QUERY, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+              connection.prepareStatement(query, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
       verifyException(
           () ->
               connection
-                  .prepareStatement(DML_RETURNING, Statement.NO_GENERATED_KEYS)
+                  .prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS)
                   .executeLargeUpdate(),
           Code.INVALID_ARGUMENT);
     }
@@ -466,35 +501,35 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   @Test
   public void testPreparedStatementExecuteLargeUpdateReturnColumnNames() throws SQLException {
     try (Connection connection = createConnection()) {
-      assertEquals(1, connection.prepareStatement(DML, new String[] {"id"}).executeLargeUpdate());
+      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());
+          connection.prepareStatement(largeDml, new String[] {"id"}).executeLargeUpdate());
       verifyException(
-          () -> connection.prepareStatement(QUERY, new String[] {"id"}).executeLargeUpdate(),
+          () -> connection.prepareStatement(query, new String[] {"id"}).executeLargeUpdate(),
           Code.FAILED_PRECONDITION);
       assertEquals(
-          1L, connection.prepareStatement(DML_RETURNING, new String[] {"id"}).executeLargeUpdate());
+          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());
+      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());
+          connection.prepareStatement(largeDml, new int[] {1}).executeLargeUpdate());
+      verifyException(() -> connection.prepareStatement(query, new int[] {1}).executeLargeUpdate());
       verifyException(
-          () -> connection.prepareStatement(DML_RETURNING, new int[] {1}).executeLargeUpdate(),
+          () -> connection.prepareStatement(dmlReturning, new int[] {1}).executeLargeUpdate(),
           Code.INVALID_ARGUMENT);
     }
   }
@@ -502,14 +537,14 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   @Test
   public void testPreparedStatementExecute() throws SQLException {
     try (Connection connection = createConnection()) {
-      verifyPreparedUpdateCount(connection.prepareStatement(DML), PreparedStatement::execute, 1L);
+      verifyPreparedUpdateCount(connection.prepareStatement(dml), PreparedStatement::execute, 1L);
       verifyPreparedUpdateCount(
-          connection.prepareStatement(LARGE_DML), PreparedStatement::execute, LARGE_UPDATE_COUNT);
+          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(query), PreparedStatement::execute);
       verifyPreparedResultSet(
-          connection.prepareStatement(DML_RETURNING), PreparedStatement::execute);
+          connection.prepareStatement(dmlReturning), PreparedStatement::execute);
     }
   }
 
@@ -518,11 +553,11 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
     try (Connection connection = createConnection()) {
       // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
       verifyPreparedUpdateCount(
-          connection.prepareStatement(DML, Statement.NO_GENERATED_KEYS),
+          connection.prepareStatement(dml, Statement.NO_GENERATED_KEYS),
           PreparedStatement::execute,
           1L);
       verifyPreparedUpdateCount(
-          connection.prepareStatement(LARGE_DML, Statement.NO_GENERATED_KEYS),
+          connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS),
           PreparedStatement::execute,
           LARGE_UPDATE_COUNT);
       verifyPreparedUpdateCount(
@@ -530,10 +565,10 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
           PreparedStatement::execute,
           Statement.SUCCESS_NO_INFO);
       verifyPreparedResultSet(
-          connection.prepareStatement(QUERY, Statement.NO_GENERATED_KEYS),
+          connection.prepareStatement(query, Statement.NO_GENERATED_KEYS),
           PreparedStatement::execute);
       verifyPreparedResultSet(
-          connection.prepareStatement(DML_RETURNING, Statement.NO_GENERATED_KEYS),
+          connection.prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS),
           PreparedStatement::execute);
     }
   }
@@ -542,9 +577,9 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testPreparedStatementExecuteReturnColumnNames() throws SQLException {
     try (Connection connection = createConnection()) {
       verifyPreparedUpdateCount(
-          connection.prepareStatement(DML, new String[] {"id"}), PreparedStatement::execute, 1L);
+          connection.prepareStatement(dml, new String[] {"id"}), PreparedStatement::execute, 1L);
       verifyPreparedUpdateCount(
-          connection.prepareStatement(LARGE_DML, new String[] {"id"}),
+          connection.prepareStatement(largeDml, new String[] {"id"}),
           PreparedStatement::execute,
           LARGE_UPDATE_COUNT);
       verifyPreparedUpdateCount(
@@ -552,9 +587,9 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
           PreparedStatement::execute,
           Statement.SUCCESS_NO_INFO);
       verifyPreparedResultSet(
-          connection.prepareStatement(QUERY, new String[] {"id"}), PreparedStatement::execute);
+          connection.prepareStatement(query, new String[] {"id"}), PreparedStatement::execute);
       verifyPreparedResultSet(
-          connection.prepareStatement(DML_RETURNING, new String[] {"id"}),
+          connection.prepareStatement(dmlReturning, new String[] {"id"}),
           PreparedStatement::execute);
     }
   }
@@ -563,9 +598,9 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
   public void testPreparedStatementExecuteReturnColumnIndexes() throws SQLException {
     try (Connection connection = createConnection()) {
       verifyPreparedUpdateCount(
-          connection.prepareStatement(DML, new int[] {1}), PreparedStatement::execute, 1L);
+          connection.prepareStatement(dml, new int[] {1}), PreparedStatement::execute, 1L);
       verifyPreparedUpdateCount(
-          connection.prepareStatement(LARGE_DML, new int[] {1}),
+          connection.prepareStatement(largeDml, new int[] {1}),
           PreparedStatement::execute,
           LARGE_UPDATE_COUNT);
       verifyPreparedUpdateCount(
@@ -573,9 +608,9 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
           PreparedStatement::execute,
           Statement.SUCCESS_NO_INFO);
       verifyPreparedResultSet(
-          connection.prepareStatement(QUERY, new int[] {1}), PreparedStatement::execute);
+          connection.prepareStatement(query, new int[] {1}), PreparedStatement::execute);
       verifyPreparedResultSet(
-          connection.prepareStatement(DML_RETURNING, new int[] {1}), PreparedStatement::execute);
+          connection.prepareStatement(dmlReturning, new int[] {1}), PreparedStatement::execute);
     }
   }
 
@@ -671,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());
+      }
+    }
+  }
+}