From 8bdc5695ee271d6ea8e5f3f8fc9bc9a68fb64ffa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 27 May 2024 09:46:13 +0200 Subject: [PATCH] feat: support DISCARD statements (#1863) * feat: support DISCARD statements Adds support for DISCARD ALL|PLANS|SEQUENCES|TEMP|TEMPORARY statements. All but ALL are no-ops. * chore: run code formatter --- .../spanner/pgadapter/ConnectionHandler.java | 3 +- .../spanner/pgadapter/error/SQLState.java | 3 + .../statements/DiscardStatement.java | 152 ++++++++++++++++++ .../pgadapter/wireprotocol/ParseMessage.java | 9 ++ .../pgadapter/wireprotocol/QueryMessage.java | 1 + .../JdbcSimpleModeMockServerTest.java | 38 +++++ .../statements/DiscardStatementTest.java | 44 +++++ 7 files changed, 249 insertions(+), 1 deletion(-) create mode 100644 src/main/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatement.java create mode 100644 src/test/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatementTest.java 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 932e21692..aab093435 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java +++ b/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); } 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 724cd1fbd..b2afd47ee 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 @@ -158,6 +158,9 @@ public enum SQLState { InFailedSqlTransaction("25P02"), IdleInTransactionSessionTimeout("25P03"), + // Class 26 - Invalid SQL statement name + InvalidSqlStatementName("26000"), + // Class 34 - Cursor InvalidCursorName("34000"), diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatement.java new file mode 100644 index 000000000..09b1d9ab0 --- /dev/null +++ b/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 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 parameterFormatCodes, + List 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); + } +} diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java index a10a767a7..b01fc53f3 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java +++ b/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, diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java index 454288a5a..61135469d 100644 --- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java +++ b/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"; diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java index fe173f6e2..f1620f7e6 100644 --- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java +++ b/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"; diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatementTest.java new file mode 100644 index 000000000..331c61c4e --- /dev/null +++ b/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")); + } +}