1 | #pragma once |
2 | |
3 | #include <Common/typeid_cast.h> |
4 | #include <Parsers/ASTQueryWithTableAndOutput.h> |
5 | #include <Parsers/ASTRenameQuery.h> |
6 | #include <Parsers/ASTIdentifier.h> |
7 | #include <Parsers/ASTSelectQuery.h> |
8 | #include <Parsers/ASTSubquery.h> |
9 | #include <Parsers/ASTSelectWithUnionQuery.h> |
10 | #include <Parsers/ASTTablesInSelectQuery.h> |
11 | #include <Parsers/ASTFunction.h> |
12 | #include <Parsers/DumpASTNode.h> |
13 | #include <Interpreters/DatabaseAndTableWithAlias.h> |
14 | #include <Interpreters/IdentifierSemantic.h> |
15 | |
16 | namespace DB |
17 | { |
18 | |
19 | /// Visitors consist of functions with unified interface 'void visit(Casted & x, ASTPtr & y)', there x is y, successfully casted to Casted. |
20 | /// Both types and fuction could have const specifiers. The second argument is used by visitor to replaces AST node (y) if needed. |
21 | |
22 | /// Visits AST nodes, add default database to tables if not set. There's different logic for DDLs and selects. |
23 | class AddDefaultDatabaseVisitor |
24 | { |
25 | public: |
26 | AddDefaultDatabaseVisitor(const String & database_name_, std::ostream * ostr_ = nullptr) |
27 | : database_name(database_name_), |
28 | visit_depth(0), |
29 | ostr(ostr_) |
30 | {} |
31 | |
32 | void visitDDL(ASTPtr & ast) const |
33 | { |
34 | visitDDLChildren(ast); |
35 | |
36 | if (!tryVisitDynamicCast<ASTQueryWithTableAndOutput>(ast) && |
37 | !tryVisitDynamicCast<ASTRenameQuery>(ast)) |
38 | {} |
39 | } |
40 | |
41 | void visit(ASTPtr & ast) const |
42 | { |
43 | if (!tryVisit<ASTSelectQuery>(ast) && |
44 | !tryVisit<ASTSelectWithUnionQuery>(ast) && |
45 | !tryVisit<ASTFunction>(ast)) |
46 | visitChildren(*ast); |
47 | } |
48 | |
49 | void visit(ASTSelectQuery & select) const |
50 | { |
51 | ASTPtr unused; |
52 | visit(select, unused); |
53 | } |
54 | |
55 | void visit(ASTSelectWithUnionQuery & select) const |
56 | { |
57 | ASTPtr unused; |
58 | visit(select, unused); |
59 | } |
60 | |
61 | private: |
62 | const String database_name; |
63 | mutable size_t visit_depth; |
64 | std::ostream * ostr; |
65 | |
66 | void visit(ASTSelectWithUnionQuery & select, ASTPtr &) const |
67 | { |
68 | for (auto & child : select.list_of_selects->children) |
69 | tryVisit<ASTSelectQuery>(child); |
70 | } |
71 | |
72 | void visit(ASTSelectQuery & select, ASTPtr &) const |
73 | { |
74 | if (select.tables()) |
75 | tryVisit<ASTTablesInSelectQuery>(select.refTables()); |
76 | |
77 | visitChildren(select); |
78 | } |
79 | |
80 | void visit(ASTTablesInSelectQuery & tables, ASTPtr &) const |
81 | { |
82 | for (auto & child : tables.children) |
83 | tryVisit<ASTTablesInSelectQueryElement>(child); |
84 | } |
85 | |
86 | void visit(ASTTablesInSelectQueryElement & tables_element, ASTPtr &) const |
87 | { |
88 | if (tables_element.table_expression) |
89 | tryVisit<ASTTableExpression>(tables_element.table_expression); |
90 | } |
91 | |
92 | void visit(ASTTableExpression & table_expression, ASTPtr &) const |
93 | { |
94 | if (table_expression.database_and_table_name) |
95 | tryVisit<ASTIdentifier>(table_expression.database_and_table_name); |
96 | else if (table_expression.subquery) |
97 | tryVisit<ASTSubquery>(table_expression.subquery); |
98 | } |
99 | |
100 | /// @note It expects that only table (not column) identifiers are visited. |
101 | void visit(const ASTIdentifier & identifier, ASTPtr & ast) const |
102 | { |
103 | if (!identifier.compound()) |
104 | ast = createTableIdentifier(database_name, identifier.name); |
105 | } |
106 | |
107 | void visit(ASTSubquery & subquery, ASTPtr &) const |
108 | { |
109 | tryVisit<ASTSelectWithUnionQuery>(subquery.children[0]); |
110 | } |
111 | |
112 | void visit(ASTFunction & function, ASTPtr &) const |
113 | { |
114 | bool is_operator_in = false; |
115 | for (auto name : {"in" , "notIn" , "globalIn" , "globalNotIn" }) |
116 | { |
117 | if (function.name == name) |
118 | { |
119 | is_operator_in = true; |
120 | break; |
121 | } |
122 | } |
123 | |
124 | for (auto & child : function.children) |
125 | { |
126 | if (child.get() == function.arguments.get()) |
127 | { |
128 | for (size_t i = 0; i < child->children.size(); ++i) |
129 | { |
130 | if (is_operator_in && i == 1) |
131 | { |
132 | /// Second argument of the "in" function (or similar) may be a table name or a subselect. |
133 | /// Rewrite the table name or descend into subselect. |
134 | if (!tryVisit<ASTIdentifier>(child->children[i])) |
135 | visit(child->children[i]); |
136 | } |
137 | else |
138 | visit(child->children[i]); |
139 | } |
140 | } |
141 | else |
142 | visit(child); |
143 | } |
144 | } |
145 | |
146 | void visitChildren(IAST & ast) const |
147 | { |
148 | for (auto & child : ast.children) |
149 | visit(child); |
150 | } |
151 | |
152 | template <typename T> |
153 | bool tryVisit(ASTPtr & ast) const |
154 | { |
155 | if (T * t = typeid_cast<T *>(ast.get())) |
156 | { |
157 | DumpASTNode dump(*ast, ostr, visit_depth, "addDefaultDatabaseName" ); |
158 | visit(*t, ast); |
159 | return true; |
160 | } |
161 | return false; |
162 | } |
163 | |
164 | |
165 | void visitDDL(ASTQueryWithTableAndOutput & node, ASTPtr &) const |
166 | { |
167 | if (node.database.empty()) |
168 | node.database = database_name; |
169 | } |
170 | |
171 | void visitDDL(ASTRenameQuery & node, ASTPtr &) const |
172 | { |
173 | for (ASTRenameQuery::Element & elem : node.elements) |
174 | { |
175 | if (elem.from.database.empty()) |
176 | elem.from.database = database_name; |
177 | if (elem.to.database.empty()) |
178 | elem.to.database = database_name; |
179 | } |
180 | } |
181 | |
182 | void visitDDLChildren(ASTPtr & ast) const |
183 | { |
184 | for (auto & child : ast->children) |
185 | visitDDL(child); |
186 | } |
187 | |
188 | template <typename T> |
189 | bool tryVisitDynamicCast(ASTPtr & ast) const |
190 | { |
191 | if (T * t = dynamic_cast<T *>(ast.get())) |
192 | { |
193 | visitDDL(*t, ast); |
194 | return true; |
195 | } |
196 | return false; |
197 | } |
198 | }; |
199 | |
200 | } |
201 | |