Analyzer AutoFinalOnQueryPass fix

src/Analyzer/Passes/AutoFinalOnQueryPass.cpp

@@ -1,8 +1,11 @@
 #include "AutoFinalOnQueryPass.h"
 
-#include <Analyzer/TableNode.h>
-#include <Analyzer/TableExpressionModifiers.h>
 #include <Storages/IStorage.h>
+
+#include <Analyzer/Utils.h>
+#include <Analyzer/TableNode.h>
+#include <Analyzer/TableFunctionNode.h>
+#include <Analyzer/TableExpressionModifiers.h>
 #include <Analyzer/InDepthQueryTreeVisitor.h>
 
 namespace DB
@@ -10,52 +13,64 @@ namespace DB
 
 namespace
 {
-    class AutoFinalOnQueryPassVisitor : public InDepthQueryTreeVisitorWithContext<AutoFinalOnQueryPassVisitor>
+
+class AutoFinalOnQueryPassVisitor : public InDepthQueryTreeVisitorWithContext<AutoFinalOnQueryPassVisitor>
+{
+public:
+    using Base = InDepthQueryTreeVisitorWithContext<AutoFinalOnQueryPassVisitor>;
+    using Base::Base;
+
+    void visitImpl(QueryTreeNodePtr & node)
     {
-    public:
-        using Base = InDepthQueryTreeVisitorWithContext<AutoFinalOnQueryPassVisitor>;
-        using Base::Base;
+        const auto & context = getContext();
+        if (!context->getSettingsRef().final)
+            return;
 
-        void visitImpl(QueryTreeNodePtr & node)
+        const auto * query_node = node->as<QueryNode>();
+        if (!query_node)
+            return;
+
+        auto table_expressions = extractTableExpressions(query_node->getJoinTree());
+        for (auto & table_expression : table_expressions)
+            applyFinalIfNeeded(table_expression);
+    }
+private:
+    static void applyFinalIfNeeded(QueryTreeNodePtr & node)
+    {
+        auto * table_node = node->as<TableNode>();
+        auto * table_function_node = node->as<TableFunctionNode>();
+        if (!table_node && !table_function_node)
+            return;
+
+        const auto & storage = table_node ? table_node->getStorage() : table_function_node->getStorage();
+        bool is_final_supported = storage && storage->supportsFinal() && !storage->isRemote();
+        if (!is_final_supported)
+            return;
+
+        TableExpressionModifiers table_expression_modifiers_with_final(true /*has_final*/, {}, {});
+
+        if (table_node)
         {
-            if (auto * table_node = node->as<TableNode>())
-            {
-                if (autoFinalOnQuery(*table_node, table_node->getStorage(), getContext()))
-                {
-                    auto modifier = TableExpressionModifiers(true, std::nullopt, std::nullopt);
-                    table_node->setTableExpressionModifiers(modifier);
-                }
-            }
+            if (table_node->hasTableExpressionModifiers())
+                table_node->getTableExpressionModifiers()->setHasFinal(true);
+            else
+                table_node->setTableExpressionModifiers(table_expression_modifiers_with_final);
         }
-
-    private:
-        static bool autoFinalOnQuery(TableNode & table_node, StoragePtr storage, ContextPtr context)
+        else if (table_function_node)
         {
-            bool is_auto_final_setting_on = context->getSettingsRef().final;
-            bool is_final_supported = storage && storage->supportsFinal() && !storage->isRemote();
-            bool is_query_already_final = table_node.hasTableExpressionModifiers() ? table_node.getTableExpressionModifiers().has_value() : false;
-
-            return is_auto_final_setting_on && !is_query_already_final && is_final_supported;
+            if (table_function_node->hasTableExpressionModifiers())
+                table_function_node->getTableExpressionModifiers()->setHasFinal(true);
+            else
+                table_function_node->setTableExpressionModifiers(table_expression_modifiers_with_final);
         }
+    }
+};
 
-    };
-
-}
-
-String AutoFinalOnQueryPass::getName()
-{
-    return "AutoFinalOnQueryPass";
-}
-
-String AutoFinalOnQueryPass::getDescription()
-{
-    return "Automatically applies final modifier to queries if it is supported and if user level final setting is set.";
 }
 
 void AutoFinalOnQueryPass::run(QueryTreeNodePtr query_tree_node, ContextPtr context)
 {
     auto visitor = AutoFinalOnQueryPassVisitor(std::move(context));
-
     visitor.visit(query_tree_node);
 }
 

src/Analyzer/Passes/AutoFinalOnQueryPass.h

@@ -7,13 +7,23 @@
 namespace DB
 {
 
-
+/** Automatically applies final modifier to table expressions in queries if it is supported and if user level final setting is set.
+  *
+  * Example: SELECT id, value FROM test_table;
+  * Result: SELECT id, value FROM test_table FINAL;
+  */
 class AutoFinalOnQueryPass final : public IQueryTreePass
 {
 public:
-    String getName() override;
+    String getName() override
+    {
+        return "AutoFinalOnQueryPass";
+    }
 
-    String getDescription() override;
+    String getDescription() override
+    {
+        return "Automatically applies final modifier to table expressions in queries if it is supported and if user level final setting is set";
+    }
 
     void run(QueryTreeNodePtr query_tree_node, ContextPtr context) override;
 };

src/Analyzer/TableExpressionModifiers.h

@@ -28,6 +28,12 @@ public:
         return has_final;
     }
 
+    /// Set has final value
+    void setHasFinal(bool value)
+    {
+        has_final = value;
+    }
+
     /// Returns true if sample size ratio is specified, false otherwise
     bool hasSampleSizeRatio() const
     {

src/Analyzer/TableFunctionNode.h

@@ -116,6 +116,12 @@ public:
         return table_expression_modifiers;
     }
 
+    /// Get table expression modifiers
+    std::optional<TableExpressionModifiers> & getTableExpressionModifiers()
+    {
+        return table_expression_modifiers;
+    }
+
     /// Set table expression modifiers
     void setTableExpressionModifiers(TableExpressionModifiers table_expression_modifiers_value)
     {

src/Analyzer/TableNode.h

@@ -68,6 +68,12 @@ public:
         return table_expression_modifiers;
     }
 
+    /// Get table expression modifiers
+    std::optional<TableExpressionModifiers> & getTableExpressionModifiers()
+    {
+        return table_expression_modifiers;
+    }
+
     /// Set table expression modifiers
     void setTableExpressionModifiers(TableExpressionModifiers table_expression_modifiers_value)
     {