feat: support DISCARD statements

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

@@ -665,7 +665,8 @@ public boolean cancelActiveStatement(int connectionId, int secret) {
   public IntermediatePreparedStatement getStatement(String statementName) {
     if (!hasStatement(statementName)) {
       throw PGExceptionFactory.newPGException(
-          "prepared statement " + statementName + " does not exist");
+          "prepared statement " + statementName + " does not exist",
+          SQLState.InvalidSqlStatementName);
     }
     return this.statementsMap.get(statementName);
   }

src/main/java/com/google/cloud/spanner/pgadapter/error/SQLState.java

@@ -158,6 +158,9 @@ public enum SQLState {
   InFailedSqlTransaction("25P02"),
   IdleInTransactionSessionTimeout("25P03"),
 
+  // Class 26 - Invalid SQL statement name
+  InvalidSqlStatementName("26000"),
+
   // Class 34 - Cursor
   InvalidCursorName("34000"),
 

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

@@ -0,0 +1,152 @@
+// Copyright 2024 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.statements;
+
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
+import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
+import com.google.cloud.spanner.pgadapter.error.SQLState;
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection.ConnectionState;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
+import com.google.cloud.spanner.pgadapter.statements.DiscardStatement.ParsedDiscardStatement.DiscardType;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.Futures;
+import java.util.List;
+import java.util.concurrent.Future;
+
+public class DiscardStatement extends IntermediatePortalStatement {
+  static final class ParsedDiscardStatement {
+    enum DiscardType {
+      PLANS,
+      SEQUENCES,
+      TEMPORARY,
+      ALL,
+    }
+
+    final DiscardType type;
+
+    private ParsedDiscardStatement(DiscardType type) {
+      this.type = type;
+    }
+  }
+
+  private final ParsedDiscardStatement discardStatement;
+
+  public DiscardStatement(
+      ConnectionHandler connectionHandler,
+      OptionsMetadata options,
+      String name,
+      ParsedStatement parsedStatement,
+      Statement originalStatement) {
+    super(
+        name,
+        new IntermediatePreparedStatement(
+            connectionHandler,
+            options,
+            name,
+            NO_PARAMETER_TYPES,
+            parsedStatement,
+            originalStatement),
+        NO_PARAMS,
+        ImmutableList.of(),
+        ImmutableList.of());
+    this.discardStatement = parse(originalStatement.getSql());
+  }
+
+  @Override
+  public String getCommandTag() {
+    return "DISCARD";
+  }
+
+  @Override
+  public StatementType getStatementType() {
+    return StatementType.CLIENT_SIDE;
+  }
+
+  @Override
+  public void executeAsync(BackendConnection backendConnection) {
+    this.executed = true;
+    try {
+      switch (discardStatement.type) {
+        case ALL:
+          if (backendConnection.getConnectionState() == ConnectionState.TRANSACTION) {
+            setFutureStatementResult(
+                Futures.immediateFailedFuture(
+                    PGExceptionFactory.newPGException(
+                        "DISCARD ALL cannot run inside a transaction block",
+                        SQLState.ActiveSqlTransaction)));
+            return;
+          }
+          connectionHandler.closeAllStatements();
+          break;
+        case PLANS:
+        case SEQUENCES:
+        case TEMPORARY:
+        default:
+          // The default is no-op as PGAdapter does not support clearing plans, sequences or
+          // temporary tables.
+      }
+    } catch (Exception exception) {
+      setFutureStatementResult(Futures.immediateFailedFuture(exception));
+      return;
+    }
+    setFutureStatementResult(Futures.immediateFuture(new NoResult()));
+  }
+
+  @Override
+  public Future<StatementResult> describeAsync(BackendConnection backendConnection) {
+    // Return null to indicate that this DISCARD statement does not return any
+    // RowDescriptionResponse.
+    return Futures.immediateFuture(null);
+  }
+
+  @Override
+  public IntermediatePortalStatement createPortal(
+      String name,
+      byte[][] parameters,
+      List<Short> parameterFormatCodes,
+      List<Short> resultFormatCodes) {
+    // DISCARD does not support binding any parameters, so we just return the same statement.
+    return this;
+  }
+
+  static ParsedDiscardStatement parse(String sql) {
+    Preconditions.checkNotNull(sql);
+
+    SimpleParser parser = new SimpleParser(sql);
+    if (!parser.eatKeyword("discard")) {
+      throw PGExceptionFactory.newPGException("not a valid DISCARD statement: " + sql);
+    }
+    DiscardType type;
+    if (parser.eatKeyword("all")) {
+      type = DiscardType.ALL;
+    } else if (parser.eatKeyword("plans")) {
+      type = DiscardType.PLANS;
+    } else if (parser.eatKeyword("sequences")) {
+      type = DiscardType.SEQUENCES;
+    } else if (parser.eatKeyword("temp") || parser.eatKeyword("temporary")) {
+      type = DiscardType.TEMPORARY;
+    } else {
+      throw PGExceptionFactory.newPGException("Invalid DISCARD statement: " + parser.getSql());
+    }
+    parser.throwIfHasMoreTokens();
+    return new ParsedDiscardStatement(type);
+  }
+}

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

@@ -19,6 +19,7 @@
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.COPY;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.DEALLOCATE;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.DECLARE;
+import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.DISCARD;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.EXECUTE;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.FETCH;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.MOVE;
@@ -44,6 +45,7 @@
 import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
 import com.google.cloud.spanner.pgadapter.statements.DeallocateStatement;
 import com.google.cloud.spanner.pgadapter.statements.DeclareStatement;
+import com.google.cloud.spanner.pgadapter.statements.DiscardStatement;
 import com.google.cloud.spanner.pgadapter.statements.ExecuteStatement;
 import com.google.cloud.spanner.pgadapter.statements.FetchStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
@@ -148,6 +150,13 @@ static IntermediatePreparedStatement createStatement(
             name,
             parsedStatement,
             originalStatement);
+      } else if (isCommand(DISCARD, originalStatement.getSql())) {
+        return new DiscardStatement(
+            connectionHandler,
+            connectionHandler.getServer().getOptions(),
+            name,
+            parsedStatement,
+            originalStatement);
       } else if (isCommand(DECLARE, originalStatement.getSql())) {
         return new DeclareStatement(
             connectionHandler,

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

@@ -29,6 +29,7 @@ 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 DISCARD = "DISCARD";
   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/JdbcSimpleModeMockServerTest.java

@@ -16,12 +16,14 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 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.SpannerException;
+import com.google.cloud.spanner.pgadapter.error.SQLState;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ListValue;
@@ -54,6 +56,7 @@
 import org.junit.runners.Parameterized.Parameter;
 import org.junit.runners.Parameterized.Parameters;
 import org.postgresql.jdbc.TimestampUtils;
+import org.postgresql.util.PSQLException;
 
 /**
  * Tests the native PG JDBC driver in simple query mode. This is similar to the protocol that is
@@ -571,6 +574,41 @@ public void testExecuteUnknownStatement() throws SQLException {
     }
   }
 
+  @Test
+  public void testDiscard() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      connection.setAutoCommit(true);
+      // Verify that all variants are supported.
+      connection.createStatement().execute("discard all");
+      connection.createStatement().execute("discard plans");
+      connection.createStatement().execute("discard sequences");
+      connection.createStatement().execute("discard temp");
+      connection.createStatement().execute("discard temporary");
+
+      // Create a prepared statement verify that it is dropped by DISCARD ALL.
+      connection.createStatement().execute("prepare foo as SELECT 1");
+      connection.createStatement().execute("execute foo");
+      connection.createStatement().execute("discard all");
+      PSQLException exception =
+          assertThrows(
+              PSQLException.class, () -> connection.createStatement().execute("execute foo"));
+      assertNotNull(exception.getServerErrorMessage());
+      assertEquals(
+          SQLState.InvalidSqlStatementName.toString(),
+          exception.getServerErrorMessage().getSQLState());
+
+      // Verify that 'discard all' is not accepted in a transaction block.
+      connection.setAutoCommit(false);
+      exception =
+          assertThrows(
+              PSQLException.class, () -> connection.createStatement().execute("discard all"));
+      assertNotNull(exception.getServerErrorMessage());
+      assertEquals(
+          SQLState.ActiveSqlTransaction.toString(),
+          exception.getServerErrorMessage().getSQLState());
+    }
+  }
+
   @Test
   public void testGetTimezoneStringUtc() throws SQLException {
     String sql = "select '2022-01-01 10:00:00+01'::timestamptz";

src/test/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatementTest.java

@@ -0,0 +1,44 @@
+// Copyright 2022 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.statements;
+
+import static com.google.cloud.spanner.pgadapter.statements.DiscardStatement.parse;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import com.google.cloud.spanner.pgadapter.error.PGException;
+import com.google.cloud.spanner.pgadapter.statements.DiscardStatement.ParsedDiscardStatement.DiscardType;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class DiscardStatementTest {
+
+  @Test
+  public void testParse() {
+    assertEquals(DiscardType.ALL, parse("discard all").type);
+    assertEquals(DiscardType.PLANS, parse("discard plans").type);
+    assertEquals(DiscardType.SEQUENCES, parse("discard sequences").type);
+    assertEquals(DiscardType.TEMPORARY, parse("discard temp").type);
+    assertEquals(DiscardType.TEMPORARY, parse("discard temporary").type);
+    assertEquals(DiscardType.TEMPORARY, parse("discard/*comment*/temporary").type);
+
+    assertThrows(PGException.class, () -> parse("discard foo"));
+    assertThrows(PGException.class, () -> parse("discard"));
+    assertThrows(PGException.class, () -> parse("deallocate all foo"));
+    assertThrows(PGException.class, () -> parse("discard temp all"));
+  }
+}