feat: support RESET ALL for the Spanner connection (#1904)

Adds support for the RESET ALL command to reset both the PGAdapter
state and the Connection state.

pom.xml

@@ -77,7 +77,7 @@
  <dependency>
  <groupId>com.google.cloud</groupId>
  <artifactId>libraries-bom</artifactId>
- <version>26.40.0</version>
+ <version>26.41.0</version>
  <type>pom</type>
  <scope>import</scope>
  </dependency>

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

@@ -17,7 +17,10 @@
 import static com.google.cloud.spanner.pgadapter.error.PGExceptionFactory.toPGException;
 import static com.google.cloud.spanner.pgadapter.statements.IntermediateStatement.PARSER;
 import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.addLimitIfParameterizedOffset;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.isCommand;
 import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.replaceForUpdate;
+import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.ROLLBACK;
+import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.SHOW;
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.ByteArray;
@@ -302,7 +305,8 @@ void doExecute() {
  // SELECT statements, then we should create a read-only transaction. Also, if a transaction
  // block always ends with a ROLLBACK, PGAdapter should skip the entire execution of that
  // block.
- SessionStatement sessionStatement = getSessionManagementStatement(parsedStatement);
+ SessionStatement sessionStatement =
+ getSessionManagementStatement(updatedStatement, parsedStatement);
  if (!localStatements.get().isEmpty()
  && localStatements.get().containsKey(statement.getSql())
  && localStatements.get().get(statement.getSql()) != null
@@ -312,7 +316,7 @@ void doExecute() {
  Objects.requireNonNull(localStatements.get().get(statement.getSql()));
  result.set(localStatement.execute(BackendConnection.this));
  } else if (sessionStatement != null) {
- result.set(sessionStatement.execute(sessionState));
+ result.set(sessionStatement.execute(sessionState, spannerConnection));
  } else if (connectionState == ConnectionState.ABORTED
  && !spannerConnection.isInTransaction()
  && (isRollback(parsedStatement) || isCommit(parsedStatement))) {
@@ -519,14 +523,14 @@ boolean isUnsupportedConcurrencyModeException(SpannerException spannerException)
  }
 
  @Nullable
- SessionStatement getSessionManagementStatement(ParsedStatement parsedStatement) {
+ SessionStatement getSessionManagementStatement(
+ Statement statement, ParsedStatement parsedStatement) {
  if (parsedStatement.getType() == StatementType.UNKNOWN
  || (parsedStatement.getType() == StatementType.QUERY
- && parsedStatement.getSqlWithoutComments().length() >= 4
- && parsedStatement
- .getSqlWithoutComments()
- .substring(0, 4)
- .equalsIgnoreCase("show"))) {
+ && isCommand(SHOW, statement.getSql()))
+ || (parsedStatement.getType() == StatementType.CLIENT_SIDE
+ && parsedStatement.getClientSideStatementType()
+ == ClientSideStatementType.RESET_ALL)) {
  return SessionStatementParser.parse(parsedStatement);
  }
  return null;

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

@@ -113,7 +113,9 @@ public StatementType getStatementType() {
  public void executeAsync(BackendConnection backendConnection) {
  this.executed = true;
  setFutureStatementResult(
- Futures.immediateFuture(showStatement.execute(backendConnection.getSessionState())));
+ Futures.immediateFuture(
+ showStatement.execute(
+ backendConnection.getSessionState(), backendConnection.getSpannerConnection())));
  }
 
  @Override

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

@@ -118,7 +118,8 @@ public StatementType getStatementType() {
  public void executeAsync(BackendConnection backendConnection) {
  this.executed = true;
  try {
- setStatement.execute(backendConnection.getSessionState());
+ setStatement.execute(
+ backendConnection.getSessionState(), backendConnection.getSpannerConnection());
  } catch (Throwable throwable) {
  setFutureStatementResult(Futures.immediateFailedFuture(throwable));
  return;

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

@@ -23,7 +23,9 @@
 import com.google.cloud.spanner.Type.StructField;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
+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.cloud.spanner.pgadapter.session.PGSetting;
 import com.google.cloud.spanner.pgadapter.session.SessionState;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
@@ -63,7 +65,7 @@ String getKey() {
  return "all";
  }
 
- abstract StatementResult execute(SessionState sessionState);
+ abstract StatementResult execute(SessionState sessionState, Connection connection);
  }
 
  static class SetStatement extends SessionStatement {
@@ -103,7 +105,7 @@ SetStatement build() {
  }
 
  @Override
- public StatementResult execute(SessionState sessionState) {
+ public StatementResult execute(SessionState sessionState, Connection connection) {
  if (local) {
  sessionState.setLocal(extension, name, value);
  } else {
@@ -145,14 +147,15 @@ static ResetStatement createResetAll() {
  }
 
  @Override
- public StatementResult execute(SessionState sessionState) {
+ public StatementResult execute(SessionState sessionState, Connection connection) {
  if (extension == null && name != null) {
  PGSetting setting = sessionState.get(null, name);
  sessionState.set(null, name, setting.getResetVal());
  } else if (extension != null && name != null) {
  sessionState.set(extension, name, null);
  } else {
  sessionState.resetAll();
+ connection.reset();
  }
  return RESET_RESULT;
  }
@@ -190,7 +193,7 @@ static ShowStatement createShowAll() {
  }
 
  @Override
- public StatementResult execute(SessionState sessionState) {
+ public StatementResult execute(SessionState sessionState, Connection connection) {
  if (name != null) {
  String value;
  if (missingOk) {
@@ -263,7 +266,8 @@ private static String unquote(String value) {
  }
 
  public static @Nullable SessionStatement parse(ParsedStatement parsedStatement) {
- if (parsedStatement.getType() == StatementType.CLIENT_SIDE) {
+ if (parsedStatement.getType() == StatementType.CLIENT_SIDE
+ && parsedStatement.getClientSideStatementType() != ClientSideStatementType.RESET_ALL) {
  // This statement is handled by the Connection API.
  return null;
  }

src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java

@@ -29,7 +29,9 @@ public class QueryMessage extends ControlMessage {
  public static final String PREPARE = "PREPARE";
  public static final String EXECUTE = "EXECUTE";
  public static final String DEALLOCATE = "DEALLOCATE";
+ public static final String SHOW = "SHOW";
  public static final String DISCARD = "DISCARD";
+ public static final String RESET = "RESET";
  public static final String VACUUM = "VACUUM";
  public static final String TRUNCATE = "TRUNCATE";
  public static final String SAVEPOINT = "SAVEPOINT";

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

@@ -4150,8 +4150,7 @@ public void testResetAll() throws SQLException {
 
  verifySettingIsNull(connection, "application_name");
  verifySettingValue(connection, "search_path", "public");
- // TODO: Change this when the Spanner Connection API also supports RESET ALL.
- verifySettingValue(connection, "spanner.autocommit_dml_mode", "PARTITIONED_NON_ATOMIC");
+ verifySettingValue(connection, "spanner.autocommit_dml_mode", "TRANSACTIONAL");
  }
  }