Analyzer: fix storage replacement with insertion block

src/Analyzer/Utils.cpp

@@ -326,6 +326,68 @@ void addTableExpressionOrJoinIntoTablesInSelectQuery(ASTPtr & tables_in_select_q
     }
 }
 
+QueryTreeNodes extractAllTableReferences(const QueryTreeNodePtr & tree)
+{
+    QueryTreeNodes result;
+
+    QueryTreeNodes nodes_to_process;
+    nodes_to_process.push_back(tree);
+
+    while (!nodes_to_process.empty())
+    {
+        auto node_to_process = std::move(nodes_to_process.back());
+        nodes_to_process.pop_back();
+
+        auto node_type = node_to_process->getNodeType();
+
+        switch (node_type)
+        {
+            case QueryTreeNodeType::TABLE:
+            {
+                result.push_back(std::move(node_to_process));
+                break;
+            }
+            case QueryTreeNodeType::QUERY:
+            {
+                nodes_to_process.push_back(node_to_process->as<QueryNode>()->getJoinTree());
+                break;
+            }
+            case QueryTreeNodeType::UNION:
+            {
+                for (const auto & union_node : node_to_process->as<UnionNode>()->getQueries().getNodes())
+                    nodes_to_process.push_back(union_node);
+                break;
+            }
+            case QueryTreeNodeType::TABLE_FUNCTION:
+            {
+                // Arguments of table function can't contain TableNodes.
+                break;
+            }
+            case QueryTreeNodeType::ARRAY_JOIN:
+            {
+                nodes_to_process.push_back(node_to_process->as<ArrayJoinNode>()->getTableExpression());
+                break;
+            }
+            case QueryTreeNodeType::JOIN:
+            {
+                auto & join_node = node_to_process->as<JoinNode &>();
+                nodes_to_process.push_back(join_node.getRightTableExpression());
+                nodes_to_process.push_back(join_node.getLeftTableExpression());
+                break;
+            }
+            default:
+            {
+                throw Exception(ErrorCodes::LOGICAL_ERROR,
+                                "Unexpected node type for table expression. "
+                                "Expected table, table function, query, union, join or array join. Actual {}",
+                                node_to_process->getNodeTypeName());
+            }
+        }
+    }
+
+    return result;
+}
+
 QueryTreeNodes extractTableExpressions(const QueryTreeNodePtr & join_tree_node, bool add_array_join)
 {
     QueryTreeNodes result;

src/Analyzer/Utils.h

@@ -50,10 +50,13 @@ std::optional<bool> tryExtractConstantFromConditionNode(const QueryTreeNodePtr &
   */
 void addTableExpressionOrJoinIntoTablesInSelectQuery(ASTPtr & tables_in_select_query_ast, const QueryTreeNodePtr & table_expression, const IQueryTreeNode::ConvertToASTOptions & convert_to_ast_options);
 
-/// Extract table, table function, query, union from join tree
+/// Extract all TableNodes from the query tree.
+QueryTreeNodes extractAllTableReferences(const QueryTreeNodePtr & tree);
+
+/// Extract table, table function, query, union from join tree.
 QueryTreeNodes extractTableExpressions(const QueryTreeNodePtr & join_tree_node, bool add_array_join = false);
 
-/// Extract left table expression from join tree
+/// Extract left table expression from join tree.
 QueryTreeNodePtr extractLeftTableExpression(const QueryTreeNodePtr & join_tree_node);
 
 /** Build table expressions stack that consists from table, table function, query, union, join, array join from join tree.

src/Interpreters/InterpreterSelectQueryAnalyzer.cpp

@@ -74,60 +74,25 @@ ContextMutablePtr buildContext(const ContextPtr & context, const SelectQueryOpti
 
 void replaceStorageInQueryTree(QueryTreeNodePtr & query_tree, const ContextPtr & context, const StoragePtr & storage)
 {
-    auto query_to_replace_table_expression = query_tree;
-    QueryTreeNodePtr table_expression_to_replace;
+    auto nodes = extractAllTableReferences(query_tree);
+    IQueryTreeNode::ReplacementMap replacement_map;
 
-    while (!table_expression_to_replace)
+    for (auto & node : nodes)
     {
-        if (auto * union_node = query_to_replace_table_expression->as<UnionNode>())
-            query_to_replace_table_expression = union_node->getQueries().getNodes().at(0);
-
-        auto & query_to_replace_table_expression_typed = query_to_replace_table_expression->as<QueryNode &>();
-        auto left_table_expression = extractLeftTableExpression(query_to_replace_table_expression_typed.getJoinTree());
-        auto left_table_expression_node_type = left_table_expression->getNodeType();
-
-        switch (left_table_expression_node_type)
-        {
-            case QueryTreeNodeType::QUERY:
-            case QueryTreeNodeType::UNION:
-            {
-                query_to_replace_table_expression = std::move(left_table_expression);
-                break;
-            }
-            case QueryTreeNodeType::TABLE:
-            case QueryTreeNodeType::TABLE_FUNCTION:
-            case QueryTreeNodeType::IDENTIFIER:
-            {
-                table_expression_to_replace = std::move(left_table_expression);
-                break;
-            }
-            default:
-            {
-                throw Exception(ErrorCodes::UNSUPPORTED_METHOD,
-                    "Expected table, table function or identifier node to replace with storage. Actual {}",
-                    left_table_expression->formatASTForErrorMessage());
-            }
-        }
-    }
-
-    /// Don't replace storage if table name differs
-    if (auto * table_node = table_expression_to_replace->as<TableNode>(); table_node && table_node->getStorageID().getFullNameNotQuoted() != storage->getStorageID().getFullNameNotQuoted())
-        return;
+        auto & table_node = node->as<TableNode &>();
 
-    auto replacement_table_expression = std::make_shared<TableNode>(storage, context);
-    std::optional<TableExpressionModifiers> table_expression_modifiers;
+        /// Don't replace storage if table name differs
+        if (table_node.getStorageID().getFullNameNotQuoted() != storage->getStorageID().getFullNameNotQuoted())
+            continue;
 
-    if (auto * table_node = table_expression_to_replace->as<TableNode>())
-        table_expression_modifiers = table_node->getTableExpressionModifiers();
-    else if (auto * table_function_node = table_expression_to_replace->as<TableFunctionNode>())
-        table_expression_modifiers = table_function_node->getTableExpressionModifiers();
-    else if (auto * identifier_node = table_expression_to_replace->as<IdentifierNode>())
-        table_expression_modifiers = identifier_node->getTableExpressionModifiers();
+        auto replacement_table_expression = std::make_shared<TableNode>(storage, context);
 
-    if (table_expression_modifiers)
-        replacement_table_expression->setTableExpressionModifiers(*table_expression_modifiers);
+        if (auto table_expression_modifiers = table_node.getTableExpressionModifiers())
+            replacement_table_expression->setTableExpressionModifiers(*table_expression_modifiers);
 
-    query_tree = query_tree->cloneAndReplace(table_expression_to_replace, std::move(replacement_table_expression));
+        replacement_map.emplace(node.get(), std::move(replacement_table_expression));
+    }
+    query_tree = query_tree->cloneAndReplace(replacement_map);
 }
 
 QueryTreeNodePtr buildQueryTreeAndRunPasses(const ASTPtr & query,