Repository: googleapis/java-spanner-jdbc
Pull request #1304 (2023-08-02 02:29:19)

build(deps): bump cryptography from 41.0.2 to 41.0.3 in /.kokoro
Bumps [cryptography](https://github.com/pyca/cryptography) from 41.0.2 to 41.0.3.
<details>
<summary>Changelog</summary>
<p><em>Sourced from <a href="https://github.com/pyca/cryptography/blob/main/CHANGELOG.rst">cryptography's changelog</a>.</em></p>
<blockquote>
<p>41.0.3 - 2023-08-01</p>
<pre><code>
* Fixed performance regression loading DH public keys.
* Fixed a memory leak when using
  :class:`~cryptography.hazmat.primitives.ciphers.aead.ChaCha20Poly1305`.
* Updated Windows, macOS, and Linux wheels to be compiled with OpenSSL 3.1.2.
<p>.. _v41-0-2:
</code></pre></p>
</blockquote>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a href="https://github.com/pyca/cryptography/commit/b22271cf3c3dd8dc8978f8f4b00b5c7060b6538d"><code>b22271c</code></a> bump for 41.0.3 (<a href="https://redirect.github.com/pyca/cryptography/issues/9330">#9330</a>)</li>
<li><a href="https://github.com/pyca/cryptography/commit/774a4a16cbd22a89fdb4195ade9e4fcee27a7afa"><code>774a4a1</code></a> Only check DH key validity when loading a private key. (<a href="https://redirect.github.com/pyca/cryptography/issues/9071">#9071</a>) (<a href="https://redirect.github.com/pyca/cryptography/issues/9319">#9319</a>)</li>
<li><a href="https://github.com/pyca/cryptography/commit/bfa4d95f0f356f2d535efd5c775e0fb3efe90ef2"><code>bfa4d95</code></a> changelog for 41.0.3 (<a href="https://redirect.github.com/pyca/cryptography/issues/9320">#9320</a>)</li>
<li><a href="https://github.com/pyca/cryptography/commit/0da7165aa73c0a4865b0a4d9e019db3c16eea55a"><code>0da7165</code></a> backport fix the memory leak in fixedpool (<a href="https://redirect.github.com/pyca/cryptography/issues/9272">#9272</a>) (<a href="https://redirect.github.com/pyca/cryptography/issues/9309">#9309</a>)</li>
<li>See full diff in <a href="https://github.com/pyca/cryptography/compare/41.0.2...41.0.3">compare view</a></li>
</ul>
</details>
<br />


[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=cryptography&package-manager=pip&previous-version=41.0.2&new-version=41.0.3)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores)

You can trigger a rebase of this PR by commenting `@dependabot rebase`.

[//]: # (dependabot-automerge-start)
[//]: # (dependabot-automerge-end)

---

<details>
<summary>Dependabot commands and options</summary>
<br />

You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot merge` will merge this PR after your CI passes on it
- `@dependabot squash and merge` will squash and merge this PR after your CI passes on it
- `@dependabot cancel merge` will cancel a previously requested merge and block automerging
- `@dependabot reopen` will reopen this PR if it is closed
- `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
You can disable automated security fix PRs for this repo from the [Security Alerts page](https://github.com/googleapis/java-spanner-jdbc/network/alerts).

</details>

> **Note**
> Automatic rebases have been disabled on this pull request as it has been open for over 30 days.


olavloite: @dependabot rebase
dependabot[bot]: Looks like this PR has been edited by someone other than Dependabot. That means Dependabot can't rebase it - sorry!

If you're happy for Dependabot to recreate it from scratch, overwriting any edits, you can request `@dependabot recreate`.

olavloite: @dependabot recreate
dependabot[bot]: OK, I won't notify you again about this release, but will get in touch when a new version is available. If you'd rather skip all updates until the next major or minor version, let me know by commenting `@dependabot ignore this major version` or `@dependabot ignore this minor version`. You can also ignore all major, minor, or patch releases for a dependency by adding an [`ignore` condition](https://docs.github.com/en/code-security/supply-chain-security/configuration-options-for-dependency-updates#ignore) with the desired `update_types` to your config file.

If you change your mind, just re-open this PR and I'll resolve any conflicts on it.
olavloite: @dependabot rebase

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index beda76d..af67a88 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,28 @@
 # Changelog
 
+## [2.13.0](https://github.com/googleapis/java-spanner-jdbc/compare/v2.12.1...v2.13.0) (2023-09-15)
+
+
+### Features
+
+* Support partitioned queries ([#1300](https://github.com/googleapis/java-spanner-jdbc/issues/1300)) ([c50da41](https://github.com/googleapis/java-spanner-jdbc/commit/c50da41e688ff48f8967c0f114f5bac8eaac49f9))
+
+
+### Bug Fixes
+
+* Comments should be sent to Spanner for PostgreSQL databases ([#1331](https://github.com/googleapis/java-spanner-jdbc/issues/1331)) ([7c9e781](https://github.com/googleapis/java-spanner-jdbc/commit/7c9e781bf45b112266e278e1df1586e56043698e))
+
+
+### Documentation
+
+* Create Spring Data JDBC sample ([#1334](https://github.com/googleapis/java-spanner-jdbc/issues/1334)) ([cefea55](https://github.com/googleapis/java-spanner-jdbc/commit/cefea55086eb191f71a1a493e046cb136f9f9f87))
+
+
+### Dependencies
+
+* Update actions/checkout action to v4 - abandoned ([#1333](https://github.com/googleapis/java-spanner-jdbc/issues/1333)) ([ce82b42](https://github.com/googleapis/java-spanner-jdbc/commit/ce82b42d3abb8de0f8b3ee2915c2008673775ea1))
+* Update dependency org.springframework.data:spring-data-bom to v2023.0.4 ([#1347](https://github.com/googleapis/java-spanner-jdbc/issues/1347)) ([893f61a](https://github.com/googleapis/java-spanner-jdbc/commit/893f61ab04e32c690f1ff9fc813bd2ba6ebca328))
+
 ## [2.12.1](https://github.com/googleapis/java-spanner-jdbc/compare/v2.12.0...v2.12.1) (2023-09-12)
 
 
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/pom.xml b/pom.xml
index 5d2bd1a..cbdc4b5 100644
--- a/pom.xml
+++ b/pom.xml
@@ -4,7 +4,7 @@
   xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
   <modelVersion>4.0.0</modelVersion>
   <artifactId>google-cloud-spanner-jdbc</artifactId>
-  <version>2.12.2-SNAPSHOT</version><!-- {x-version-update:google-cloud-spanner-jdbc:current} -->
+  <version>2.13.1-SNAPSHOT</version><!-- {x-version-update:google-cloud-spanner-jdbc:current} -->
   <packaging>jar</packaging>
   <name>Google Cloud Spanner JDBC</name>
   <url>https://github.com/googleapis/java-spanner-jdbc</url>
@@ -420,7 +420,7 @@
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-javadoc-plugin</artifactId>
-        <version>3.5.0</version>
+        <version>3.6.0</version>
         <reportSets>
           <reportSet>
             <id>html</id>
diff --git a/samples/install-without-bom/pom.xml b/samples/install-without-bom/pom.xml
index 0e0e2d7..2cf650a 100644
--- a/samples/install-without-bom/pom.xml
+++ b/samples/install-without-bom/pom.xml
@@ -29,7 +29,7 @@
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner-jdbc</artifactId>
-      <version>2.12.1</version>
+      <version>2.13.0</version>
     </dependency>
     <!-- [END spanner-jdbc_install_without_bom] -->
 
diff --git a/samples/snapshot/pom.xml b/samples/snapshot/pom.xml
index f0f4191..69d2a1b 100644
--- a/samples/snapshot/pom.xml
+++ b/samples/snapshot/pom.xml
@@ -28,7 +28,7 @@
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner-jdbc</artifactId>
-      <version>2.12.2-SNAPSHOT</version>
+      <version>2.13.1-SNAPSHOT</version>
     </dependency>
     <!-- {x-version-update-end} -->
 
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..6816cff 100644
--- a/samples/spring-data-jdbc/pom.xml
+++ b/samples/spring-data-jdbc/pom.xml
@@ -23,14 +23,14 @@
       <dependency>
         <groupId>org.springframework.data</groupId>
         <artifactId>spring-data-bom</artifactId>
-        <version>2023.0.3</version>
+        <version>2023.0.4</version>
         <scope>import</scope>
         <type>pom</type>
       </dependency>
       <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>
@@ -48,7 +48,7 @@
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner-jdbc</artifactId>
-      <version>2.12.1</version>
+      <version>2.13.0</version>
     </dependency>
     <dependency>
       <groupId>org.postgresql</groupId>
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
index e7b1a47..9713172 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
@@ -24,12 +24,15 @@ import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
+import com.google.rpc.Code;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 /** Base class for Cloud Spanner JDBC {@link Statement}s */
 abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Statement {
@@ -117,6 +120,25 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
     }
   }
 
+  /** Functional interface that throws {@link SQLException}. */
+  interface JdbcFunction<T, R> {
+    R apply(T t) throws SQLException;
+  }
+
+  /** Runs the given function with the timeout that has been set on this statement. */
+  protected <T> T runWithStatementTimeout(JdbcFunction<Connection, T> function)
+      throws SQLException {
+    checkClosed();
+    StatementTimeout originalTimeout = setTemporaryStatementTimeout();
+    try {
+      return function.apply(getConnection().getSpannerConnection());
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    } finally {
+      resetStatementTimeout(originalTimeout);
+    }
+  }
+
   /**
    * Sets the statement timeout of the Spanner {@link Connection} to the query timeout of this JDBC
    * {@link Statement} and returns the original timeout of the Spanner {@link Connection} so it can
@@ -183,21 +205,74 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
       QueryAnalyzeMode analyzeMode,
       QueryOption... options)
       throws SQLException {
+    Options.QueryOption[] queryOptions = getQueryOptions(options);
+    return doWithStatementTimeout(
+        () -> {
+          com.google.cloud.spanner.ResultSet resultSet;
+          if (analyzeMode == null) {
+            resultSet = connection.getSpannerConnection().executeQuery(statement, queryOptions);
+          } else {
+            resultSet = connection.getSpannerConnection().analyzeQuery(statement, analyzeMode);
+          }
+          return JdbcResultSet.of(this, resultSet);
+        });
+  }
+
+  private <T> T doWithStatementTimeout(Supplier<T> runnable) throws SQLException {
+    return doWithStatementTimeout(runnable, ignore -> Boolean.TRUE);
+  }
+
+  private <T> T doWithStatementTimeout(
+      Supplier<T> runnable, Function<T, Boolean> shouldResetTimeout) throws SQLException {
     StatementTimeout originalTimeout = setTemporaryStatementTimeout();
+    T result = null;
     try {
-      com.google.cloud.spanner.ResultSet resultSet;
-      if (analyzeMode == null) {
-        resultSet =
-            connection.getSpannerConnection().executeQuery(statement, getQueryOptions(options));
-      } else {
-        resultSet = connection.getSpannerConnection().analyzeQuery(statement, analyzeMode);
-      }
-      return JdbcResultSet.of(this, resultSet);
-    } catch (SpannerException e) {
-      throw JdbcSqlExceptionFactory.of(e);
+      result = runnable.get();
+      return result;
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
     } finally {
-      resetStatementTimeout(originalTimeout);
+      if (shouldResetTimeout.apply(result)) {
+        resetStatementTimeout(originalTimeout);
+      }
+    }
+  }
+
+  /**
+   * Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
+   * statement.
+   *
+   * @param statement The SQL statement to execute
+   * @return the number of rows that was inserted/updated/deleted
+   * @throws SQLException if a database error occurs, or if the number of rows affected is larger
+   *     than {@link Integer#MAX_VALUE}
+   */
+  int executeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
+    return checkedCast(executeLargeUpdate(statement));
+  }
+
+  /**
+   * Do a checked cast from long to int. Throws a {@link SQLException} with code {@link
+   * Code#OUT_OF_RANGE} if the update count is too big to fit in an int.
+   */
+  int checkedCast(long updateCount) throws SQLException {
+    if (updateCount > Integer.MAX_VALUE) {
+      throw JdbcSqlExceptionFactory.of(
+          "update count too large for executeUpdate: " + updateCount, Code.OUT_OF_RANGE);
     }
+    return (int) updateCount;
+  }
+
+  /**
+   * Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
+   * statement.
+   *
+   * @param statement The SQL statement to execute
+   * @return the number of rows that was inserted/updated/deleted
+   * @throws SQLException if a database error occurs
+   */
+  long executeLargeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
+    return doWithStatementTimeout(() -> connection.getSpannerConnection().executeUpdate(statement));
   }
 
   /**
@@ -210,24 +285,16 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
    * @throws SQLException if a database error occurs.
    */
   StatementResult execute(com.google.cloud.spanner.Statement statement) throws SQLException {
-    StatementTimeout originalTimeout = setTemporaryStatementTimeout();
-    boolean mustResetTimeout = false;
-    try {
-      StatementResult result = connection.getSpannerConnection().execute(statement);
-      mustResetTimeout = !resultIsSetStatementTimeout(result);
-      if (mustResetTimeout && resultIsShowStatementTimeout(result)) {
-        // it was a 'SHOW STATEMENT_TIMEOUT statement, we need to re-run to get the correct value
-        mustResetTimeout = false;
-        result = rerunShowStatementTimeout(statement, originalTimeout);
-      }
-      return result;
-    } catch (SpannerException e) {
-      throw JdbcSqlExceptionFactory.of(e);
-    } finally {
-      if (mustResetTimeout) {
-        resetStatementTimeout(originalTimeout);
-      }
+    StatementResult statementResult =
+        doWithStatementTimeout(
+            () -> connection.getSpannerConnection().execute(statement),
+            result -> !resultIsSetStatementTimeout(result));
+    if (resultIsShowStatementTimeout(statementResult)) {
+      // We can safely re-run it without first resetting the timeout to the original value, as that
+      // has already been done by the 'doWithStatementTimeout' function.
+      return rerunShowStatementTimeout(statement);
     }
+    return statementResult;
   }
 
   /**
@@ -250,18 +317,22 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
    *     executed was a SET STATEMENT_TIMEOUT statement.
    */
   private boolean resultIsSetStatementTimeout(StatementResult result) {
-    return result.getClientSideStatementType() == ClientSideStatementType.SET_STATEMENT_TIMEOUT;
+    return result != null
+        && result.getClientSideStatementType() == ClientSideStatementType.SET_STATEMENT_TIMEOUT;
   }
 
   private boolean resultIsShowStatementTimeout(StatementResult result) {
-    return result.getClientSideStatementType() == ClientSideStatementType.SHOW_STATEMENT_TIMEOUT;
+    return result != null
+        && result.getClientSideStatementType() == ClientSideStatementType.SHOW_STATEMENT_TIMEOUT;
   }
 
-  private StatementResult rerunShowStatementTimeout(
-      com.google.cloud.spanner.Statement statement, StatementTimeout originalTimeout)
+  private StatementResult rerunShowStatementTimeout(com.google.cloud.spanner.Statement statement)
       throws SQLException {
-    resetStatementTimeout(originalTimeout);
-    return connection.getSpannerConnection().execute(statement);
+    try {
+      return connection.getSpannerConnection().execute(statement);
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    }
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
index cf368f6..c161ee5 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
@@ -22,12 +22,14 @@ import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.CommitStats;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.Mutation;
-import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.AutocommitDmlMode;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.connection.TransactionMode;
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Timestamp;
 import java.util.Iterator;
@@ -334,6 +336,73 @@ public interface CloudSpannerJdbcConnection extends Connection {
     return Dialect.GOOGLE_STANDARD_SQL;
   }
 
+  /**
+   * Enable data boost for partitioned queries. See also {@link
+   * CloudSpannerJdbcStatement#partitionQuery(String, PartitionOptions, QueryOption...)} and {@link
+   * CloudSpannerJdbcPreparedStatement#partitionQuery(PartitionOptions, QueryOption...)}.
+   */
+  default void setDataBoostEnabled(boolean dataBoostEnabled) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Returns whether data boost is enabled for partitioned queries. See also {@link
+   * CloudSpannerJdbcStatement#partitionQuery(String, PartitionOptions, QueryOption...)} and {@link
+   * CloudSpannerJdbcPreparedStatement#partitionQuery(PartitionOptions, QueryOption...)}.
+   */
+  default boolean isDataBoostEnabled() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Sets whether this connection should always use partitioned queries when a query is executed on
+   * this connection. Setting this flag to <code>true</code> and then executing a query that cannot
+   * be partitioned, or executing a query in a read/write transaction, will cause an error. Use this
+   * flag in combination with {@link #setDataBoostEnabled(boolean)} to force all queries on this
+   * connection to use data boost.
+   */
+  default void setAutoPartitionMode(boolean alwaysUsePartitionedQueries) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /** Returns whether this connection will execute all queries as partitioned queries. */
+  default boolean isAutoPartitionMode() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Sets the maximum number of partitions that should be included as a hint to Cloud Spanner when
+   * partitioning a query on this connection. Note that this is only a hint and Cloud Spanner might
+   * choose to ignore the hint.
+   */
+  default void setMaxPartitions(int maxPartitions) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Gets the maximum number of partitions that should be included as a hint to Cloud Spanner when
+   * partitioning a query on this connection. Note that this is only a hint and Cloud Spanner might
+   * choose to ignore the hint.
+   */
+  default int getMaxPartitions() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Sets the maximum degree of parallelism that is used when executing a partitioned query. A
+   * partitioned query will use up to <code>maxThreads</code> to execute and retrieve the results
+   * from Cloud Spanner. Set this value to <code>0</code>> to use the number of available processors
+   * as returned by {@link Runtime#availableProcessors()}.
+   */
+  default void setMaxPartitionedParallelism(int maxThreads) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /** Returns the maximum degree of parallelism that is used for partitioned queries. */
+  default int getMaxPartitionedParallelism() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
   /**
    * @see
    *     com.google.cloud.spanner.connection.Connection#addTransactionRetryListener(com.google.cloud.spanner.connection.TransactionRetryListener)
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPartitionedQueryResultSet.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPartitionedQueryResultSet.java
new file mode 100644
index 0000000..3b6118f
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPartitionedQueryResultSet.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.spanner.jdbc;
+
+import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
+import java.sql.ResultSet;
+
+/**
+ * Result set that is returned for partitioned queries, e.g. for 'run partitioned query select ...'
+ * or for {@link CloudSpannerJdbcPreparedStatement#runPartitionedQuery(PartitionOptions,
+ * QueryOption...)}.
+ */
+public interface CloudSpannerJdbcPartitionedQueryResultSet extends ResultSet {
+  /** Returns the number of partitions that this result set contains. */
+  int getNumPartitions();
+
+  /** Returns the degree of parallelism that this result set uses. */
+  int getParallelism();
+}
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPreparedStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPreparedStatement.java
new file mode 100644
index 0000000..8e72d95
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcPreparedStatement.java
@@ -0,0 +1,61 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.jdbc;
+
+import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+
+/**
+ * This interface is implemented by {@link PreparedStatement}s that are created on Cloud Spanner
+ * JDBC connections.
+ */
+public interface CloudSpannerJdbcPreparedStatement extends PreparedStatement {
+
+  /**
+   * Partitions this query, so it can be executed in parallel. This method returns a {@link
+   * ResultSet} with a string-representation of the partitions that were created. These strings can
+   * be used to execute a partition either on this connection or an any other connection (on this
+   * host or an any other host) by calling the method {@link #runPartition()}. This method will
+   * automatically enable data boost for the query if {@link
+   * CloudSpannerJdbcConnection#isDataBoostEnabled()} returns true.
+   */
+  ResultSet partitionQuery(PartitionOptions partitionOptions, QueryOption... options)
+      throws SQLException;
+
+  /**
+   * Executes the given partition of a query. The partition that should be executed must be set as a
+   * string parameter on this {@link PreparedStatement} using {@link #setString(int, String)}. The
+   * value should be a string that was returned by {@link #partitionQuery(PartitionOptions,
+   * QueryOption...)}.
+   */
+  ResultSet runPartition() throws SQLException;
+
+  /**
+   * Executes the given query as a partitioned query. The query will first be partitioned using the
+   * {@link #partitionQuery(PartitionOptions, QueryOption...)} method. Each of the partitions will
+   * then be executed in the background, and the results will be merged into a single result set.
+   *
+   * <p>This method will use {@link CloudSpannerJdbcConnection#getMaxPartitionedParallelism()}
+   * threads to execute the partitioned query. Set this variable to a higher/lower value to
+   * increase/decrease the degree of parallelism used for execution.
+   */
+  CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
+      PartitionOptions partitionOptions, QueryOption... options) throws SQLException;
+}
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcStatement.java
new file mode 100644
index 0000000..6115c0b
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcStatement.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.jdbc;
+
+import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+/**
+ * This interface is implemented by {@link Statement}s that are created on Cloud Spanner JDBC
+ * connections.
+ */
+public interface CloudSpannerJdbcStatement extends Statement {
+
+  /**
+   * Partitions the given query, so it can be executed in parallel. This method returns a {@link
+   * ResultSet} with a string-representation of the partitions that were created. These strings can
+   * be used to execute a partition either on this connection or an any other connection (on this
+   * host or an any other host) by calling the method {@link #runPartition(String)}. This method
+   * will automatically enable data boost for the query if {@link
+   * CloudSpannerJdbcConnection#isDataBoostEnabled()} returns true.
+   */
+  ResultSet partitionQuery(String query, PartitionOptions partitionOptions, QueryOption... options)
+      throws SQLException;
+
+  /**
+   * Executes the given partition of a query. The encodedPartitionId should be a string that was
+   * returned by {@link #partitionQuery(String, PartitionOptions, QueryOption...)}.
+   */
+  ResultSet runPartition(String encodedPartitionId) throws SQLException;
+
+  /**
+   * Executes the given query as a partitioned query. The query will first be partitioned using the
+   * {@link #partitionQuery(String, PartitionOptions, QueryOption...)} method. Each of the
+   * partitions will then be executed in the background, and the results will be merged into a
+   * single result set.
+   *
+   * <p>This method will use {@link CloudSpannerJdbcConnection#getMaxPartitionedParallelism()}
+   * threads to execute the partitioned query. Set this variable to a higher/lower value to
+   * increase/decrease the degree of parallelism used for execution.
+   */
+  CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
+      String query, PartitionOptions partitionOptions, QueryOption... options) throws SQLException;
+}
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
index 6d6bd80..6031afe 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
@@ -25,6 +25,7 @@ import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.AutocommitDmlMode;
+import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.connection.TransactionMode;
@@ -44,6 +45,8 @@ import java.sql.Timestamp;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
 
 /** Jdbc Connection class for Google Cloud Spanner */
 class JdbcConnection extends AbstractJdbcConnection {
@@ -592,6 +595,72 @@ class JdbcConnection extends AbstractJdbcConnection {
     }
   }
 
+  /**
+   * Convenience method for calling a setter and translating any {@link SpannerException} to a
+   * {@link SQLException}.
+   */
+  private <T> void set(BiConsumer<Connection, T> setter, T value) throws SQLException {
+    checkClosed();
+    try {
+      setter.accept(getSpannerConnection(), value);
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    }
+  }
+
+  /**
+   * Convenience method for calling a getter and translating any {@link SpannerException} to a
+   * {@link SQLException}.
+   */
+  private <R> R get(Function<Connection, R> getter) throws SQLException {
+    checkClosed();
+    try {
+      return getter.apply(getSpannerConnection());
+    } catch (SpannerException spannerException) {
+      throw JdbcSqlExceptionFactory.of(spannerException);
+    }
+  }
+
+  @Override
+  public void setDataBoostEnabled(boolean dataBoostEnabled) throws SQLException {
+    set(Connection::setDataBoostEnabled, dataBoostEnabled);
+  }
+
+  @Override
+  public boolean isDataBoostEnabled() throws SQLException {
+    return get(Connection::isDataBoostEnabled);
+  }
+
+  @Override
+  public void setAutoPartitionMode(boolean autoPartitionMode) throws SQLException {
+    set(Connection::setAutoPartitionMode, autoPartitionMode);
+  }
+
+  @Override
+  public boolean isAutoPartitionMode() throws SQLException {
+    return get(Connection::isAutoPartitionMode);
+  }
+
+  @Override
+  public void setMaxPartitions(int maxPartitions) throws SQLException {
+    set(Connection::setMaxPartitions, maxPartitions);
+  }
+
+  @Override
+  public int getMaxPartitions() throws SQLException {
+    return get(Connection::getMaxPartitions);
+  }
+
+  @Override
+  public void setMaxPartitionedParallelism(int maxThreads) throws SQLException {
+    set(Connection::setMaxPartitionedParallelism, maxThreads);
+  }
+
+  @Override
+  public int getMaxPartitionedParallelism() throws SQLException {
+    return get(Connection::getMaxPartitionedParallelism);
+  }
+
   @SuppressWarnings("deprecation")
   private static final class JdbcToSpannerTransactionRetryListener
       implements com.google.cloud.spanner.connection.TransactionRetryListener {
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcPartitionedQueryResultSet.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPartitionedQueryResultSet.java
new file mode 100644
index 0000000..1d5ef66
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPartitionedQueryResultSet.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.jdbc;
+
+import com.google.cloud.spanner.connection.PartitionedQueryResultSet;
+import com.google.common.base.Preconditions;
+import java.sql.Statement;
+
+/**
+ * {@link java.sql.ResultSet} implementation that is returned for queries that are executed with
+ * `RUN PARTITIONED QUERY ...`.
+ */
+class JdbcPartitionedQueryResultSet extends JdbcResultSet
+    implements CloudSpannerJdbcPartitionedQueryResultSet {
+  static JdbcPartitionedQueryResultSet of(
+      Statement statement, PartitionedQueryResultSet resultSet) {
+    return new JdbcPartitionedQueryResultSet(
+        Preconditions.checkNotNull(statement), Preconditions.checkNotNull(resultSet));
+  }
+
+  private final PartitionedQueryResultSet partitionedQueryResultSet;
+
+  private JdbcPartitionedQueryResultSet(Statement statement, PartitionedQueryResultSet resultSet) {
+    super(statement, resultSet);
+    this.partitionedQueryResultSet = resultSet;
+  }
+
+  @Override
+  public int getNumPartitions() {
+    return partitionedQueryResultSet.getNumPartitions();
+  }
+
+  @Override
+  public int getParallelism() {
+    return partitionedQueryResultSet.getParallelism();
+  }
+}
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
index 3ab302c..518807d 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
@@ -16,7 +16,9 @@
 
 package com.google.cloud.spanner.jdbc;
 
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 import com.google.cloud.spanner.ResultSets;
 import com.google.cloud.spanner.SpannerException;
@@ -33,10 +35,10 @@ import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 
 /** Implementation of {@link PreparedStatement} for Cloud Spanner. */
-class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
+class JdbcPreparedStatement extends AbstractJdbcPreparedStatement
+    implements CloudSpannerJdbcPreparedStatement {
   private static final char POS_PARAM_CHAR = '?';
   private final String sql;
-  private final String sqlWithoutComments;
   private final ParametersInfo parameters;
   private final ImmutableList<String> generatedKeysColumns;
 
@@ -46,9 +48,15 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
     super(connection);
     this.sql = sql;
     try {
-      this.sqlWithoutComments = parser.removeCommentsAndTrim(this.sql);
+      // The PostgreSQL parser allows comments to be present in the SQL string that is used to parse
+      // the query parameters.
+      String sqlForParameterExtraction =
+          getConnection().getDialect() == Dialect.POSTGRESQL
+              ? this.sql
+              : parser.removeCommentsAndTrim(this.sql);
       this.parameters =
-          parser.convertPositionalParametersToNamedParameters(POS_PARAM_CHAR, sqlWithoutComments);
+          parser.convertPositionalParametersToNamedParameters(
+              POS_PARAM_CHAR, sqlForParameterExtraction);
     } catch (SpannerException e) {
       throw JdbcSqlExceptionFactory.of(e);
     }
@@ -128,4 +136,37 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
       return rs.getMetaData();
     }
   }
+
+  @Override
+  public ResultSet partitionQuery(PartitionOptions partitionOptions, QueryOption... options)
+      throws SQLException {
+    return runWithStatementTimeout(
+        connection ->
+            JdbcResultSet.of(
+                this, connection.partitionQuery(createStatement(), partitionOptions, options)));
+  }
+
+  @Override
+  public ResultSet runPartition() throws SQLException {
+    return runWithStatementTimeout(
+        connection -> {
+          if (getParameters().getHighestIndex() < 1 || getParameters().getParameter(1) == null) {
+            throw JdbcSqlExceptionFactory.of(
+                "No query parameter has been set. runPartition() requires the partition ID to be set as a query parameter with index 1. Call PreparedStatement#setString(1, \"some-partition-id\") before calling runPartition().",
+                Code.FAILED_PRECONDITION);
+          }
+          String partitionId = getParameters().getParameter(1).toString();
+          return JdbcResultSet.of(this, connection.runPartition(partitionId));
+        });
+  }
+
+  @Override
+  public CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
+      PartitionOptions partitionOptions, QueryOption... options) throws SQLException {
+    return runWithStatementTimeout(
+        connection ->
+            JdbcPartitionedQueryResultSet.of(
+                this,
+                connection.runPartitionedQuery(createStatement(), partitionOptions, options)));
+  }
 }
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
index 94b4604..4152fa5 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
@@ -21,6 +21,7 @@ import com.google.cloud.spanner.Struct;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.Type.Code;
 import com.google.cloud.spanner.Value;
+import com.google.cloud.spanner.connection.PartitionedQueryResultSet;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import java.io.ByteArrayInputStream;
@@ -59,9 +60,13 @@ class JdbcResultSet extends AbstractJdbcResultSet {
   }
 
   static JdbcResultSet of(Statement statement, com.google.cloud.spanner.ResultSet resultSet) {
-    Preconditions.checkNotNull(statement);
-    Preconditions.checkNotNull(resultSet);
-    return new JdbcResultSet(statement, resultSet);
+    // Return a JDBC version of a PartitionedQueryResultSet if the Cloud Spanner Java client
+    // returned a PartitionedQueryResultSet.
+    if (resultSet instanceof PartitionedQueryResultSet) {
+      return JdbcPartitionedQueryResultSet.of(statement, (PartitionedQueryResultSet) resultSet);
+    }
+    return new JdbcResultSet(
+        Preconditions.checkNotNull(statement), Preconditions.checkNotNull(resultSet));
   }
 
   /**
@@ -123,7 +128,7 @@ class JdbcResultSet extends AbstractJdbcResultSet {
   private boolean nextCalledForMetaDataResult = false;
   private long currentRow = 0L;
 
-  private JdbcResultSet(Statement statement, com.google.cloud.spanner.ResultSet spanner) {
+  JdbcResultSet(Statement statement, com.google.cloud.spanner.ResultSet spanner) {
     super(spanner);
     this.statement = statement;
   }
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
index 6c6a45a..60cc3e5 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
@@ -19,6 +19,8 @@ package com.google.cloud.spanner.jdbc;
 import static com.google.cloud.spanner.jdbc.JdbcConnection.NO_GENERATED_KEY_COLUMNS;
 
 import com.google.cloud.spanner.Options;
+import com.google.cloud.spanner.Options.QueryOption;
+import com.google.cloud.spanner.PartitionOptions;
 import com.google.cloud.spanner.ResultSets;
 import com.google.cloud.spanner.SpannerBatchUpdateException;
 import com.google.cloud.spanner.SpannerException;
@@ -40,7 +42,7 @@ import java.util.List;
 import java.util.stream.Collectors;
 
 /** Implementation of {@link java.sql.Statement} for Google Cloud Spanner. */
-class JdbcStatement extends AbstractJdbcStatement {
+class JdbcStatement extends AbstractJdbcStatement implements CloudSpannerJdbcStatement {
   static final ImmutableList<String> ALL_COLUMNS = ImmutableList.of("*");
 
   enum BatchType {
@@ -79,11 +81,7 @@ class JdbcStatement extends AbstractJdbcStatement {
 
   private int executeUpdate(String sql, ImmutableList<String> generatedKeysColumns)
       throws SQLException {
-    long result = executeLargeUpdate(sql, generatedKeysColumns);
-    if (result > Integer.MAX_VALUE) {
-      throw JdbcSqlExceptionFactory.of("update count too large: " + result, Code.OUT_OF_RANGE);
-    }
-    return (int) result;
+    return checkedCast(executeLargeUpdate(sql, generatedKeysColumns));
   }
 
   /**
@@ -300,11 +298,7 @@ class JdbcStatement extends AbstractJdbcStatement {
   @Override
   public int getUpdateCount() throws SQLException {
     checkClosed();
-    if (currentUpdateCount > Integer.MAX_VALUE) {
-      throw JdbcSqlExceptionFactory.of(
-          "update count too large: " + currentUpdateCount, Code.OUT_OF_RANGE);
-    }
-    return (int) currentUpdateCount;
+    return checkedCast(currentUpdateCount);
   }
 
   /**
@@ -481,12 +475,7 @@ class JdbcStatement extends AbstractJdbcStatement {
   int[] convertUpdateCounts(long[] updateCounts) throws SQLException {
     int[] res = new int[updateCounts.length];
     for (int index = 0; index < updateCounts.length; index++) {
-      if (updateCounts[index] > Integer.MAX_VALUE) {
-        throw JdbcSqlExceptionFactory.of(
-            String.format("Update count too large for int: %d", updateCounts[index]),
-            Code.OUT_OF_RANGE);
-      }
-      res[index] = (int) updateCounts[index];
+      res[index] = checkedCast(updateCounts[index]);
     }
     return res;
   }
@@ -598,4 +587,29 @@ class JdbcStatement extends AbstractJdbcStatement {
   static boolean isNullOrEmpty(String[] columnNames) {
     return columnNames == null || columnNames.length == 0;
   }
+
+  @Override
+  public ResultSet partitionQuery(
+      String query, PartitionOptions partitionOptions, QueryOption... options) throws SQLException {
+    return runWithStatementTimeout(
+        connection ->
+            JdbcResultSet.of(
+                this, connection.partitionQuery(Statement.of(query), partitionOptions, options)));
+  }
+
+  @Override
+  public ResultSet runPartition(String encodedPartitionId) throws SQLException {
+    return runWithStatementTimeout(
+        connection -> JdbcResultSet.of(this, connection.runPartition(encodedPartitionId)));
+  }
+
+  @Override
+  public CloudSpannerJdbcPartitionedQueryResultSet runPartitionedQuery(
+      String query, PartitionOptions partitionOptions, QueryOption... options) throws SQLException {
+    return runWithStatementTimeout(
+        connection ->
+            JdbcPartitionedQueryResultSet.of(
+                this,
+                connection.runPartitionedQuery(Statement.of(query), partitionOptions, options)));
+  }
 }
diff --git a/versions.txt b/versions.txt
index 2da2362..3ea0291 100644
--- a/versions.txt
+++ b/versions.txt
@@ -1,4 +1,4 @@
 # Format:
 # module:released-version:current-version
 
-google-cloud-spanner-jdbc:2.12.1:2.12.2-SNAPSHOT
+google-cloud-spanner-jdbc:2.13.0:2.13.1-SNAPSHOT


Test patch:
diff --git a/.github/workflows/integration-tests-against-emulator.yaml b/.github/workflows/integration-tests-against-emulator.yaml
index 8ee7c4d..f299be4 100644
--- a/.github/workflows/integration-tests-against-emulator.yaml
+++ b/.github/workflows/integration-tests-against-emulator.yaml
@@ -16,7 +16,7 @@ jobs:
           - 9020:9020
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4
     - uses: stCarolas/setup-maven@v4
       with:
         maven-version: 3.8.1
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
index f74de66..65a54a8 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
@@ -17,39 +17,97 @@
 package com.google.cloud.spanner.jdbc;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.cloud.spanner.SessionPoolOptions;
 import com.google.cloud.spanner.connection.AbstractMockServerTest;
 import com.google.cloud.spanner.connection.ConnectionOptions;
+import com.google.cloud.spanner.connection.SpannerPool;
+import com.google.longrunning.Operation;
+import com.google.protobuf.Any;
+import com.google.protobuf.Empty;
 import com.google.protobuf.ListValue;
 import com.google.protobuf.Value;
 import com.google.rpc.Code;
+import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
 import com.google.spanner.v1.ResultSetMetadata;
+import com.google.spanner.v1.ResultSetStats;
 import com.google.spanner.v1.StructType;
 import com.google.spanner.v1.StructType.Field;
 import com.google.spanner.v1.Type;
 import com.google.spanner.v1.TypeCode;
 import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
 
 /**
  * Test class for verifying that the methods execute, executeQuery, and executeUpdate work as
- * intended.
+ * intended. It also verifies that they always also include any comments in the statement for the
+ * PostgreSQL dialect, as these may contain hints.
  */
-@RunWith(JUnit4.class)
+@RunWith(Parameterized.class)
 public class ExecuteMockServerTest extends AbstractMockServerTest {
-  private static final String QUERY = "select * from my_table";
+
+  @Parameters(name = "dialect = {0}")
+  public static Object[] parameters() {
+    return Dialect.values();
+  }
+
+  @Parameter public Dialect dialect;
+
+  private static final String DDL = "create table my_table";
+  private static final long LARGE_UPDATE_COUNT = 2L * Integer.MAX_VALUE;
+
+  private String query;
+  private String dml;
+  private String largeDml;
+  private String dmlReturning;
 
   @Before
   public void setupResults() {
+    query =
+        dialect == Dialect.POSTGRESQL
+            ? "/*@ lock_scanned_ranges = exclusive */ select * from my_table"
+            : "select * from my_table";
+    dml =
+        dialect == Dialect.POSTGRESQL
+            ? "/*@ lock_scanned_ranges = exclusive */ insert into my_table (id, value) values (1, 'One')"
+            : "insert into my_table (id, value) values (1, 'One')";
+    String DML_THEN_RETURN_ID =
+        dml + (dialect == Dialect.POSTGRESQL ? "\nRETURNING \"id\"" : "\nTHEN RETURN `id`");
+    largeDml =
+        dialect == Dialect.POSTGRESQL
+            ? "/*@ lock_scanned_ranges = exclusive */ update my_table set value='new value' where true"
+            : "update my_table set value='new value' where true";
+    String LARGE_DML_THEN_RETURN_ID =
+        largeDml + (dialect == Dialect.POSTGRESQL ? "\nRETURNING \"id\"" : "\nTHEN RETURN `id`");
+    dmlReturning =
+        dialect == Dialect.POSTGRESQL
+            ? "/*@ lock_scanned_ranges = exclusive */ insert into my_table (id, value) values (1, 'One') RETURNING *"
+            : "insert into my_table (id, value) values (1, 'One') THEN RETURN *";
+
+    // This forces a refresh of the Spanner instance that is used for a connection, which again is
+    // needed in order to refresh the dialect of the database.
+    SpannerPool.closeSpannerPool();
+    mockSpanner.putStatementResult(StatementResult.detectDialectResult(dialect));
     super.setupResults();
+
     com.google.spanner.v1.ResultSet resultSet =
         com.google.spanner.v1.ResultSet.newBuilder()
             .setMetadata(
@@ -74,8 +132,58 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
                     .addValues(Value.newBuilder().setStringValue("One").build())
                     .build())
             .build();
+    com.google.spanner.v1.ResultSet returnIdResultSet =
+        com.google.spanner.v1.ResultSet.newBuilder()
+            .setMetadata(
+                ResultSetMetadata.newBuilder()
+                    .setRowType(
+                        StructType.newBuilder()
+                            .addFields(
+                                Field.newBuilder()
+                                    .setType(Type.newBuilder().setCode(TypeCode.INT64).build())
+                                    .setName("id")
+                                    .build())
+                            .build())
+                    .build())
+            .addRows(
+                ListValue.newBuilder()
+                    .addValues(Value.newBuilder().setStringValue("1").build())
+                    .build())
+            .build();
+    mockSpanner.putStatementResult(
+        StatementResult.query(com.google.cloud.spanner.Statement.of(query), resultSet));
+    mockSpanner.putStatementResult(
+        StatementResult.update(com.google.cloud.spanner.Statement.of(dml), 1L));
     mockSpanner.putStatementResult(
-        StatementResult.query(com.google.cloud.spanner.Statement.of(QUERY), resultSet));
+        StatementResult.update(
+            com.google.cloud.spanner.Statement.of(largeDml), LARGE_UPDATE_COUNT));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            com.google.cloud.spanner.Statement.of(dmlReturning),
+            resultSet
+                .toBuilder()
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            com.google.cloud.spanner.Statement.of(DML_THEN_RETURN_ID),
+            returnIdResultSet
+                .toBuilder()
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            com.google.cloud.spanner.Statement.of(LARGE_DML_THEN_RETURN_ID),
+            returnIdResultSet
+                .toBuilder()
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(LARGE_UPDATE_COUNT).build())
+                .build()));
+    mockDatabaseAdmin.addResponse(
+        Operation.newBuilder()
+            .setDone(true)
+            .setResponse(Any.pack(Empty.getDefaultInstance()))
+            .setMetadata(Any.pack(UpdateDatabaseDdlMetadata.getDefaultInstance()))
+            .build());
   }
 
   private String createUrl() {
@@ -84,6 +192,501 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
         getPort(), "proj", "inst", "db");
   }
 
+  private Connection createConnection() throws SQLException {
+    return DriverManager.getConnection(createUrl());
+  }
+
+  @Test
+  public void testStatementExecuteQuery() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      try (ResultSet resultSet = statement.executeQuery(query)) {
+        verifyResultSet(resultSet);
+      }
+      try (ResultSet resultSet = statement.executeQuery(dmlReturning)) {
+        verifyResultSet(resultSet);
+      }
+      verifyException(() -> statement.executeQuery(dml));
+      verifyException(() -> statement.executeQuery(largeDml));
+      verifyException(() -> statement.executeQuery(DDL));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteUpdate() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeUpdate(dml));
+      assertEquals(0, statement.executeUpdate(DDL));
+      verifyOverflow(() -> statement.executeUpdate(largeDml));
+      verifyException(() -> statement.executeUpdate(query));
+      verifyException(() -> statement.executeUpdate(dmlReturning));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteUpdateReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      assertEquals(1, statement.executeUpdate(dml, Statement.NO_GENERATED_KEYS));
+      assertEquals(0, statement.executeUpdate(DDL, Statement.NO_GENERATED_KEYS));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeUpdate(query, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeUpdate(dmlReturning, Statement.NO_GENERATED_KEYS));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteUpdateReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeUpdate(dml, new String[] {"id"}));
+      assertEquals(0, statement.executeUpdate(DDL, new String[] {"id"}));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, new String[] {"id"}));
+      verifyException(
+          () -> statement.executeUpdate(query, new String[] {"id"}), Code.FAILED_PRECONDITION);
+      assertEquals(1, statement.executeUpdate(dmlReturning, new String[] {"id"}));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteUpdateReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeUpdate(dml, new int[] {1}));
+      assertEquals(0, statement.executeUpdate(DDL, new int[] {1}));
+      verifyOverflow(() -> statement.executeUpdate(largeDml, new int[] {1}));
+      verifyException(() -> statement.executeUpdate(query, new int[] {1}));
+      verifyException(() -> statement.executeUpdate(dmlReturning, new int[] {1}));
+    }
+  }
+
+  @Test
+  public void testStatementLargeExecuteUpdate() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1L, statement.executeLargeUpdate(dml));
+      assertEquals(0L, statement.executeLargeUpdate(DDL));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml));
+      verifyException(() -> statement.executeLargeUpdate(query));
+      verifyException(() -> statement.executeLargeUpdate(dmlReturning));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteLargeUpdateReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      assertEquals(1, statement.executeLargeUpdate(dml, Statement.NO_GENERATED_KEYS));
+      assertEquals(0, statement.executeLargeUpdate(DDL, Statement.NO_GENERATED_KEYS));
+      assertEquals(
+          LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, Statement.NO_GENERATED_KEYS));
+      verifyException(() -> statement.executeLargeUpdate(query, Statement.NO_GENERATED_KEYS));
+      verifyException(
+          () -> statement.executeLargeUpdate(dmlReturning, Statement.NO_GENERATED_KEYS));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteLargeUpdateReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeLargeUpdate(dml, new String[] {"id"}));
+      assertEquals(0, statement.executeLargeUpdate(DDL, new String[] {"id"}));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, new String[] {"id"}));
+      verifyException(
+          () -> statement.executeLargeUpdate(query, new String[] {"id"}), Code.FAILED_PRECONDITION);
+      assertEquals(1L, statement.executeLargeUpdate(dmlReturning, new String[] {"id"}));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteLargeUpdateReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      assertEquals(1, statement.executeLargeUpdate(dml, new int[] {1}));
+      assertEquals(0, statement.executeLargeUpdate(DDL, new int[] {1}));
+      assertEquals(LARGE_UPDATE_COUNT, statement.executeLargeUpdate(largeDml, new int[] {1}));
+      verifyException(() -> statement.executeLargeUpdate(query, new int[] {1}));
+      verifyException(() -> statement.executeLargeUpdate(dmlReturning, new int[] {1}));
+    }
+  }
+
+  @Test
+  public void testStatementExecute() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      verifyUpdateCount(statement, () -> statement.execute(dml), 1L);
+      verifyUpdateCount(statement, () -> statement.execute(largeDml), LARGE_UPDATE_COUNT);
+      verifyUpdateCount(statement, () -> statement.execute(DDL), Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(query));
+      verifyResultSet(statement, () -> statement.execute(dmlReturning));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      verifyUpdateCount(statement, () -> statement.execute(dml, Statement.NO_GENERATED_KEYS), 1L);
+      verifyUpdateCount(
+          statement,
+          () -> statement.execute(largeDml, Statement.NO_GENERATED_KEYS),
+          LARGE_UPDATE_COUNT);
+      verifyUpdateCount(
+          statement,
+          () -> statement.execute(DDL, Statement.NO_GENERATED_KEYS),
+          Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(query, Statement.NO_GENERATED_KEYS));
+      verifyResultSet(
+          statement, () -> statement.execute(dmlReturning, Statement.NO_GENERATED_KEYS));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      verifyUpdateCount(statement, () -> statement.execute(dml, new String[] {"id"}), 1L);
+      verifyUpdateCount(
+          statement, () -> statement.execute(largeDml, new String[] {"id"}), LARGE_UPDATE_COUNT);
+      verifyUpdateCount(
+          statement, () -> statement.execute(DDL, new String[] {"id"}), Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(query, new String[] {"id"}));
+      verifyResultSet(statement, () -> statement.execute(dmlReturning, new String[] {"id"}));
+    }
+  }
+
+  @Test
+  public void testStatementExecuteReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection();
+        Statement statement = connection.createStatement()) {
+      verifyUpdateCount(statement, () -> statement.execute(dml, new int[] {1}), 1L);
+      verifyUpdateCount(
+          statement, () -> statement.execute(largeDml, new int[] {1}), LARGE_UPDATE_COUNT);
+      verifyUpdateCount(
+          statement, () -> statement.execute(DDL, new int[] {1}), Statement.SUCCESS_NO_INFO);
+      verifyResultSet(statement, () -> statement.execute(query, new int[] {1}));
+      verifyResultSet(statement, () -> statement.execute(dmlReturning, new int[] {1}));
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteQuery() throws SQLException {
+    try (Connection connection = createConnection()) {
+      try (ResultSet resultSet = connection.prepareStatement(query).executeQuery()) {
+        verifyResultSet(resultSet);
+      }
+      try (ResultSet resultSet = connection.prepareStatement(dmlReturning).executeQuery()) {
+        verifyResultSet(resultSet);
+      }
+      verifyException(() -> connection.prepareStatement(dml).executeQuery());
+      verifyException(() -> connection.prepareStatement(largeDml).executeQuery());
+      verifyException(() -> connection.prepareStatement(DDL).executeQuery());
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdate() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml).executeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL).executeUpdate());
+      verifyOverflow(() -> connection.prepareStatement(largeDml).executeUpdate());
+      verifyException(() -> connection.prepareStatement(query).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(dmlReturning).executeUpdate(), Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdateReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      assertEquals(
+          1, connection.prepareStatement(dml, Statement.NO_GENERATED_KEYS).executeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL,
+      // Statement.NO_GENERATED_KEYS).executeUpdate());
+      verifyOverflow(
+          () -> connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(query, Statement.NO_GENERATED_KEYS).executeUpdate());
+      verifyException(
+          () ->
+              connection
+                  .prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS)
+                  .executeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdateReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml, new String[] {"id"}).executeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL, new String[] {"id"}).executeUpdate());
+      verifyOverflow(
+          () -> connection.prepareStatement(largeDml, new String[] {"id"}).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(query, new String[] {"id"}).executeUpdate(),
+          Code.FAILED_PRECONDITION);
+      assertEquals(
+          1, connection.prepareStatement(dmlReturning, new String[] {"id"}).executeUpdate());
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteUpdateReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml, new int[] {1}).executeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL, new int[] {1}).executeUpdate());
+      verifyOverflow(() -> connection.prepareStatement(largeDml, new int[] {1}).executeUpdate());
+      verifyException(() -> connection.prepareStatement(query, new int[] {1}).executeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(dmlReturning, new int[] {1}).executeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementLargeExecuteUpdate() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1L, connection.prepareStatement(dml).executeLargeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0L, connection.prepareStatement(DDL).executeLargeUpdate());
+      assertEquals(LARGE_UPDATE_COUNT, connection.prepareStatement(largeDml).executeLargeUpdate());
+      verifyException(() -> connection.prepareStatement(query).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(dmlReturning).executeLargeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteLargeUpdateReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      assertEquals(
+          1, connection.prepareStatement(dml, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL,
+      // Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      assertEquals(
+          LARGE_UPDATE_COUNT,
+          connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      verifyException(
+          () ->
+              connection.prepareStatement(query, Statement.NO_GENERATED_KEYS).executeLargeUpdate());
+      verifyException(
+          () ->
+              connection
+                  .prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS)
+                  .executeLargeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteLargeUpdateReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml, new String[] {"id"}).executeLargeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL, new String[]
+      // {"id"}).executeLargeUpdate());
+      assertEquals(
+          LARGE_UPDATE_COUNT,
+          connection.prepareStatement(largeDml, new String[] {"id"}).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(query, new String[] {"id"}).executeLargeUpdate(),
+          Code.FAILED_PRECONDITION);
+      assertEquals(
+          1L, connection.prepareStatement(dmlReturning, new String[] {"id"}).executeLargeUpdate());
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteLargeUpdateReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertEquals(1, connection.prepareStatement(dml, new int[] {1}).executeLargeUpdate());
+      // TODO: Enable the next statement once PreparedStatement supports executing DDL through the
+      //       executeUpdate() method.
+      // assertEquals(0, connection.prepareStatement(DDL, new int[] {1}).executeLargeUpdate());
+      assertEquals(
+          LARGE_UPDATE_COUNT,
+          connection.prepareStatement(largeDml, new int[] {1}).executeLargeUpdate());
+      verifyException(() -> connection.prepareStatement(query, new int[] {1}).executeLargeUpdate());
+      verifyException(
+          () -> connection.prepareStatement(dmlReturning, new int[] {1}).executeLargeUpdate(),
+          Code.INVALID_ARGUMENT);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecute() throws SQLException {
+    try (Connection connection = createConnection()) {
+      verifyPreparedUpdateCount(connection.prepareStatement(dml), PreparedStatement::execute, 1L);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(largeDml), PreparedStatement::execute, LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL), PreparedStatement::execute, Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(connection.prepareStatement(query), PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(dmlReturning), PreparedStatement::execute);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteReturnGeneratedKeys() throws SQLException {
+    try (Connection connection = createConnection()) {
+      // TODO: Add tests for RETURN_GENERATED_KEYS when that is supported.
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(dml, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute,
+          1L);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(largeDml, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute,
+          LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute,
+          Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(
+          connection.prepareStatement(query, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(dmlReturning, Statement.NO_GENERATED_KEYS),
+          PreparedStatement::execute);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteReturnColumnNames() throws SQLException {
+    try (Connection connection = createConnection()) {
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(dml, new String[] {"id"}), PreparedStatement::execute, 1L);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(largeDml, new String[] {"id"}),
+          PreparedStatement::execute,
+          LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL, new String[] {"id"}),
+          PreparedStatement::execute,
+          Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(
+          connection.prepareStatement(query, new String[] {"id"}), PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(dmlReturning, new String[] {"id"}),
+          PreparedStatement::execute);
+    }
+  }
+
+  @Test
+  public void testPreparedStatementExecuteReturnColumnIndexes() throws SQLException {
+    try (Connection connection = createConnection()) {
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(dml, new int[] {1}), PreparedStatement::execute, 1L);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(largeDml, new int[] {1}),
+          PreparedStatement::execute,
+          LARGE_UPDATE_COUNT);
+      verifyPreparedUpdateCount(
+          connection.prepareStatement(DDL, new int[] {1}),
+          PreparedStatement::execute,
+          Statement.SUCCESS_NO_INFO);
+      verifyPreparedResultSet(
+          connection.prepareStatement(query, new int[] {1}), PreparedStatement::execute);
+      verifyPreparedResultSet(
+          connection.prepareStatement(dmlReturning, new int[] {1}), PreparedStatement::execute);
+    }
+  }
+
+  private void verifyResultSet(ResultSet resultSet) throws SQLException {
+    assertNotNull(resultSet.getMetaData());
+    assertEquals(2, resultSet.getMetaData().getColumnCount());
+    assertTrue(resultSet.next());
+    assertFalse(resultSet.next());
+  }
+
+  private void verifyException(ThrowingRunnable runnable) {
+    verifyException(runnable, Code.INVALID_ARGUMENT);
+  }
+
+  private void verifyOverflow(ThrowingRunnable runnable) {
+    verifyException(runnable, Code.OUT_OF_RANGE);
+  }
+
+  private void verifyException(ThrowingRunnable runnable, Code code) {
+    SQLException exception = assertThrows(SQLException.class, runnable);
+    assertTrue(exception instanceof JdbcSqlException);
+    JdbcSqlException sqlException = (JdbcSqlException) exception;
+    assertEquals(code, sqlException.getCode());
+  }
+
+  interface ThrowingFunction<T> {
+    T apply() throws SQLException;
+  }
+
+  interface ThrowingPreparedFunction<T> {
+    T apply(PreparedStatement statement) throws SQLException;
+  }
+
+  private void verifyPreparedUpdateCount(
+      PreparedStatement statement, ThrowingPreparedFunction<Boolean> function, long updateCount)
+      throws SQLException {
+    verifyUpdateCount(statement, () -> function.apply(statement), updateCount);
+  }
+
+  private void verifyUpdateCount(
+      Statement statement, ThrowingFunction<Boolean> function, long updateCount)
+      throws SQLException {
+    assertFalse(function.apply());
+    if (updateCount > Integer.MAX_VALUE) {
+      verifyOverflow(statement::getUpdateCount);
+    } else {
+      assertEquals((int) updateCount, statement.getUpdateCount());
+    }
+    assertEquals(updateCount, statement.getLargeUpdateCount());
+    assertNull(statement.getResultSet());
+    assertFalse(statement.getMoreResults());
+    assertEquals(-1, statement.getUpdateCount());
+    assertEquals(-1L, statement.getLargeUpdateCount());
+  }
+
+  private void verifyPreparedResultSet(
+      PreparedStatement statement, ThrowingPreparedFunction<Boolean> function) throws SQLException {
+    verifyResultSet(statement, () -> function.apply(statement));
+  }
+
+  private void verifyResultSet(Statement statement, ThrowingFunction<Boolean> function)
+      throws SQLException {
+    assertTrue(function.apply());
+    assertEquals(-1, statement.getUpdateCount());
+    assertEquals(-1L, statement.getLargeUpdateCount());
+    assertNotNull(statement.getResultSet());
+    try (ResultSet resultSet = statement.getResultSet()) {
+      assertTrue(resultSet.next());
+      assertFalse(resultSet.next());
+    }
+
+    assertFalse(statement.getMoreResults());
+    assertEquals(-1, statement.getUpdateCount());
+    assertEquals(-1L, statement.getLargeUpdateCount());
+  }
+
   @Test
   public void testInvalidExecuteUpdate_shouldNotLeakSession() throws SQLException {
     int maxSessions = 1;
@@ -103,7 +706,7 @@ public class ExecuteMockServerTest extends AbstractMockServerTest {
       for (int i = 0; i < (maxSessions + 1); i++) {
         SQLException exception =
             assertThrows(
-                SQLException.class, () -> connection.createStatement().executeUpdate(QUERY));
+                SQLException.class, () -> connection.createStatement().executeUpdate(query));
         assertTrue(exception instanceof JdbcSqlException);
         JdbcSqlException jdbcSqlException = (JdbcSqlException) exception;
         // This would be RESOURCE_EXHAUSTED if the query leaked a session.
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/PartitionedQueryMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/PartitionedQueryMockServerTest.java
new file mode 100644
index 0000000..7e4e749
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/jdbc/PartitionedQueryMockServerTest.java
@@ -0,0 +1,502 @@
+/*
+ * 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(
+                      "/*\n"
+                          + " * Copyright 2022 Google LLC\n"
+                          + " *\n"
+                          + " * Licensed under the Apache License, Version 2.0 (the \"License\");\n"
+                          + " * you may not use this file except in compliance with the License.\n"
+                          + " * You may obtain a copy of the License at\n"
+                          + " *\n"
+                          + " *       http://www.apache.org/licenses/LICENSE-2.0\n"
+                          + " *\n"
+                          + " * Unless required by applicable law or agreed to in writing, software\n"
+                          + " * distributed under the License is distributed on an \"AS IS\" BASIS,\n"
+                          + " * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n"
+                          + " * See the License for the specific language governing permissions and\n"
+                          + " * limitations under the License.\n"
+                          + " */\n"
+                          + "\n"
+                          + "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());
+      }
+    }
+  }
+}