1 | #include <Interpreters/getTableExpressions.h> |
2 | #include <Interpreters/InterpreterSelectWithUnionQuery.h> |
3 | #include <Parsers/ASTTablesInSelectQuery.h> |
4 | #include <Parsers/ASTSelectQuery.h> |
5 | #include <Storages/IStorage.h> |
6 | |
7 | namespace DB |
8 | { |
9 | |
10 | NameSet removeDuplicateColumns(NamesAndTypesList & columns) |
11 | { |
12 | NameSet names; |
13 | for (auto it = columns.begin(); it != columns.end();) |
14 | { |
15 | if (names.emplace(it->name).second) |
16 | ++it; |
17 | else |
18 | columns.erase(it++); |
19 | } |
20 | return names; |
21 | } |
22 | |
23 | std::vector<const ASTTableExpression *> getTableExpressions(const ASTSelectQuery & select_query) |
24 | { |
25 | if (!select_query.tables()) |
26 | return {}; |
27 | |
28 | std::vector<const ASTTableExpression *> tables_expression; |
29 | |
30 | for (const auto & child : select_query.tables()->children) |
31 | { |
32 | const auto * tables_element = child->as<ASTTablesInSelectQueryElement>(); |
33 | |
34 | if (tables_element && tables_element->table_expression) |
35 | tables_expression.emplace_back(tables_element->table_expression->as<ASTTableExpression>()); |
36 | } |
37 | |
38 | return tables_expression; |
39 | } |
40 | |
41 | const ASTTableExpression * getTableExpression(const ASTSelectQuery & select, size_t table_number) |
42 | { |
43 | if (!select.tables()) |
44 | return {}; |
45 | |
46 | const auto & tables_in_select_query = select.tables()->as<ASTTablesInSelectQuery &>(); |
47 | if (tables_in_select_query.children.size() <= table_number) |
48 | return {}; |
49 | |
50 | const auto & tables_element = tables_in_select_query.children[table_number]->as<ASTTablesInSelectQueryElement &>(); |
51 | |
52 | if (!tables_element.table_expression) |
53 | return {}; |
54 | |
55 | return tables_element.table_expression->as<ASTTableExpression>(); |
56 | } |
57 | |
58 | ASTPtr (const ASTSelectQuery & select, size_t table_number) |
59 | { |
60 | if (const ASTTableExpression * table_expression = getTableExpression(select, table_number)) |
61 | { |
62 | if (table_expression->database_and_table_name) |
63 | return table_expression->database_and_table_name; |
64 | |
65 | if (table_expression->table_function) |
66 | return table_expression->table_function; |
67 | |
68 | if (table_expression->subquery) |
69 | return table_expression->subquery->children[0]; |
70 | } |
71 | |
72 | return nullptr; |
73 | } |
74 | |
75 | static NamesAndTypesList getColumnsFromTableExpression(const ASTTableExpression & table_expression, const Context & context, |
76 | NamesAndTypesList & materialized, NamesAndTypesList & aliases, NamesAndTypesList & virtuals) |
77 | { |
78 | NamesAndTypesList names_and_type_list; |
79 | if (table_expression.subquery) |
80 | { |
81 | const auto & subquery = table_expression.subquery->children.at(0); |
82 | names_and_type_list = InterpreterSelectWithUnionQuery::getSampleBlock(subquery, context).getNamesAndTypesList(); |
83 | } |
84 | else if (table_expression.table_function) |
85 | { |
86 | const auto table_function = table_expression.table_function; |
87 | auto query_context = const_cast<Context *>(&context.getQueryContext()); |
88 | const auto & function_storage = query_context->executeTableFunction(table_function); |
89 | auto & columns = function_storage->getColumns(); |
90 | names_and_type_list = columns.getOrdinary(); |
91 | materialized = columns.getMaterialized(); |
92 | aliases = columns.getAliases(); |
93 | virtuals = columns.getVirtuals(); |
94 | } |
95 | else if (table_expression.database_and_table_name) |
96 | { |
97 | DatabaseAndTableWithAlias database_table(table_expression.database_and_table_name); |
98 | const auto & table = context.getTable(database_table.database, database_table.table); |
99 | auto & columns = table->getColumns(); |
100 | names_and_type_list = columns.getOrdinary(); |
101 | materialized = columns.getMaterialized(); |
102 | aliases = columns.getAliases(); |
103 | virtuals = columns.getVirtuals(); |
104 | } |
105 | |
106 | return names_and_type_list; |
107 | } |
108 | |
109 | NamesAndTypesList getColumnsFromTableExpression(const ASTTableExpression & table_expression, const Context & context) |
110 | { |
111 | NamesAndTypesList materialized; |
112 | NamesAndTypesList aliases; |
113 | NamesAndTypesList virtuals; |
114 | return getColumnsFromTableExpression(table_expression, context, materialized, aliases, virtuals); |
115 | } |
116 | |
117 | std::vector<TableWithColumnNames> getDatabaseAndTablesWithColumnNames(const std::vector<const ASTTableExpression *> & table_expressions, |
118 | const Context & context, bool remove_duplicates) |
119 | { |
120 | std::vector<TableWithColumnNames> tables_with_columns; |
121 | |
122 | if (!table_expressions.empty()) |
123 | { |
124 | String current_database = context.getCurrentDatabase(); |
125 | |
126 | for (const ASTTableExpression * table_expression : table_expressions) |
127 | { |
128 | DatabaseAndTableWithAlias table_name(*table_expression, current_database); |
129 | |
130 | NamesAndTypesList materialized; |
131 | NamesAndTypesList aliases; |
132 | NamesAndTypesList virtuals; |
133 | NamesAndTypesList names_and_types = getColumnsFromTableExpression(*table_expression, context, materialized, aliases, virtuals); |
134 | |
135 | if (remove_duplicates) |
136 | removeDuplicateColumns(names_and_types); |
137 | |
138 | tables_with_columns.emplace_back(std::move(table_name), names_and_types.getNames()); |
139 | auto & table = tables_with_columns.back(); |
140 | table.addHiddenColumns(materialized); |
141 | table.addHiddenColumns(aliases); |
142 | table.addHiddenColumns(virtuals); |
143 | } |
144 | } |
145 | |
146 | return tables_with_columns; |
147 | } |
148 | |
149 | } |
150 | |