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,6 +70,7 @@
 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;
@@ -198,10 +201,13 @@ else if (statement instanceof ExecuteImmediate executeImmediate) {
             ImmutableSet.Builder<String> catalogBuilder = ImmutableSet.builder();
             ImmutableSet.Builder<String> schemaBuilder = ImmutableSet.builder();
             ImmutableSet.Builder<String> catalogSchemaBuilder = ImmutableSet.builder();
+            Set<QualifiedName> temporaryTables = new HashSet<>();
 
-            getNames(statement, tableBuilder, catalogBuilder, schemaBuilder, catalogSchemaBuilder);
+            getNames(statement, tableBuilder, catalogBuilder, schemaBuilder, catalogSchemaBuilder, temporaryTables);
             tables = tableBuilder.build();
+
             catalogBuilder.addAll(tables.stream().map(q -> q.getParts().getFirst()).iterator());
+
             catalogs = catalogBuilder.build();
             schemaBuilder.addAll(tables.stream().map(q -> q.getParts().get(1)).iterator());
             schemas = schemaBuilder.build();
@@ -263,7 +269,8 @@ private String decodePreparedStatementFromHeader(String headerValue)
     private void getNames(Node node, ImmutableSet.Builder<QualifiedName> tableBuilder,
             ImmutableSet.Builder<String> catalogBuilder,
             ImmutableSet.Builder<String> schemaBuilder,
-            ImmutableSet.Builder<String> catalogSchemaBuilder)
+            ImmutableSet.Builder<String> catalogSchemaBuilder,
+            Set<QualifiedName> temporaryTables)
             throws RequestParsingException
     {
         switch (node) {
@@ -278,6 +285,7 @@ private void getNames(Node node, ImmutableSet.Builder<QualifiedName> tableBuilde
             case DropCatalog s -> catalogBuilder.add(s.getCatalogName().getValue());
             case DropSchema s -> setCatalogAndSchemaNameFromSchemaQualifiedName(Optional.of(s.getSchemaName()), catalogBuilder, schemaBuilder, catalogSchemaBuilder);
             case DropTable s -> tableBuilder.add(qualifyName(s.getTableName()));
+            case Query q -> q.getWith().ifPresent(with -> temporaryTables.addAll(with.getQueries().stream().map(WithQuery::getName).map(Identifier::getValue).map(QualifiedName::of).toList()));
             case RenameMaterializedView s -> {
                 tableBuilder.add(qualifyName(s.getSource()));
                 tableBuilder.add(qualifyName(s.getTarget()));
@@ -336,13 +344,18 @@ private void getNames(Node node, ImmutableSet.Builder<QualifiedName> tableBuilde
             case SetSchemaAuthorization s -> setCatalogAndSchemaNameFromSchemaQualifiedName(Optional.of(s.getSource()), catalogBuilder, schemaBuilder, catalogSchemaBuilder);
             case SetTableAuthorization s -> tableBuilder.add(qualifyName(s.getSource()));
             case SetViewAuthorization s -> tableBuilder.add(qualifyName(s.getSource()));
-            case Table s -> tableBuilder.add(qualifyName(s.getName()));
+            case Table s -> {
+                // ignore temporary tables as they can have various table parts
+                if (!temporaryTables.contains(s.getName())) {
+                    tableBuilder.add(qualifyName(s.getName()));
+                }
+            }
             case TableFunctionInvocation s -> tableBuilder.add(qualifyName(s.getName()));
             default -> {}
         }
 
         for (Node child : node.getChildren()) {
-            getNames(child, tableBuilder, catalogBuilder, schemaBuilder, catalogSchemaBuilder);
+            getNames(child, tableBuilder, catalogBuilder, schemaBuilder, catalogSchemaBuilder, temporaryTables);
         }
     }
 

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

@@ -404,6 +404,33 @@ 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
+                """;
+        HttpServletRequest mockRequestWithDefaults = prepareMockRequest();
+        when(mockRequestWithDefaults.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
+        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);
+
+        Set<QualifiedName> tablesWithDefaults = trinoQueryPropertiesWithDefaults.getTables();
+        assertThat(tablesWithDefaults).containsExactly(QualifiedName.of("cat", "schem", "tbl1"));
+
+        HttpServletRequest mockRequestNoDefaults = prepareMockRequest();
+        when(mockRequestNoDefaults.getReader()).thenReturn(new BufferedReader(new StringReader(query)));
+
+        TrinoQueryProperties trinoQueryPropertiesNoDefaults = new TrinoQueryProperties(mockRequestNoDefaults, requestAnalyzerConfig);
+        Set<QualifiedName> tablesNoDefaults = trinoQueryPropertiesNoDefaults.getTables();
+        assertThat(tablesNoDefaults).containsExactly(QualifiedName.of("cat", "schem", "tbl1"));
+    }
+
     private HttpServletRequest prepareMockRequest()
     {
         HttpServletRequest mockRequest = mock(HttpServletRequest.class);