feat: support (ignore) show and set search_path (#288)

* feat: support show search_path * fix: only include public in search path * feat: add support for set search_path * feat: add support for set search_path * fix: remove special handling of UNKNOWN We can safely remove the special handling of UNKNOWN statements, as the case where there is no command tag is already handled at the beginning of the method. * docs: add todo to remove later
GoogleCloudPlatform · Jul 24, 2022 · 93d8c14 · 93d8c14
1 parent ae1f2f3
commit 93d8c14
Show file tree

Hide file tree

Showing 8 changed files with 231 additions and 12 deletions.
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
@@ -806,7 +806,7 @@ public static final class NoResult implements StatementResult {
       this.commandTag = null;
     }
 
-    private NoResult(String commandTag) {
+    public NoResult(String commandTag) {
       this.commandTag = commandTag;
     }
 

diff --git a/...main/java/com/google/cloud/spanner/pgadapter/statements/local/SetSearchPathStatement.java b/...main/java/com/google/cloud/spanner/pgadapter/statements/local/SetSearchPathStatement.java
@@ -0,0 +1,69 @@
+// 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.local;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
+
+/**
+ * A no-op SET search_path implementation. This should be removed once support has been added to the
+ * Connection API.
+ */
+// TODO: Remove this once search_path support has been added to the Connection API.
+@InternalApi
+public class SetSearchPathStatement implements LocalStatement {
+  public static final SetSearchPathStatement INSTANCE = new SetSearchPathStatement();
+
+  private SetSearchPathStatement() {}
+
+  @Override
+  public String[] getSql() {
+    return new String[] {
+      "set search_path to public",
+      "SET search_path TO public",
+      "SET SEARCH_PATH TO public",
+      "set search_path to \"public\"",
+      "SET search_path TO \"public\"",
+      "SET SEARCH_PATH TO \"public\"",
+      "set search_path to \"$user\", public",
+      "SET search_path TO \"$user\", public",
+      "SET SEARCH_PATH TO \"$user\", public",
+      "set search_path to \"$user\", \"public\"",
+      "SET search_path TO \"$user\", \"public\"",
+      "SET SEARCH_PATH TO \"$user\", \"public\"",
+      // Some tools prepend the current search_path with 'public' without properly checking whether
+      // 'public' is already part of the existing search path.
+      "set search_path to public, public",
+      "SET search_path TO public, public",
+      "SET SEARCH_PATH TO public, public",
+      "set search_path to public, \"public\"",
+      "SET search_path TO public, \"public\"",
+      "SET SEARCH_PATH TO public, \"public\"",
+      "set search_path to public, \"$user\", public",
+      "SET search_path TO public, \"$user\", public",
+      "SET SEARCH_PATH TO public, \"$user\", public",
+      "set search_path to public, \"$user\", \"public\"",
+      "SET search_path TO public, \"$user\", \"public\"",
+      "SET SEARCH_PATH TO public, \"$user\", \"public\"",
+    };
+  }
+
+  @Override
+  public StatementResult execute(BackendConnection backendConnection) {
+    return new NoResult("SET");
+  }
+}
diff --git a/...ain/java/com/google/cloud/spanner/pgadapter/statements/local/ShowSearchPathStatement.java b/...ain/java/com/google/cloud/spanner/pgadapter/statements/local/ShowSearchPathStatement.java
@@ -0,0 +1,53 @@
+// 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.local;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.ResultSets;
+import com.google.cloud.spanner.Struct;
+import com.google.cloud.spanner.Type;
+import com.google.cloud.spanner.Type.StructField;
+import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
+import com.google.common.collect.ImmutableList;
+
+@InternalApi
+public class ShowSearchPathStatement implements LocalStatement {
+  public static final ShowSearchPathStatement INSTANCE = new ShowSearchPathStatement();
+
+  private ShowSearchPathStatement() {}
+
+  @Override
+  public String[] getSql() {
+    return new String[] {
+      "show search_path", "SHOW search_path", "SHOW SEARCH_PATH",
+    };
+  }
+
+  @Override
+  public StatementResult execute(BackendConnection backendConnection) {
+    ResultSet resultSet =
+        ResultSets.forRows(
+            Type.struct(StructField.of("search_path", Type.string())),
+            ImmutableList.of(
+                Struct.newBuilder()
+                    .set("search_path")
+                    .to(backendConnection.getCurrentSchema())
+                    .build()));
+    return new QueryResult(resultSet);
+  }
+}
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java
@@ -21,6 +21,8 @@
 import com.google.cloud.spanner.pgadapter.statements.local.SelectCurrentCatalogStatement;
 import com.google.cloud.spanner.pgadapter.statements.local.SelectCurrentDatabaseStatement;
 import com.google.cloud.spanner.pgadapter.statements.local.SelectCurrentSchemaStatement;
+import com.google.cloud.spanner.pgadapter.statements.local.SetSearchPathStatement;
+import com.google.cloud.spanner.pgadapter.statements.local.ShowSearchPathStatement;
 import com.google.cloud.spanner.pgadapter.statements.local.ShowServerVersionStatement;
 import com.google.common.collect.ImmutableList;
 import java.util.List;
@@ -39,7 +41,9 @@ public class ClientAutoDetector {
           SelectCurrentSchemaStatement.INSTANCE,
           SelectCurrentDatabaseStatement.INSTANCE,
           SelectCurrentCatalogStatement.INSTANCE,
-          ShowServerVersionStatement.INSTANCE);
+          ShowSearchPathStatement.INSTANCE,
+          ShowServerVersionStatement.INSTANCE,
+          SetSearchPathStatement.INSTANCE);
 
   public enum WellKnownClient {
     PSQL {

diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
@@ -250,6 +250,7 @@ public void sendSpannerResult(IntermediateStatement statement, QueryMode mode, l
     switch (statement.getStatementType()) {
       case DDL:
       case CLIENT_SIDE:
+      case UNKNOWN:
         new CommandCompleteResponse(this.outputStream, command).send(false);
         break;
       case QUERY:
@@ -270,7 +271,6 @@ public void sendSpannerResult(IntermediateStatement statement, QueryMode mode, l
         command += ("INSERT".equals(command) ? " 0 " : " ") + statement.getUpdateCount();
         new CommandCompleteResponse(this.outputStream, command).send(false);
         break;
-      case UNKNOWN:
       default:
         throw new IllegalStateException("Unknown statement type: " + statement.getStatement());
     }

diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java
@@ -15,7 +15,6 @@
 package com.google.cloud.spanner.pgadapter;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -105,19 +104,34 @@ public void testInsertResult() throws Exception {
   }
 
   @Test
-  public void testUnknownStatementTypeThrowsError() {
+  public void testUnknownStatementTypeDoesNotThrowError() throws Exception {
+    ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+    DataOutputStream outputStream = new DataOutputStream(buffer);
+    DataInputStream inputStream =
+        new DataInputStream(
+            new ByteArrayInputStream(new byte[] {(byte) QUERY_IDENTIFIER, 0, 0, 0, 5, 0}));
+
+    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     ExecuteMessage executeMessage =
         new ExecuteMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN);
     IntermediateStatement intermediateStatement = mock(IntermediateStatement.class);
-    when(intermediateStatement.getCommandTag()).thenReturn("PARSE");
+    when(intermediateStatement.getCommandTag()).thenReturn("parse");
     when(intermediateStatement.getStatementType()).thenReturn(StatementType.UNKNOWN);
-    when(intermediateStatement.getStatement()).thenReturn("parse foo from bar");
 
-    IllegalStateException exception =
-        assertThrows(
-            IllegalStateException.class,
-            () -> executeMessage.sendSpannerResult(intermediateStatement, QueryMode.SIMPLE, 0L));
-    assertEquals("Unknown statement type: parse foo from bar", exception.getMessage());
+    executeMessage.sendSpannerResult(intermediateStatement, QueryMode.SIMPLE, 0L);
+
+    DataInputStream outputReader =
+        new DataInputStream(new ByteArrayInputStream(buffer.toByteArray()));
+    // identifier
+    outputReader.readByte();
+    // length
+    outputReader.readInt();
+    final String resultMessage = "parse";
+    int numOfBytes = resultMessage.getBytes(UTF8).length;
+    byte[] bytes = new byte[numOfBytes];
+    assertEquals(numOfBytes, outputReader.read(bytes, 0, numOfBytes));
+    assertEquals(resultMessage, new String(bytes, UTF8));
   }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -183,6 +183,38 @@ public void testSelectCurrentCatalog() throws SQLException {
     assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
   }
 
+  @Test
+  public void testShowSearchPath() throws SQLException {
+    String sql = "show search_path";
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals("public", resultSet.getString("search_path"));
+        assertFalse(resultSet.next());
+      }
+    }
+
+    // The statement is handled locally and not sent to Cloud Spanner.
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
+  @Test
+  public void testSetSearchPath() throws SQLException {
+    String sql = "set search_path to public";
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        assertFalse(statement.execute(sql));
+        assertEquals(0, statement.getUpdateCount());
+        assertFalse(statement.getMoreResults());
+      }
+    }
+
+    // The statement is handled locally and not sent to Cloud Spanner.
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
   @Test
   public void testShowServerVersion() throws SQLException {
     String sql = "show server_version";

diff --git a/...java/com/google/cloud/spanner/pgadapter/statements/local/ShowSearchPathStatementTest.java b/...java/com/google/cloud/spanner/pgadapter/statements/local/ShowSearchPathStatementTest.java
@@ -0,0 +1,47 @@
+// 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.local;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ShowSearchPathStatementTest {
+
+  @Test
+  public void testExecute() {
+    for (String schema : new String[] {"public", "my_schema"}) {
+      BackendConnection backendConnection = mock(BackendConnection.class);
+      when(backendConnection.getCurrentSchema()).thenReturn(schema);
+
+      try (ResultSet resultSet =
+          ShowSearchPathStatement.INSTANCE.execute(backendConnection).getResultSet()) {
+        assertTrue(resultSet.next());
+        assertEquals(1, resultSet.getColumnCount());
+        assertEquals(schema, resultSet.getString("search_path"));
+        assertFalse(resultSet.next());
+      }
+    }
+  }
+}