Fix parsing failure on WITH clauses

Author: Chaho12

gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java

@@ -44,6 +44,7 @@
 import io.trino.sql.tree.Node;
 import io.trino.sql.tree.NodeLocation;
 import io.trino.sql.tree.QualifiedName;
+import io.trino.sql.tree.Query;
 import io.trino.sql.tree.RenameMaterializedView;
 import io.trino.sql.tree.RenameSchema;
 import io.trino.sql.tree.RenameTable;
@@ -59,6 +60,7 @@
 import io.trino.sql.tree.Statement;
 import io.trino.sql.tree.Table;
 import io.trino.sql.tree.TableFunctionInvocation;
+import io.trino.sql.tree.WithQuery;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.ws.rs.HttpMethod;
 
@@ -68,12 +70,14 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Enumeration;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.io.BaseEncoding.base64Url;
 import static io.airlift.json.JsonCodec.jsonCodec;
 import static java.lang.Math.toIntExact;
@@ -90,6 +94,7 @@ public class TrinoQueryProperties
     private String queryType = "";
     private String resourceGroupQueryType = "";
     private Set<QualifiedName> tables = ImmutableSet.of();
+    private final Set<QualifiedName> temporaryTables = new HashSet<>();
     private final Optional<String> defaultCatalog;
     private final Optional<String> defaultSchema;
     private Set<String> catalogs = ImmutableSet.of();
@@ -201,12 +206,17 @@ else if (statement instanceof ExecuteImmediate executeImmediate) {
 
             getNames(statement, tableBuilder, catalogBuilder, schemaBuilder, catalogSchemaBuilder);
             tables = tableBuilder.build();
-            catalogBuilder.addAll(tables.stream().map(q -> q.getParts().getFirst()).iterator());
+
+            Set<QualifiedName> filteredTables = tables.stream()
+                    .filter(table -> !temporaryTables.contains(table))
+                    .collect(toImmutableSet());
+
+            catalogBuilder.addAll(filteredTables.stream().map(q -> q.getParts().getFirst()).iterator());
             catalogs = catalogBuilder.build();
-            schemaBuilder.addAll(tables.stream().map(q -> q.getParts().get(1)).iterator());
+            schemaBuilder.addAll(filteredTables.stream().map(q -> q.getParts().get(1)).iterator());
             schemas = schemaBuilder.build();
             catalogSchemaBuilder.addAll(
-                    tables.stream().map(qualifiedName -> format("%s.%s", qualifiedName.getParts().getFirst(), qualifiedName.getParts().get(1))).iterator());
+                    filteredTables.stream().map(qualifiedName -> format("%s.%s", qualifiedName.getParts().getFirst(), qualifiedName.getParts().get(1))).iterator());
             catalogSchemas = catalogSchemaBuilder.build();
         }
         catch (IOException e) {
@@ -338,6 +348,7 @@ private void getNames(Node node, ImmutableSet.Builder<QualifiedName> tableBuilde
             case SetViewAuthorization s -> tableBuilder.add(qualifyName(s.getSource()));
             case Table s -> tableBuilder.add(qualifyName(s.getName()));
             case TableFunctionInvocation s -> tableBuilder.add(qualifyName(s.getName()));
+            case WithQuery withQuery -> temporaryTables.add(QualifiedName.of(withQuery.getName().getValue()));
             default -> {}
         }
 
@@ -385,6 +396,12 @@ private QualifiedName qualifyName(QualifiedName table)
             throws RequestParsingException
     {
         List<String> tableParts = table.getParts();
+
+        // ignore temporary tables created by WITH clause as they can have various table parts
+        if (temporaryTables.contains(table)) {
+            return table;
+        }
+
         return switch (tableParts.size()) {
             case 1 -> QualifiedName.of(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), defaultSchema.orElseThrow(this::unsetDefaultExceptionSupplier), tableParts.getFirst());
             case 2 -> QualifiedName.of(defaultCatalog.orElseThrow(this::unsetDefaultExceptionSupplier), tableParts.getFirst(), tableParts.get(1));

gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelector.java

@@ -40,6 +40,7 @@
 import java.util.stream.Stream;
 
 import static io.trino.gateway.ha.router.RoutingGroupSelector.ROUTING_GROUP_HEADER;
+import static java.lang.String.format;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
@@ -404,6 +405,36 @@ void testTrinoQueryPropertiesTableExtraction(String query, Set<String> catalogs,
         assertThat(trinoQueryProperties.getCatalogs()).isEqualTo(catalogs);
     }
 
+    @Test
+    void testWithQueryNameExcluded()
+            throws IOException
+    {
+        String query = """
+                WITH dos AS (SELECT c1 from cat.schem.tbl1),
+                uno as (SELECT c1 FROM dos)
+                SELECT c1 FROM uno, dos
+                """;
+        BufferedReader bufferedReader = new BufferedReader(new StringReader(query));
+        HttpServletRequest mockRequestWithDefaults = prepareMockRequest();
+        when(mockRequestWithDefaults.getReader()).thenReturn(bufferedReader);
+        when(mockRequestWithDefaults.getHeader(TrinoQueryProperties.TRINO_CATALOG_HEADER_NAME)).thenReturn(DEFAULT_CATALOG);
+        when(mockRequestWithDefaults.getHeader(TrinoQueryProperties.TRINO_SCHEMA_HEADER_NAME)).thenReturn(DEFAULT_SCHEMA);
+
+        TrinoQueryProperties trinoQueryPropertiesWithDefaults = new TrinoQueryProperties(mockRequestWithDefaults, requestAnalyzerConfig);
+
+        assertThat(trinoQueryPropertiesWithDefaults.tablesContains(format("%s.%s.%s", DEFAULT_CATALOG, DEFAULT_SCHEMA, "uno"))).isFalse();
+        assertThat(trinoQueryPropertiesWithDefaults.tablesContains(format("%s.%s.%s", DEFAULT_CATALOG, DEFAULT_SCHEMA, "does"))).isFalse();
+        assertThat(trinoQueryPropertiesWithDefaults.tablesContains("cat.schem.tbl1")).isTrue();
+
+        HttpServletRequest mockRequestNoDefaults = prepareMockRequest();
+        bufferedReader = new BufferedReader(new StringReader(query));
+        when(mockRequestNoDefaults.getReader()).thenReturn(bufferedReader);
+
+        TrinoQueryProperties trinoQueryPropertiesNoDefaults = new TrinoQueryProperties(mockRequestNoDefaults, requestAnalyzerConfig);
+        assertThat(trinoQueryPropertiesNoDefaults.getTables()).hasSize(3);
+        assertThat(trinoQueryPropertiesNoDefaults.tablesContains("cat.schem.tbl1")).isTrue();
+    }
+
     private HttpServletRequest prepareMockRequest()
     {
         HttpServletRequest mockRequest = mock(HttpServletRequest.class);