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
16namespace 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.
23class AddDefaultDatabaseVisitor
24{
25public:
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
61private:
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