Fix Pre Process Metadata Call

temp
prestodb · Apr 15, 2024 · 8f80ac1 · 8f80ac1
1 parent 94e53c4
commit 8f80ac1
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 10 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/analyzer/MetadataExtractor.java b/presto-main/src/main/java/com/facebook/presto/sql/analyzer/MetadataExtractor.java
@@ -24,9 +24,11 @@
 import com.facebook.presto.sql.tree.Analyze;
 import com.facebook.presto.sql.tree.DefaultTraversalVisitor;
 import com.facebook.presto.sql.tree.Delete;
+import com.facebook.presto.sql.tree.Identifier;
 import com.facebook.presto.sql.tree.Insert;
 import com.facebook.presto.sql.tree.Statement;
 import com.facebook.presto.sql.tree.Table;
+import com.facebook.presto.sql.tree.With;
 
 import java.util.HashSet;
 import java.util.Optional;
@@ -125,17 +127,20 @@ private class MetadataExtractorContext
     {
         private final Optional<MetadataExtractorContext> parent;
         private final Set<QualifiedObjectName> tableNames;
+        private final Set<Identifier> tableNamesToSkipProcessing;  // Table names, such as query with CTE, that we don't need to get the metadata info so we should skip sending request
 
         public MetadataExtractorContext()
         {
             this.parent = Optional.empty();
             this.tableNames = new HashSet<>();
+            this.tableNamesToSkipProcessing = new HashSet<>();
         }
 
         public MetadataExtractorContext(Optional<MetadataExtractorContext> parent)
         {
             this.parent = parent;
             this.tableNames = new HashSet<>();
+            this.tableNamesToSkipProcessing = new HashSet<>();
         }
 
         public void addTable(QualifiedObjectName tableName)
@@ -147,21 +152,27 @@ public void addTable(QualifiedObjectName tableName)
 
         private boolean tableExists(QualifiedObjectName tableName)
         {
-            if (tableNames.contains(tableName)) {
-                return true;
-            }
-
-            if (parent.isPresent()) {
-                return parent.get().tableExists(tableName);
-            }
-
-            return false;
+            return tableNames.contains((tableName)) ||
+                    parent.map(metadataExtractorContext -> metadataExtractorContext.tableExists(tableName)).orElse(false);
         }
 
         public Set<QualifiedObjectName> getTableNames()
         {
             return tableNames;
         }
+
+        public void addTableNameToSkipProcessing(Identifier commonTableExpressionName)
+        {
+            if (!shouldSkipProcessing(commonTableExpressionName)) {
+                tableNamesToSkipProcessing.add(commonTableExpressionName);
+            }
+        }
+
+        public boolean shouldSkipProcessing(Identifier tableName)
+        {
+            return tableNamesToSkipProcessing.contains(tableName) ||
+                    parent.map(metadataExtractorContext -> metadataExtractorContext.shouldSkipProcessing(tableName)).orElse(false);
+        }
     }
 
     private class Visitor
@@ -185,8 +196,10 @@ protected Void visitTable(Table table, MetadataExtractorContext context)
                 throw new SemanticException(MISSING_SCHEMA, table, "Schema name is empty");
             }
 
+            if (!context.shouldSkipProcessing(new Identifier(tableName.getObjectName()))) {
+                context.addTable(tableName);
+            }
             // This could be either tableName, view, or MView
-            context.addTable(tableName);
             return super.visitTable(table, context);
         }
 
@@ -236,5 +249,15 @@ protected Void visitAnalyze(Analyze node, MetadataExtractorContext context)
             context.addTable(tableName);
             return super.visitAnalyze(node, context);
         }
+
+        @Override
+        protected Void visitWith(With node, MetadataExtractorContext context)
+        {
+            node.getQueries().forEach(query -> {
+                context.addTableNameToSkipProcessing(query.getName());
+                process(query, context);
+            });
+            return null;
+        }
     }
 }
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
@@ -6887,6 +6887,12 @@ public void testPreProcessMetastoreCalls()
                 "WHERE\n" +
                 "    c.nationkey = 1\n";
         assertEqualsIgnoreOrder(computeActual(enablePreProcessMetadataCalls, query), computeActual(getSession(), query));
+
+        query = "WITH temp_cte AS ( \n" +
+                "   SELECT name from nation \n" +
+                "   ) \n" +
+                "   SELECT * FROM temp_cte";
+        assertEqualsIgnoreOrder(computeActual(enablePreProcessMetadataCalls, query), computeActual(getSession(), query));
     }
 
     @Test