1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
|
#include <Common/checkStackSize.h>
#include <Interpreters/ApplyWithSubqueryVisitor.h>
#include <Interpreters/IdentifierSemantic.h>
#include <Interpreters/StorageID.h>
#include <Interpreters/misc.h>
#include <Parsers/ASTFunction.h>
#include <Parsers/ASTIdentifier.h>
#include <Parsers/ASTSelectQuery.h>
#include <Parsers/ASTSelectWithUnionQuery.h>
#include <Parsers/ASTSubquery.h>
#include <Parsers/ASTTablesInSelectQuery.h>
#include <Parsers/ASTWithElement.h>
namespace DB
{
void ApplyWithSubqueryVisitor::visit(ASTPtr & ast, const Data & data)
{
checkStackSize();
if (auto * node_select = ast->as<ASTSelectQuery>())
visit(*node_select, data);
else
{
for (auto & child : ast->children)
visit(child, data);
if (auto * node_func = ast->as<ASTFunction>())
visit(*node_func, data);
else if (auto * node_table = ast->as<ASTTableExpression>())
visit(*node_table, data);
}
}
void ApplyWithSubqueryVisitor::visit(ASTSelectQuery & ast, const Data & data)
{
std::optional<Data> new_data;
if (auto with = ast.with())
{
for (auto & child : with->children)
{
visit(child, new_data ? *new_data : data);
if (auto * ast_with_elem = child->as<ASTWithElement>())
{
if (!new_data)
new_data = data;
new_data->subqueries[ast_with_elem->name] = ast_with_elem->subquery;
}
}
}
for (auto & child : ast.children)
{
if (child != ast.with())
visit(child, new_data ? *new_data : data);
}
}
void ApplyWithSubqueryVisitor::visit(ASTSelectWithUnionQuery & ast, const Data & data)
{
for (auto & child : ast.children)
visit(child, data);
}
void ApplyWithSubqueryVisitor::visit(ASTTableExpression & table, const Data & data)
{
if (table.database_and_table_name)
{
auto table_id = table.database_and_table_name->as<ASTTableIdentifier>()->getTableId();
if (table_id.database_name.empty())
{
auto subquery_it = data.subqueries.find(table_id.table_name);
if (subquery_it != data.subqueries.end())
{
auto old_alias = table.database_and_table_name->tryGetAlias();
table.children.clear();
table.database_and_table_name.reset();
table.subquery = subquery_it->second->clone();
table.subquery->as<ASTSubquery &>().cte_name = table_id.table_name;
if (!old_alias.empty())
table.subquery->setAlias(old_alias);
table.children.emplace_back(table.subquery);
}
}
}
}
void ApplyWithSubqueryVisitor::visit(ASTFunction & func, const Data & data)
{
/// Special CTE case, where the right argument of IN is alias (ASTIdentifier) from WITH clause.
if (checkFunctionIsInOrGlobalInOperator(func))
{
auto & ast = func.arguments->children.at(1);
if (const auto * identifier = ast->as<ASTIdentifier>())
{
if (identifier->isShort())
{
/// Clang-tidy is wrong on this line, because `func.arguments->children.at(1)` gets replaced before last use of `name`.
auto name = identifier->shortName(); // NOLINT
auto subquery_it = data.subqueries.find(name);
if (subquery_it != data.subqueries.end())
{
auto old_alias = func.arguments->children[1]->tryGetAlias();
func.arguments->children[1] = subquery_it->second->clone();
func.arguments->children[1]->as<ASTSubquery &>().cte_name = name;
if (!old_alias.empty())
func.arguments->children[1]->setAlias(old_alias);
}
}
}
}
}
}
|
