Repository: GoogleCloudPlatform/pgadapter
Pull request #792 (2023-04-19 05:42:00)

feat: automatically add LIMIT clause
Automatically add a LIMIT clause to queries that use an OFFSET clause with a query parameter and no LIMIT clause. That is currently not supported in Cloud Spanner PostgreSQL databases.

This type of queries are generated by Prisma.

codecov[bot]: ## [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/792?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#792](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/792?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (7a4acf4) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/aa75a20b62723bc1d7acf19538d6e4479843482e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (aa75a20) will **decrease** coverage by `0.05%`.
> The diff coverage is `80.55%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #792      +/-   ##
========================================================
- Coverage                 89.54%   89.50%   -0.05%     
- Complexity                 2354     2366      +12     
========================================================
  Files                       129      129              
  Lines                      7767     7802      +35     
  Branches                   1112     1125      +13     
========================================================
+ Hits                       6955     6983      +28     
- Misses                      564      565       +1     
- Partials                    248      254       +6     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `89.50% <80.55%> (-0.05%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/792?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...oud/spanner/pgadapter/statements/SimpleParser.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/792?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvU2ltcGxlUGFyc2VyLmphdmE=) | `91.50% <77.41%> (-0.90%)` | :arrow_down: |
| [.../cloud/spanner/pgadapter/session/SessionState.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/792?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3Nlc3Npb24vU2Vzc2lvblN0YXRlLmphdmE=) | `95.74% <100.00%> (+0.02%)` | :arrow_up: |
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/792?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `93.69% <100.00%> (+0.03%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)


Patch:
diff --git a/docs/psycopg3.md b/docs/psycopg3.md
index 96cabd60..b7b120cc 100644
--- a/docs/psycopg3.md
+++ b/docs/psycopg3.md
@@ -75,13 +75,9 @@ psycopg.errors.RaiseException: Unknown statement: DECLARE "my_cursor" CURSOR FOR
 
 ### Nested Transactions
 `psycopg3` implements [nested transactions](https://www.psycopg.org/psycopg3/docs/basic/transactions.html#nested-transactions)
-using `SAVEPOINT`. This feature is currently not supported with PGAdapter.
-
-Creating a nested transaction in `psycopg3` with PGAdapter will cause an error like the following:
-
-```
-psycopg.errors.RaiseException: Unknown statement: SAVEPOINT "_pg3_2"
-```
+using `SAVEPOINT`. Rolling back to a `SAVEPOINT` can fail if the transaction contained at least one
+query that called a volatile function or if the underlying data that has been accessed by the
+transaction has been modified by another transaction.
 
 ## Performance Considerations
 
diff --git a/pom.xml b/pom.xml
index a337718d..39317efa 100644
--- a/pom.xml
+++ b/pom.xml
@@ -80,7 +80,7 @@
     <dependency>
       <groupId>org.checkerframework</groupId>
       <artifactId>checker-qual</artifactId>
-      <version>3.33.0</version>
+      <version>3.34.0</version>
     </dependency>
     <dependency>
       <groupId>com.google.cloud</groupId>
diff --git a/samples/python/sqlalchemy-sample/README.md b/samples/python/sqlalchemy-sample/README.md
index a736b768..84c4f368 100644
--- a/samples/python/sqlalchemy-sample/README.md
+++ b/samples/python/sqlalchemy-sample/README.md
@@ -83,7 +83,7 @@ The following limitations are currently known:
 | DDL Transactions             | Cloud Spanner does not support DDL statements in a transaction. Add `?options=-c spanner.ddl_transaction_mode=AutocommitExplicitTransaction` to your connection string to automatically convert DDL transactions to [non-atomic DDL batches](../../../docs/ddl.md). |
 | Generated primary keys       | Manually assign a value to the primary key column in your code. The recommended primary key type is a random UUID. Sequences / SERIAL / IDENTITY columns are currently not supported.                                                                               |
 | INSERT ... ON CONFLICT       | `INSERT ... ON CONFLICT` is not supported.                                                                                                                                                                                                                          |
-| SAVEPOINT                    | Nested transactions and savepoints are not supported.                                                                                                                                                                                                               |
+| SAVEPOINT                    | Rolling back to a `SAVEPOINT` can fail if the transaction contained at least one query that called a volatile function.                                                                                                                                             |
 | SELECT ... FOR UPDATE        | `SELECT ... FOR UPDATE` is not supported.                                                                                                                                                                                                                           |
 | Server side cursors          | Server side cursors are currently not supported.                                                                                                                                                                                                                    |
 | Transaction isolation level  | Only SERIALIZABLE and AUTOCOMMIT are supported. `postgresql_readonly=True` is also supported. It is recommended to use either autocommit or read-only for workloads that only read data and/or that do not need to be atomic to get the best possible performance.  |
@@ -116,9 +116,9 @@ or https://docs.sqlalchemy.org/en/14/dialects/postgresql.html#sqlalchemy.dialect
 will fail.
 
 ### SAVEPOINT - Nested transactions
-`SAVEPOINT`s are not supported by Cloud Spanner. Nested transactions in SQLAlchemy are translated to
-savepoints and are therefore not supported. Trying to use `Session.begin_nested()`
-(https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.Session.begin_nested) will fail.
+Rolling back to a `SAVEPOINT` can fail if the transaction contained at least one query that called a
+volatile function or if the underlying data that has been accessed by the transaction has been
+modified by another transaction.
 
 ### Locking - SELECT ... FOR UPDATE
 Locking clauses, like `SELECT ... FOR UPDATE`, are not supported (see also https://docs.sqlalchemy.org/en/20/orm/queryguide/query.html#sqlalchemy.orm.Query.with_for_update).
diff --git a/samples/python/sqlalchemy2-sample/README.md b/samples/python/sqlalchemy2-sample/README.md
index 9f0d2504..45528ea4 100644
--- a/samples/python/sqlalchemy2-sample/README.md
+++ b/samples/python/sqlalchemy2-sample/README.md
@@ -77,7 +77,7 @@ The following limitations are currently known:
 | DDL Transactions               | Cloud Spanner does not support DDL statements in a transaction. Add `?options=-c spanner.ddl_transaction_mode=AutocommitExplicitTransaction` to your connection string to automatically convert DDL transactions to [non-atomic DDL batches](../../../docs/ddl.md). |
 | Generated primary keys         | Manually assign a value to the primary key column in your code. The recommended primary key type is a random UUID. Sequences / SERIAL / IDENTITY columns are currently not supported.                                                                               |
 | INSERT ... ON CONFLICT         | `INSERT ... ON CONFLICT` is not supported.                                                                                                                                                                                                                          |
-| SAVEPOINT                      | Nested transactions and savepoints are not supported.                                                                                                                                                                                                               |
+| SAVEPOINT                      | Rolling back to a `SAVEPOINT` can fail if the transaction contained at least one query that called a volatile function.                                                                                                                                             |
 | SELECT ... FOR UPDATE          | `SELECT ... FOR UPDATE` is not supported.                                                                                                                                                                                                                           |
 | Server side cursors            | Server side cursors are currently not supported.                                                                                                                                                                                                                    |
 | Transaction isolation level    | Only SERIALIZABLE and AUTOCOMMIT are supported. `postgresql_readonly=True` is also supported. It is recommended to use either autocommit or read-only for workloads that only read data and/or that do not need to be atomic to get the best possible performance.  |
@@ -110,9 +110,9 @@ or https://docs.sqlalchemy.org/en/20/dialects/postgresql.html#sqlalchemy.dialect
 will fail.
 
 ### SAVEPOINT - Nested transactions
-`SAVEPOINT`s are not supported by Cloud Spanner. Nested transactions in SQLAlchemy are translated to
-savepoints and are therefore not supported. Trying to use `Session.begin_nested()`
-(https://docs.sqlalchemy.org/en/20/orm/session_api.html#sqlalchemy.orm.Session.begin_nested) will fail.
+Rolling back to a `SAVEPOINT` can fail if the transaction contained at least one query that called a
+volatile function or if the underlying data that has been accessed by the transaction has been
+modified by another transaction.
 
 ### Locking - SELECT ... FOR UPDATE
 Locking clauses, like `SELECT ... FOR UPDATE`, are not supported (see also https://docs.sqlalchemy.org/en/20/orm/queryguide/query.html#sqlalchemy.orm.Query.with_for_update).
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
index 94a0bcbc..5d62b171 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -34,6 +34,7 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.connection.ConnectionOptionsHelper;
+import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.error.SQLState;
@@ -240,6 +241,7 @@ public class ConnectionHandler extends Thread {
       spannerConnection.close();
       throw e;
     }
+    spannerConnection.setSavepointSupport(SavepointSupport.ENABLED);
     this.spannerConnection = spannerConnection;
     this.databaseId = connectionOptions.getDatabaseId();
     this.extendedQueryProtocolHandler = new ExtendedQueryProtocolHandler(this);
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/error/SQLState.java b/src/main/java/com/google/cloud/spanner/pgadapter/error/SQLState.java
index 89be8e33..21b20e01 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/error/SQLState.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/error/SQLState.java
@@ -149,6 +149,10 @@ public enum SQLState {
   InFailedSqlTransaction("25P02"),
   IdleInTransactionSessionTimeout("25P03"),
 
+  // Class 3B — Savepoint Exception
+  SavepointException("3B000"),
+  InvalidSavepointSpecification("3B001"),
+
   // Class 42 — Syntax Error or Access Rule Violation
   SyntaxErrorOrAccessRuleViolation("42000"),
   SyntaxError("42601"),
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java b/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
index d443e974..4d8c1919 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
@@ -392,6 +392,19 @@ public class SessionState {
     return getBoolSetting("spanner", "force_autocommit", false);
   }
 
+  /**
+   * Returns whether statements with an OFFSET clause that uses a parameter should be automatically
+   * appended with a LIMIT clause. The LIMIT clause will use the literal Long.MAX_VALUE for unbound
+   * statements, and Long.MAX_VALUE - offset for bound statements.
+   *
+   * <p>This method will be removed in the future.
+   */
+  // TODO: Remove when Cloud Spanner supports parametrized OFFSET clauses without a LIMIT clause.
+  @InternalApi
+  public boolean isAutoAddLimitClause() {
+    return getBoolSetting("spanner", "auto_add_limit_clause", false);
+  }
+
   /** Returns the current setting for replacing pg_catalog tables with common table expressions. */
   public boolean isReplacePgCatalogTables() {
     PGSetting setting = internalGet(toKey("spanner", "replace_pg_catalog_tables"), false);
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
index ccbd57d4..cb153959 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
@@ -295,7 +295,13 @@ public class BackendConnection {
     }
 
     Statement bindStatement(Statement statement) {
-      return statementBinder.apply(statement);
+      Statement boundStatement = statementBinder.apply(statement);
+      // Add a LIMIT clause to the statement if it contains an OFFSET clause that uses a query
+      // parameter and there is no existing LIMIT clause in the query.
+      if (sessionState.isAutoAddLimitClause()) {
+        return SimpleParser.addLimitIfParameterizedOffset(boundStatement);
+      }
+      return boundStatement;
     }
 
     StatementResult analyzeOrExecute(Statement statement) {
@@ -604,7 +610,16 @@ public class BackendConnection {
 
     @Override
     void execute() {
-      result.set(NO_RESULT);
+      try {
+        checkConnectionState();
+        spannerConnection.savepoint(savepointStatement.getSavepointName());
+        result.set(NO_RESULT);
+      } catch (Exception exception) {
+        PGException pgException =
+            PGException.newBuilder(exception).setSQLState(SQLState.SavepointException).build();
+        result.setException(pgException);
+        throw pgException;
+      }
     }
   }
 
@@ -623,7 +638,16 @@ public class BackendConnection {
 
     @Override
     void execute() {
-      result.set(NO_RESULT);
+      try {
+        checkConnectionState();
+        spannerConnection.releaseSavepoint(releaseStatement.getSavepointName());
+        result.set(NO_RESULT);
+      } catch (Exception exception) {
+        PGException pgException =
+            PGException.newBuilder(exception).setSQLState(SQLState.SavepointException).build();
+        result.setException(pgException);
+        throw pgException;
+      }
     }
   }
 
@@ -642,11 +666,15 @@ public class BackendConnection {
 
     @Override
     void execute() {
-      throw setAndReturn(
-          result,
-          PGExceptionFactory.newPGException(
-              "Statement 'ROLLBACK [WORK | TRANSACTION] TO [SAVEPOINT] savepoint_name' is not supported",
-              SQLState.FeatureNotSupported));
+      try {
+        spannerConnection.rollbackToSavepoint(rollbackToStatement.getSavepointName());
+        result.set(NO_RESULT);
+      } catch (Exception exception) {
+        PGException pgException =
+            PGException.newBuilder(exception).setSQLState(SQLState.SavepointException).build();
+        result.setException(pgException);
+        throw pgException;
+      }
     }
   }
 
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java
index 1171e968..f86d4e59 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java
@@ -24,6 +24,7 @@ import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
@@ -225,6 +226,52 @@ public class SimpleParser {
     return builder.build();
   }
 
+  /**
+   * Adds a LIMIT clause to the given statement if the statement contains a parameterized OFFSET
+   * clause and no LIMIT clause.
+   */
+  static Statement addLimitIfParameterizedOffset(Statement statement) {
+    String sqlLowerCase = statement.getSql().toLowerCase(Locale.ENGLISH);
+    // If there is no offset clause, then we know that we don't have to analyze any further.
+    if (!sqlLowerCase.contains("offset")) {
+      return statement;
+    }
+    SimpleParser parser = new SimpleParser(statement.getSql());
+    parser.parseExpressionUntilKeyword(ImmutableList.of("limit", "offset"), true, false, false);
+    if (parser.pos >= parser.getSql().length()) {
+      return statement;
+    }
+    if (parser.eatKeyword("limit")) {
+      // If the statement contains a LIMIT clause, then we're OK.
+      return statement;
+    }
+    String parameter;
+    if (parser.eatKeyword("offset") && (parameter = parser.readQueryParameter()) != null) {
+      // Check if we have a LIMIT clause after the OFFSET clause.
+      if (parser.peekKeyword("limit")) {
+        return statement;
+      }
+      // This is probably an invalid query. Do not modify it to prevent the user getting an error
+      // message for a query that they did not write.
+      if (parser.hasMoreTokens()) {
+        return statement;
+      }
+      // The statement contains an OFFSET clause using a query parameter and no LIMIT clause.
+      // Append a LIMIT clause equal to Long.MAX_VALUE / 2.
+      // We could also calculate it based on the actual OFFSET value, but that would make the
+      // query more prone to cache misses. Also, adding the LIMIT clause with a query parameter
+      // would structurally change the query, something that we also don't want.
+      long limit = Long.MAX_VALUE / 2;
+      return copyStatement(
+          statement,
+          parser.sql.substring(0, parser.pos)
+              + " limit "
+              + limit
+              + parser.sql.substring(parser.pos));
+    }
+    return statement;
+  }
+
   private String sql;
   private int pos;
 
@@ -661,6 +708,22 @@ public class SimpleParser {
     return false;
   }
 
+  String readQueryParameter() {
+    if (eat(true, false, "$")) {
+      int startPos = pos - 1;
+      if (pos == sql.length() || !Character.isDigit(sql.charAt(pos))) {
+        return null;
+      }
+      while (pos < sql.length() && Character.isDigit(sql.charAt(pos))) {
+        pos++;
+      }
+      if (Character.isDigit(sql.charAt(pos - 1))) {
+        return sql.substring(startPos, pos);
+      }
+    }
+    return null;
+  }
+
   boolean peekJoinKeyword() {
     return peekKeyword("join")
         || peekKeyword("left")


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
index fe642942..9b640e19 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
@@ -195,6 +195,28 @@ public class ITJdbcTest implements IntegrationTest {
     }
   }
 
+  @Test
+  public void testSelectParameterizedOffsetWithoutLimit() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {
+      connection.createStatement().execute("set spanner.auto_add_limit_clause=true");
+      try (PreparedStatement statement =
+          connection.prepareStatement("select * from numbers offset ?")) {
+        for (long offset : new long[] {0L, 1L}) {
+          statement.setLong(1, offset);
+          try (ResultSet resultSet = statement.executeQuery()) {
+            if (offset == 0L) {
+              assertTrue(resultSet.next());
+              assertEquals(1L, resultSet.getLong(1));
+              assertFalse(resultSet.next());
+            } else {
+              assertFalse(resultSet.next());
+            }
+          }
+        }
+      }
+    }
+  }
+
   @Test
   public void testSelectCurrentSchema() throws SQLException {
     try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
index 4c99de89..32f94612 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -81,6 +81,7 @@ import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
+import java.sql.Savepoint;
 import java.sql.Types;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
@@ -102,11 +103,9 @@ import org.junit.runners.Parameterized.Parameters;
 import org.postgresql.PGConnection;
 import org.postgresql.PGStatement;
 import org.postgresql.core.Oid;
-import org.postgresql.jdbc.PSQLSavepoint;
 import org.postgresql.jdbc.PgStatement;
 import org.postgresql.util.PGobject;
 import org.postgresql.util.PSQLException;
-import org.postgresql.util.PSQLState;
 
 @RunWith(Parameterized.class)
 public class JdbcMockServerTest extends AbstractMockServerTest {
@@ -1998,6 +1997,43 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test
+  public void testParameterizedOffsetWithoutLimit() throws SQLException {
+    // Add a result for the non-limited query that contains one row.
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.newBuilder("select * from foo offset $1").bind("p1").to(0L).build(),
+            SELECT1_RESULTSET));
+    // Add a result for the limited query that is empty.
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.newBuilder("select * from foo offset $1 limit 4611686018427387903")
+                .bind("p1")
+                .to(0L)
+                .build(),
+            EMPTY_RESULTSET));
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      for (boolean addLimit : new boolean[] {true, false}) {
+        connection.createStatement().execute("set spanner.auto_add_limit_clause=" + addLimit);
+        try (PreparedStatement statement =
+            connection.prepareStatement("select * from foo offset ?")) {
+          statement.setLong(1, 0);
+          try (ResultSet resultSet = statement.executeQuery()) {
+            // We should get the empty result set when the auto-limit feature is enabled.
+            if (addLimit) {
+              assertFalse(resultSet.next());
+            } else {
+              assertTrue(resultSet.next());
+              assertEquals(1L, resultSet.getLong(1));
+              assertFalse(resultSet.next());
+            }
+          }
+        }
+      }
+    }
+  }
+
   @Test
   public void testTwoQueries() throws SQLException {
     try (Connection connection = DriverManager.getConnection(createUrl())) {
@@ -4678,7 +4714,12 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
   public void testUnnamedSavepoint() throws SQLException {
     try (Connection connection = DriverManager.getConnection(createUrl())) {
       connection.setAutoCommit(false);
-      assertNotNull(connection.setSavepoint());
+      Savepoint savepoint = connection.setSavepoint();
+      assertNotNull(savepoint);
+      assertEquals(0, savepoint.getSavepointId());
+
+      Savepoint savepoint2 = connection.setSavepoint();
+      assertEquals(1, savepoint2.getSavepointId());
     }
   }
 
@@ -4686,7 +4727,7 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
   public void testNamedSavepoint() throws SQLException {
     try (Connection connection = DriverManager.getConnection(createUrl())) {
       connection.setAutoCommit(false);
-      assertEquals("my-savepoint", connection.setSavepoint("my-savepoint").getSavepointName());
+      assertEquals("my_savepoint", connection.setSavepoint("my_savepoint").getSavepointName());
     }
   }
 
@@ -4694,7 +4735,7 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
   public void testReleaseSavepoint() throws SQLException {
     try (Connection connection = DriverManager.getConnection(createUrl())) {
       connection.setAutoCommit(false);
-      PSQLSavepoint savepoint = new PSQLSavepoint("my-savepoint");
+      Savepoint savepoint = connection.setSavepoint("my_savepoint");
       connection.releaseSavepoint(savepoint);
     }
   }
@@ -4703,10 +4744,8 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
   public void testRollbackToSavepoint() throws SQLException {
     try (Connection connection = DriverManager.getConnection(createUrl())) {
       connection.setAutoCommit(false);
-      PSQLSavepoint savepoint = new PSQLSavepoint("my-savepoint");
-      PSQLException exception =
-          assertThrows(PSQLException.class, () -> connection.rollback(savepoint));
-      assertEquals(PSQLState.NOT_IMPLEMENTED.getState(), exception.getSQLState());
+      Savepoint savepoint = connection.setSavepoint("my_savepoint");
+      connection.rollback(savepoint);
     }
   }
 
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/SavepointMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/SavepointMockServerTest.java
new file mode 100644
index 00000000..ce84aa79
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/SavepointMockServerTest.java
@@ -0,0 +1,598 @@
+// 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.pgadapter;
+
+import static com.google.cloud.spanner.connection.AbstractMockServerTest.RANDOM_RESULT_SET_ROW_COUNT;
+import static com.google.cloud.spanner.connection.AbstractMockServerTest.SELECT_RANDOM_STATEMENT;
+import static org.junit.Assert.assertEquals;
+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.Statement;
+import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.cloud.spanner.pgadapter.error.SQLState;
+import com.google.common.base.Strings;
+import com.google.protobuf.AbstractMessage;
+import com.google.spanner.v1.BeginTransactionRequest;
+import com.google.spanner.v1.CommitRequest;
+import com.google.spanner.v1.ExecuteBatchDmlRequest;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.RollbackRequest;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Savepoint;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.postgresql.jdbc.PSQLSavepoint;
+
+@RunWith(JUnit4.class)
+public class SavepointMockServerTest extends AbstractMockServerTest {
+
+  @BeforeClass
+  public static void loadPgJdbcDriver() throws Exception {
+    // Make sure the PG JDBC driver is loaded.
+    Class.forName("org.postgresql.Driver");
+  }
+
+  @Before
+  public void setupRandomResults() {
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            SELECT_RANDOM_STATEMENT,
+            new RandomResultSetGenerator(RANDOM_RESULT_SET_ROW_COUNT, Dialect.POSTGRESQL)
+                .generate()));
+  }
+
+  @After
+  public void clearRequests() {
+    mockSpanner.clearRequests();
+  }
+
+  private String createUrl() {
+    return String.format("jdbc:postgresql://localhost:%d/", pgServer.getLocalPort());
+  }
+
+  private Connection createConnection() throws SQLException {
+    return DriverManager.getConnection(createUrl());
+  }
+
+  @Test
+  public void testCreateSavepoint() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      connection.setSavepoint("s1");
+
+      // PostgreSQL allows multiple savepoints with the same name.
+      connection.setSavepoint("s1");
+
+      // The PostgreSQL JDBC driver also allows named savepoints without a name...
+      PSQLSavepoint savepoint = (PSQLSavepoint) connection.setSavepoint(null);
+      assertTrue(savepoint.getPGName(), savepoint.getPGName().startsWith("JDBC_SAVEPOINT_"));
+
+      // Test invalid identifiers.
+      assertThrows(SQLException.class, () -> connection.setSavepoint(""));
+      assertThrows(SQLException.class, () -> connection.setSavepoint("1"));
+      assertThrows(SQLException.class, () -> connection.setSavepoint("-foo"));
+      assertThrows(SQLException.class, () -> connection.setSavepoint(Strings.repeat("t", 129)));
+    }
+  }
+
+  @Test
+  public void testReleaseSavepoint() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      Savepoint s1 = connection.setSavepoint("s1");
+      connection.releaseSavepoint(s1);
+      assertThrows(SQLException.class, () -> connection.releaseSavepoint(s1));
+
+      Savepoint s1_2 = connection.setSavepoint("s1");
+      Savepoint s2 = connection.setSavepoint("s2");
+      connection.releaseSavepoint(s1_2);
+      // Releasing a savepoint also removes all savepoints after it.
+      SQLException exception =
+          assertThrows(SQLException.class, () -> connection.releaseSavepoint(s2));
+      assertEquals(SQLState.SavepointException.toString(), exception.getSQLState());
+      // Rollback the entire transaction to make the connection usable again.
+      connection.rollback();
+
+      // PostgreSQL allows multiple savepoints with the same name.
+      Savepoint savepoint1 = connection.setSavepoint("s1");
+      Savepoint savepoint2 = connection.setSavepoint("s2");
+      Savepoint savepoint1_2 = connection.setSavepoint("s1");
+      connection.releaseSavepoint(savepoint1_2);
+      connection.releaseSavepoint(savepoint2);
+      connection.releaseSavepoint(savepoint1);
+      assertThrows(SQLException.class, () -> connection.releaseSavepoint(savepoint1));
+    }
+  }
+
+  @Test
+  public void testRollbackToSavepoint() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      Savepoint s1 = connection.setSavepoint("s1");
+      connection.rollback(s1);
+      // Rolling back to a savepoint does not remove it, so we can roll back multiple times to the
+      // same savepoint.
+      connection.rollback(s1);
+
+      Savepoint s2 = connection.setSavepoint("s2");
+      connection.rollback(s1);
+      // Rolling back to a savepoint removes all savepoints after it.
+      assertThrows(SQLException.class, () -> connection.rollback(s2));
+      // Rollback again to make the connection usable again.
+      connection.rollback();
+
+      // PostgreSQL allows multiple savepoints with the same name.
+      connection.setSavepoint("s1");
+      Savepoint savepoint2 = connection.setSavepoint("s2");
+      Savepoint savepoint1 = connection.setSavepoint("s1");
+      connection.rollback(savepoint1);
+      connection.rollback(savepoint2);
+      connection.rollback(savepoint1);
+      connection.rollback(savepoint1);
+      connection.releaseSavepoint(savepoint1);
+      assertThrows(SQLException.class, () -> connection.rollback(savepoint1));
+    }
+  }
+
+  @Test
+  public void testSavepointInAutoCommit() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(true);
+      assertThrows(SQLException.class, () -> connection.setSavepoint("s1"));
+
+      // Starting a 'manual' transaction in autocommit mode does not enable savepoints in the PG
+      // JDBC driver.
+      connection.createStatement().execute("begin transaction");
+      SQLException exception =
+          assertThrows(SQLException.class, () -> connection.setSavepoint("s1"));
+      assertEquals(SQLState.NoActiveSqlTransaction.toString(), exception.getSQLState());
+    }
+  }
+
+  @Test
+  public void testRollbackToSavepointInReadOnlyTransaction() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setReadOnly(true);
+      connection.setAutoCommit(false);
+
+      // Read-only transactions also support savepoints, but they do not do anything. This feature
+      // is here purely for compatibility.
+      Savepoint s1 = connection.setSavepoint("s1");
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(RANDOM_RESULT_SET_ROW_COUNT, count);
+      }
+
+      connection.rollback(s1);
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(RANDOM_RESULT_SET_ROW_COUNT, count);
+      }
+      // Committing a read-only transaction is necessary to mark the end of the transaction.
+      // It is a no-op on Cloud Spanner.
+      connection.commit();
+
+      assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
+      BeginTransactionRequest beginRequest =
+          mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0);
+      assertTrue(beginRequest.getOptions().hasReadOnly());
+      assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+    }
+  }
+
+  @Test
+  public void testRollbackToSavepointInReadWriteTransaction() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      Savepoint s1 = connection.setSavepoint("s1");
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(RANDOM_RESULT_SET_ROW_COUNT, count);
+      }
+
+      connection.rollback(s1);
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(RANDOM_RESULT_SET_ROW_COUNT, count);
+      }
+      connection.commit();
+
+      // Read/write transactions are started with inlined Begin transaction options.
+      assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
+      assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
+      assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+      assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+
+      List<AbstractMessage> requests =
+          mockSpanner.getRequests().stream()
+              .filter(
+                  request ->
+                      request instanceof ExecuteSqlRequest
+                          || request instanceof RollbackRequest
+                          || request instanceof CommitRequest)
+              .collect(Collectors.toList());
+      assertEquals(4, requests.size());
+      int index = 0;
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(RollbackRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(CommitRequest.class, requests.get(index++).getClass());
+    }
+  }
+
+  @Test
+  public void testRollbackToSavepointWithDmlStatements() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      // First do a query that is included in the transaction.
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(RANDOM_RESULT_SET_ROW_COUNT, count);
+      }
+      // Set a savepoint and execute a couple of DML statements.
+      Savepoint s1 = connection.setSavepoint("s1");
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+      Savepoint s2 = connection.setSavepoint("s2");
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+      // Rollback the last DML statement and commit.
+      connection.rollback(s2);
+
+      connection.commit();
+
+      // Read/write transactions are started with inlined Begin transaction options.
+      assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
+      assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
+      assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+      assertEquals(5, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+
+      List<AbstractMessage> requests =
+          mockSpanner.getRequests().stream()
+              .filter(
+                  request ->
+                      request instanceof ExecuteSqlRequest
+                          || request instanceof RollbackRequest
+                          || request instanceof CommitRequest)
+              .collect(Collectors.toList());
+      assertEquals(7, requests.size());
+      int index = 0;
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(RollbackRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(CommitRequest.class, requests.get(index++).getClass());
+    }
+  }
+
+  @Test
+  public void testRollbackToSavepointFails() throws SQLException {
+    Statement statement = Statement.of("select * from foo where bar=true");
+    int numRows = 10;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows, Dialect.POSTGRESQL);
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      try (ResultSet resultSet = connection.createStatement().executeQuery(statement.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(numRows, count);
+      }
+      // Set a savepoint and execute a couple of DML statements.
+      Savepoint s1 = connection.setSavepoint("s1");
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+      // Change the result of the initial query.
+      mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+      // Rollback to before the DML statements.
+      // This will succeed as long as we don't execute any further statements.
+      connection.rollback(s1);
+
+      // Trying to commit the transaction or execute any other statements on the transaction will
+      // fail.
+      assertThrows(SQLException.class, connection::commit);
+
+      // Read/write transactions are started with inlined Begin transaction options.
+      assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
+      assertEquals(2, mockSpanner.countRequestsOfType(RollbackRequest.class));
+      assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+      assertEquals(4, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+
+      List<AbstractMessage> requests =
+          mockSpanner.getRequests().stream()
+              .filter(
+                  request ->
+                      request instanceof ExecuteSqlRequest
+                          || request instanceof RollbackRequest
+                          || request instanceof CommitRequest)
+              .collect(Collectors.toList());
+      assertEquals(6, requests.size());
+      int index = 0;
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(RollbackRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(RollbackRequest.class, requests.get(index++).getClass());
+    }
+  }
+
+  @Test
+  public void testRollbackToSavepointSucceedsWithRollback() throws SQLException {
+    Statement statement = Statement.of("select * from foo where bar=true");
+    int numRows = 10;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows, Dialect.POSTGRESQL);
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      try (ResultSet resultSet = connection.createStatement().executeQuery(statement.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(numRows, count);
+      }
+      // Change the result of the initial query and set a savepoint.
+      Savepoint s1 = connection.setSavepoint("s1");
+      mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+      // This will succeed as long as we don't execute any further statements.
+      connection.rollback(s1);
+
+      // Rolling back the transaction should now be a no-op, as it has already been rolled back.
+      connection.rollback();
+
+      // Read/write transactions are started with inlined Begin transaction options.
+      assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
+      assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
+      assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+      assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    }
+  }
+
+  @Test
+  public void testMultipleRollbacksWithChangedResults() throws SQLException {
+    Statement statement = Statement.of("select * from foo where bar=true");
+    int numRows = 10;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows, Dialect.POSTGRESQL);
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      try (ResultSet resultSet = connection.createStatement().executeQuery(statement.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(numRows, count);
+      }
+      Savepoint s1 = connection.setSavepoint("s1");
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+      Savepoint s2 = connection.setSavepoint("s2");
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+
+      // Change the result of the initial query to make sure that any retry will fail.
+      mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+      // This will succeed as long as we don't execute any further statements.
+      connection.rollback(s2);
+      // Rolling back one further should also work.
+      connection.rollback(s1);
+
+      // Rolling back the transaction should now be a no-op, as it has already been rolled back.
+      connection.rollback();
+
+      assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
+      assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+      assertEquals(3, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    }
+  }
+
+  @Test
+  public void testMultipleRollbacks() throws SQLException {
+    Statement statement = Statement.of("select * from foo where bar=true");
+    int numRows = 10;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows, Dialect.POSTGRESQL);
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      try (ResultSet resultSet = connection.createStatement().executeQuery(statement.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(numRows, count);
+      }
+      Savepoint s1 = connection.setSavepoint("s1");
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+      Savepoint s2 = connection.setSavepoint("s2");
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+
+      // First roll back one step and then one more.
+      connection.rollback(s2);
+      connection.rollback(s1);
+
+      // This will only commit the SELECT query.
+      connection.commit();
+
+      assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
+      assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+      assertEquals(4, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+
+      List<AbstractMessage> requests =
+          mockSpanner.getRequests().stream()
+              .filter(
+                  request ->
+                      request instanceof ExecuteSqlRequest
+                          || request instanceof RollbackRequest
+                          || request instanceof CommitRequest)
+              .collect(Collectors.toList());
+      assertEquals(6, requests.size());
+      int index = 0;
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(RollbackRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(CommitRequest.class, requests.get(index++).getClass());
+    }
+  }
+
+  @Test
+  public void testRollbackBatchDml() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+      Savepoint s1 = connection.setSavepoint("s1");
+      try (java.sql.Statement statement = connection.createStatement()) {
+        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.executeBatch();
+      }
+      Savepoint s2 = connection.setSavepoint("s2");
+
+      connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql());
+      Savepoint s3 = connection.setSavepoint("s3");
+      try (java.sql.Statement statement = connection.createStatement()) {
+        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.executeBatch();
+      }
+      connection.setSavepoint("s4");
+
+      connection.rollback(s2);
+
+      connection.commit();
+
+      assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
+      assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+      assertEquals(3, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+      assertEquals(3, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+
+      List<AbstractMessage> requests =
+          mockSpanner.getRequests().stream()
+              .filter(
+                  request ->
+                      request instanceof ExecuteSqlRequest
+                          || request instanceof RollbackRequest
+                          || request instanceof CommitRequest
+                          || request instanceof ExecuteBatchDmlRequest)
+              .collect(Collectors.toList());
+      assertEquals(8, requests.size());
+      int index = 0;
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteBatchDmlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteBatchDmlRequest.class, requests.get(index++).getClass());
+      assertEquals(RollbackRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteSqlRequest.class, requests.get(index++).getClass());
+      assertEquals(ExecuteBatchDmlRequest.class, requests.get(index++).getClass());
+      assertEquals(CommitRequest.class, requests.get(index++).getClass());
+    }
+  }
+
+  @Test
+  public void testRollbackToSavepointWithoutInternalRetries() throws SQLException {
+    Statement statement = Statement.of("select * from foo where bar=true");
+    int numRows = 10;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows, Dialect.POSTGRESQL);
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      connection.createStatement().execute("set spanner.retry_aborts_internally=false");
+
+      Savepoint s1 = connection.setSavepoint("s1");
+      try (ResultSet resultSet = connection.createStatement().executeQuery(statement.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(numRows, count);
+      }
+      // This should work.
+      connection.rollback(s1);
+      // Resuming after a rollback is not supported without internal retries enabled.
+      assertThrows(
+          SQLException.class,
+          () -> connection.createStatement().executeUpdate(INSERT_STATEMENT.getSql()));
+    }
+  }
+
+  @Test
+  public void testRollbackToSavepointWithoutInternalRetriesInReadOnlyTransaction()
+      throws SQLException {
+    Statement statement = Statement.of("select * from foo where bar=true");
+    int numRows = 10;
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(numRows, Dialect.POSTGRESQL);
+    mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      connection.createStatement().execute("set spanner.retry_aborts_internally=false");
+      connection.commit();
+      connection.setReadOnly(true);
+
+      Savepoint s1 = connection.setSavepoint("s1");
+      try (ResultSet resultSet = connection.createStatement().executeQuery(statement.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(numRows, count);
+      }
+
+      // Both rolling back and resuming after a rollback are supported in a read-only transaction,
+      // even if internal retries have been disabled.
+      connection.rollback(s1);
+      try (ResultSet resultSet = connection.createStatement().executeQuery(statement.getSql())) {
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(numRows, count);
+      }
+    }
+  }
+}
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/python/psycopg3/Psycopg3MockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/python/psycopg3/Psycopg3MockServerTest.java
index 82edd05d..e93e50f8 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/python/psycopg3/Psycopg3MockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/python/psycopg3/Psycopg3MockServerTest.java
@@ -1368,7 +1368,6 @@ public class Psycopg3MockServerTest extends AbstractMockServerTest {
   }
 
   @Test
-  @Ignore("Nested transactions are not yet supported")
   public void testNestedTransaction() throws Exception {
     String sql = "SELECT * FROM all_types WHERE col_bigint=$1";
     mockSpanner.putStatementResult(
@@ -1389,7 +1388,34 @@ public class Psycopg3MockServerTest extends AbstractMockServerTest {
             + "col_timestamptz: 2022-02-16 13:18:02.123456+00:00\n"
             + "col_date: 2022-03-29\n"
             + "col_string: test\n"
-            + "col_jsonb: {'key': 'value'}\n",
+            + "col_jsonb: {'key': 'value'}\n"
+            + "col_array_bigint: [1, None, 2]\n"
+            + "col_array_bool: [True, None, False]\n"
+            + "col_array_bytea: [b'bytes1', None, b'bytes2']\n"
+            + "col_array_float8: [3.14, None, -99.99]\n"
+            + "col_array_int: [-100, None, -200]\n"
+            + "col_array_numeric: [Decimal('6.626'), None, Decimal('-3.14')]\n"
+            + "col_array_timestamptz: [datetime.datetime(2022, 2, 16, 16, 18, 2, 123456, tzinfo=<UTC>), None, datetime.datetime(2000, 1, 1, 0, 0, tzinfo=<UTC>)]\n"
+            + "col_array_date: [datetime.date(2023, 2, 20), None, datetime.date(2000, 1, 1)]\n"
+            + "col_array_string: ['string1', None, 'string2']\n"
+            + "col_array_jsonb: [{'key': 'value1'}, None, {'key': 'value2'}]\n",
+        result);
+  }
+
+  @Test
+  public void testRollbackNestedTransaction() throws Exception {
+    String sql = "SELECT * FROM all_types WHERE col_bigint=$1";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql), ALL_TYPES_RESULTSET.toBuilder().clearRows().build()));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.newBuilder(sql).bind("p1").to(1L).build(), ALL_TYPES_RESULTSET));
+
+    String result = execute("rollback_nested_transaction");
+    assertEquals(
+        "Nested transaction failed with error: Test rollback of savepoint\n"
+            + "Rolling back to a savepoint succeeded\n",
         result);
   }
 
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/SimpleParserTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/SimpleParserTest.java
index e602a815..8a76f47a 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/SimpleParserTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/SimpleParserTest.java
@@ -15,14 +15,18 @@
 package com.google.cloud.spanner.pgadapter.statements;
 
 import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.QuotedString.unescapeQuotedStringValue;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.addLimitIfParameterizedOffset;
 import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.parseCommand;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.unquoteOrFoldIdentifier;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
+import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.statements.SimpleParser.QuotedString;
 import com.google.cloud.spanner.pgadapter.statements.SimpleParser.TableOrIndexName;
@@ -73,6 +77,7 @@ public class SimpleParserTest {
   public void testTableOrIndexName() {
     assertEquals(new TableOrIndexName(null, "foo"), new TableOrIndexName(null, "foo"));
     assertNotEquals(new TableOrIndexName(null, "foo"), new TableOrIndexName(null, "bar"));
+    assertFalse(new TableOrIndexName(null, "bar").equals(new Object()));
   }
 
   @Test
@@ -130,6 +135,16 @@ public class SimpleParserTest {
     assertFalse(new SimpleParser("").eatToken("("));
   }
 
+  @Test
+  public void testUnquoteOrFoldIdentifier() {
+    assertEquals("test", unquoteOrFoldIdentifier("test"));
+    assertEquals("test", unquoteOrFoldIdentifier("TEST"));
+    assertEquals("test", unquoteOrFoldIdentifier("Test"));
+    assertEquals("test", unquoteOrFoldIdentifier("\"test\""));
+    assertEquals("TEST", unquoteOrFoldIdentifier("\"TEST\""));
+    assertEquals("Test", unquoteOrFoldIdentifier("\"Test\""));
+  }
+
   @Test
   public void testReadTableOrIndexNamePart() {
     assertEquals("\"foo\"", new SimpleParser("\"foo\"(id)").readIdentifierPart());
@@ -426,6 +441,14 @@ public class SimpleParserTest {
 
   @Test
   public void testQuotedString() {
+    QuotedString quotedString = new QuotedString(false, '\'', "'test'");
+    assertEquals("test", quotedString.getValue());
+    assertEquals("test", quotedString.getValue());
+    assertThrows(PGException.class, () -> new QuotedString(false, '\'', "'").getValue());
+    assertThrows(PGException.class, () -> new QuotedString(false, '\'', "'foo").getValue());
+    assertThrows(PGException.class, () -> new QuotedString(true, '\'', "'").getValue());
+    assertThrows(PGException.class, () -> new QuotedString(true, '\'', "'foo").getValue());
+
     assertEquals("test", new SimpleParser("'test'").readQuotedString('\'').getValue());
     assertEquals("test", new SimpleParser("e'test'").readQuotedString('\'').getValue());
     PGException exception =
@@ -498,4 +521,195 @@ public class SimpleParserTest {
     assertEquals(
         ImmutableList.of("foo 1", "bar 2"), SimpleParser.readArrayLiteral("{foo 1, bar 2}", true));
   }
+
+  @Test
+  public void testReadQueryParameter() {
+    assertEquals("$1", new SimpleParser("$1").readQueryParameter());
+    assertEquals("$9999999", new SimpleParser("$9999999").readQueryParameter());
+    assertEquals("$1", new SimpleParser("  \n\t /*comment*/$1\n--comment\n").readQueryParameter());
+
+    assertNull(new SimpleParser("foo").readQueryParameter());
+    assertNull(new SimpleParser("   ").readQueryParameter());
+    assertNull(new SimpleParser("$foo").readQueryParameter());
+    assertNull(new SimpleParser("$").readQueryParameter());
+    assertNull(new SimpleParser("$ ").readQueryParameter());
+    assertNull(new SimpleParser("").readQueryParameter());
+  }
+
+  @Test
+  public void testAddLimitIfParameterizedOffset() {
+    Statement statement;
+
+    statement = Statement.of("select * from foo");
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement = Statement.newBuilder("select * from foo where id=$1").bind("p1").to(1L).build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement = Statement.of("select * from foo limit 1");
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement = Statement.of("select * from foo limit 1 offset 1");
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement = Statement.of("select * from foo offset 1");
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from (select id from bar limit 1) where id=$1")
+            .bind("p1")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from (select id from bar offset 1) where id=$1")
+            .bind("p1")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from (select id from bar limit 1 offset 1) where id=$1")
+            .bind("p1")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from (select id from bar limit $1 offset $2) where id=$3")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .bind("p3")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    // The replacement is only done in the outer main query and not in any sub queries.
+    statement =
+        Statement.newBuilder("select * from (select id from bar offset $1) where id=$2")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from foo where id=$1 limit $2 offset $3")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .bind("p3")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from foo where id=$1 offset $2 limit $3")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .bind("p3")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from foo where id=$1 limit $2")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from foo where id=$1 offset $2")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to((Long) null)
+            .build();
+    assertEquals(
+        Statement.newBuilder("select * from foo where id=$1 offset $2 limit 4611686018427387903")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to((Long) null)
+            .build(),
+        addLimitIfParameterizedOffset(statement));
+
+    statement =
+        Statement.newBuilder("select * from foo where id=$1 offset $2")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build();
+    assertEquals(
+        Statement.newBuilder("select * from foo where id=$1 offset $2 limit 4611686018427387903")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build(),
+        addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from foo where id=$1   offset   $2   ")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build();
+    assertEquals(
+        Statement.newBuilder(
+                "select * from foo where id=$1   offset   $2    limit 4611686018427387903")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build(),
+        addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder("select * from foo where id=$1 offset $2 /* some comment */")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build();
+    assertEquals(
+        Statement.newBuilder(
+                "select * from foo where id=$1 offset $2 /* some comment */ limit 4611686018427387903")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build(),
+        addLimitIfParameterizedOffset(statement));
+    statement =
+        Statement.newBuilder(
+                "select * from foo where id=$1 offset -- Single line comment\n$2 /* some comment */")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build();
+    assertEquals(
+        Statement.newBuilder(
+                "select * from foo where id=$1 offset -- Single line comment\n$2 /* some comment */ limit 4611686018427387903")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .build(),
+        addLimitIfParameterizedOffset(statement));
+    // This is an invalid query. That will be caught by the backend parser.
+    statement =
+        Statement.newBuilder("select * from foo where id=$1 offset $2 offset $3")
+            .bind("p1")
+            .to(1L)
+            .bind("p2")
+            .to(1L)
+            .bind("p3")
+            .to(1L)
+            .build();
+    assertSame(statement, addLimitIfParameterizedOffset(statement));
+
+    statement = Statement.of("select * from foo where id=$1 offset $2");
+    assertEquals(
+        Statement.of("select * from foo where id=$1 offset $2 limit 4611686018427387903"),
+        addLimitIfParameterizedOffset(statement));
+  }
 }
diff --git a/src/test/python/psycopg3/psycopg3_tests.py b/src/test/python/psycopg3/psycopg3_tests.py
index 9180ef54..e3ca90f3 100644
--- a/src/test/python/psycopg3/psycopg3_tests.py
+++ b/src/test/python/psycopg3/psycopg3_tests.py
@@ -19,7 +19,8 @@ from decimal import Decimal
 
 import pytz
 import psycopg
-from psycopg import Copy
+from psycopg import Copy, Rollback
+from psycopg.errors import InFailedSqlTransaction
 from psycopg.types.json import Jsonb
 
 
@@ -496,8 +497,6 @@ def named_cursor(conn_string: str):
       print_all_types(row)
 
 
-# This method is currently not being used, as nested transactions are not yet
-# supported.
 def nested_transaction(conn_string: str):
   with psycopg.connect(conn_string) as conn:
     with conn.transaction() as tx1:
@@ -507,6 +506,27 @@ def nested_transaction(conn_string: str):
         print_all_types(row)
 
 
+def rollback_nested_transaction(conn_string: str):
+  with psycopg.connect(conn_string) as conn:
+    with conn.transaction():
+      try:
+        with conn.transaction():
+          conn.execute(
+            "SELECT * FROM all_types WHERE col_bigint=%s", (1,)).fetchone()
+          raise ValueError("Test rollback of savepoint")
+        print("Nested transaction succeeded")
+      except ValueError as e:
+        # We should come here, as the inner transaction always raises an error.
+        print("Nested transaction failed with error:", e)
+      try:
+        conn.execute(
+          "SELECT * FROM all_types WHERE col_bigint=%s", (1,)).fetchone()
+        print("Rolling back to a savepoint succeeded")
+      except InFailedSqlTransaction as e:
+        print("Outer transaction failed with error:", e)
+        raise Rollback()
+
+
 def create_batch_insert_values(batch_size: int):
   values = []
   for i in range(batch_size):