feat: support savepoints (#796)

* feat: support savepoints Add support for emulated savepoints. Setting and releasing savepoints work in all cases. Rolling back to a savepoint is not guaranteed to work, as the underlying implementation will rollback the entire transaction and retry up to where the savepoint was set. Note that the retry will only be executed if the transaction is actually used after the rollback. * docs: update documentation * test: add tests and documentation * fix: update license header

Author: olavloite

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
 

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).

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).

src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java

@@ -34,6 +34,7 @@
 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 void connectToSpanner(String database, @Nullable Credentials credentials)
  spannerConnection.close();
  throw e;
  }
+ spannerConnection.setSavepointSupport(SavepointSupport.ENABLED);
  this.spannerConnection = spannerConnection;
  this.databaseId = connectionOptions.getDatabaseId();
  this.extendedQueryProtocolHandler = new ExtendedQueryProtocolHandler(this);

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"),

src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java

@@ -604,7 +604,16 @@ boolean isUpdate() {
 
  @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 +632,16 @@ boolean isUpdate() {
 
  @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 +660,15 @@ boolean isUpdate() {
 
  @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;
+ }
  }
  }
 

src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java

@@ -81,6 +81,7 @@
 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.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 {
@@ -4678,23 +4677,28 @@ public void testDescribeTruncate() throws SQLException {
  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());
  }
  }
 
  @Test
  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());
  }
  }
 
  @Test
  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 +4707,8 @@ public void testReleaseSavepoint() throws SQLException {
  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);
  }
  }