Skip to content

Commit

Permalink
feat: support DISCARD statements (#1863)
Browse files Browse the repository at this point in the history
* feat: support DISCARD statements

Adds support for DISCARD ALL|PLANS|SEQUENCES|TEMP|TEMPORARY
statements. All but ALL are no-ops.

* chore: run code formatter
  • Loading branch information
olavloite committed May 27, 2024
1 parent 0ede0a5 commit 8bdc569
Show file tree
Hide file tree
Showing 7 changed files with 249 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,9 @@ public enum SQLState {
InFailedSqlTransaction("25P02"),
IdleInTransactionSessionTimeout("25P03"),

// Class 26 - Invalid SQL statement name
InvalidSqlStatementName("26000"),

// Class 34 - Cursor
InvalidCursorName("34000"),

Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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";
Expand Down
Original file line number Diff line number Diff line change
@@ -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"));
}
}

0 comments on commit 8bdc569

Please sign in to comment.