rewrote rewriting of SELECT query in StorageMerge with WITH statement

dbms/src/Common/VirtualColumnUtils.cpp

@@ -74,25 +74,17 @@ String chooseSuffixForSet(const NamesAndTypesList & columns, const std::vector<S
 void rewriteEntityInAst(ASTPtr ast, const String & column_name, const Field & value)
 {
     ASTSelectQuery & select = typeid_cast<ASTSelectQuery &>(*ast);
-    ASTExpressionList & node = typeid_cast<ASTExpressionList &>(*select.select_expression_list);
-    ASTs & asts = node.children;
-    auto cur = std::make_shared<ASTLiteral>(StringRange(), value);
-    cur->alias = column_name;
-    ASTPtr column_value = cur;
-    bool is_replaced = false;
-    for (size_t i = 0; i < asts.size(); ++i)
+    if (!select.with_expression_list)
     {
-        if (const ASTIdentifier * identifier = typeid_cast<const ASTIdentifier *>(&* asts[i]))
-        {
-            if (identifier->kind == ASTIdentifier::Kind::Column && identifier->name == column_name)
-            {
-                asts[i] = column_value;
-                is_replaced = true;
-            }
-        }
+        select.with_expression_list = std::make_shared<ASTExpressionList>();
+        select.children.insert(select.children.begin(), select.with_expression_list);
     }
-    if (!is_replaced)
-        asts.insert(asts.begin(), column_value);
+
+    ASTExpressionList & with = typeid_cast<ASTExpressionList &>(*select.with_expression_list);
+    auto literal = std::make_shared<ASTLiteral>(StringRange(), value);
+    literal->alias = column_name;
+    literal->prefer_alias_to_column_name = true;
+    with.children.push_back(literal);
 }
 
 /// Verifying that the function depends only on the specified columns

dbms/tests/integration/test_merge_table_over_distributed/test.py

@@ -17,6 +17,7 @@ def started_cluster():
         for node in (node1, node2):
             node.query('''
 CREATE TABLE local_table(id UInt32, val String) ENGINE = TinyLog;
+CREATE TABLE local_table_2(id UInt32, val String) ENGINE = TinyLog;
 ''')
 
         node1.query("INSERT INTO local_table VALUES (1, 'node1')")
@@ -24,6 +25,7 @@ CREATE TABLE local_table(id UInt32, val String) ENGINE = TinyLog;
 
         node1.query('''
 CREATE TABLE distributed_table(id UInt32, val String) ENGINE = Distributed(test_cluster, default, local_table);
+CREATE TABLE distributed_table_2(id UInt32, val String) ENGINE = Distributed(test_cluster, default, local_table_2);
 CREATE TABLE merge_table(id UInt32, val String) ENGINE = Merge(default, '^distributed_table')
 ''')
 
@@ -49,6 +51,16 @@ def test_filtering(started_cluster):
 
     assert node1.query("SELECT id + 1 FROM merge_table WHERE val = 'node1'").rstrip() == '2'
 
+
+def test_select_table_name_from_merge_over_distributed(started_cluster):
+
+    node1.query("INSERT INTO local_table_2 VALUES (1, 'node1')")
+    node2.query("INSERT INTO local_table_2 VALUES (2, 'node2')")
+
+    node1.query("select _table == 'distributed_table' from merge_table")
+    node1.query("select * from (select _table == 'distributed_table' from merge_table limit 1)")
+
+
 if __name__ == '__main__':
     with contextmanager(started_cluster)() as cluster:
         for name, instance in cluster.instances.items():