1 | #include <Functions/IFunction.h> |
2 | #include <Functions/FunctionFactory.h> |
3 | #include <DataTypes/DataTypeArray.h> |
4 | #include <DataTypes/DataTypeString.h> |
5 | #include <DataTypes/DataTypeTuple.h> |
6 | #include <DataTypes/DataTypeUUID.h> |
7 | #include <Columns/ColumnArray.h> |
8 | #include <Columns/ColumnConst.h> |
9 | #include <Columns/ColumnString.h> |
10 | #include <Columns/ColumnTuple.h> |
11 | #include <Interpreters/Context.h> |
12 | #include <Access/RowPolicyContext.h> |
13 | #include <Access/AccessControlManager.h> |
14 | #include <ext/range.h> |
15 | |
16 | |
17 | namespace DB |
18 | { |
19 | namespace ErrorCodes |
20 | { |
21 | extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; |
22 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
23 | } |
24 | |
25 | |
26 | /// The currentRowPolicies() function can be called with 0..2 arguments: |
27 | /// currentRowPolicies() returns array of tuples (database, table_name, row_policy_name) for all the row policies applied for the current user; |
28 | /// currentRowPolicies(table_name) is equivalent to currentRowPolicies(currentDatabase(), table_name); |
29 | /// currentRowPolicies(database, table_name) returns array of names of the row policies applied to a specific table and for the current user. |
30 | class FunctionCurrentRowPolicies : public IFunction |
31 | { |
32 | public: |
33 | static constexpr auto name = "currentRowPolicies" ; |
34 | |
35 | static FunctionPtr create(const Context & context_) { return std::make_shared<FunctionCurrentRowPolicies>(context_); } |
36 | explicit FunctionCurrentRowPolicies(const Context & context_) : context(context_) {} |
37 | |
38 | String getName() const override { return name; } |
39 | size_t getNumberOfArguments() const override { return 0; } |
40 | bool isVariadic() const override { return true; } |
41 | |
42 | void checkNumberOfArgumentsIfVariadic(size_t number_of_arguments) const override |
43 | { |
44 | if (number_of_arguments > 2) |
45 | throw Exception("Number of arguments for function " + String(name) + " doesn't match: passed " |
46 | + toString(number_of_arguments) + ", should be 0..2" , |
47 | ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH); |
48 | } |
49 | |
50 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
51 | { |
52 | if (arguments.empty()) |
53 | return std::make_shared<DataTypeArray>(std::make_shared<DataTypeTuple>( |
54 | DataTypes{std::make_shared<DataTypeString>(), std::make_shared<DataTypeString>(), std::make_shared<DataTypeString>()})); |
55 | else |
56 | return std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>()); |
57 | } |
58 | |
59 | bool isDeterministic() const override { return false; } |
60 | |
61 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result_pos, size_t input_rows_count) override |
62 | { |
63 | if (arguments.empty()) |
64 | { |
65 | auto database_column = ColumnString::create(); |
66 | auto table_name_column = ColumnString::create(); |
67 | auto policy_name_column = ColumnString::create(); |
68 | for (const auto & policy_id : context.getRowPolicy()->getCurrentPolicyIDs()) |
69 | { |
70 | const auto policy = context.getAccessControlManager().tryRead<RowPolicy>(policy_id); |
71 | if (policy) |
72 | { |
73 | const String database = policy->getDatabase(); |
74 | const String table_name = policy->getTableName(); |
75 | const String policy_name = policy->getName(); |
76 | database_column->insertData(database.data(), database.length()); |
77 | table_name_column->insertData(table_name.data(), table_name.length()); |
78 | policy_name_column->insertData(policy_name.data(), policy_name.length()); |
79 | } |
80 | } |
81 | auto offset_column = ColumnArray::ColumnOffsets::create(); |
82 | offset_column->insertValue(policy_name_column->size()); |
83 | block.getByPosition(result_pos).column = ColumnConst::create( |
84 | ColumnArray::create( |
85 | ColumnTuple::create(Columns{std::move(database_column), std::move(table_name_column), std::move(policy_name_column)}), |
86 | std::move(offset_column)), |
87 | input_rows_count); |
88 | return; |
89 | } |
90 | |
91 | const IColumn * database_column = nullptr; |
92 | if (arguments.size() == 2) |
93 | { |
94 | const auto & database_column_with_type = block.getByPosition(arguments[0]); |
95 | if (!isStringOrFixedString(database_column_with_type.type)) |
96 | throw Exception{"The first argument of function " + String(name) |
97 | + " should be a string containing database name, illegal type: " |
98 | + database_column_with_type.type->getName(), |
99 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; |
100 | database_column = database_column_with_type.column.get(); |
101 | } |
102 | |
103 | const auto & table_name_column_with_type = block.getByPosition(arguments[arguments.size() - 1]); |
104 | if (!isStringOrFixedString(table_name_column_with_type.type)) |
105 | throw Exception{"The" + String(database_column ? " last" : "" ) + " argument of function " + String(name) |
106 | + " should be a string containing table name, illegal type: " + table_name_column_with_type.type->getName(), |
107 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; |
108 | const IColumn * table_name_column = table_name_column_with_type.column.get(); |
109 | |
110 | auto policy_name_column = ColumnString::create(); |
111 | auto offset_column = ColumnArray::ColumnOffsets::create(); |
112 | for (const auto i : ext::range(0, input_rows_count)) |
113 | { |
114 | String database = database_column ? database_column->getDataAt(i).toString() : context.getCurrentDatabase(); |
115 | String table_name = table_name_column->getDataAt(i).toString(); |
116 | for (const auto & policy_id : context.getRowPolicy()->getCurrentPolicyIDs(database, table_name)) |
117 | { |
118 | const auto policy = context.getAccessControlManager().tryRead<RowPolicy>(policy_id); |
119 | if (policy) |
120 | { |
121 | const String policy_name = policy->getName(); |
122 | policy_name_column->insertData(policy_name.data(), policy_name.length()); |
123 | } |
124 | } |
125 | offset_column->insertValue(policy_name_column->size()); |
126 | } |
127 | |
128 | block.getByPosition(result_pos).column = ColumnArray::create(std::move(policy_name_column), std::move(offset_column)); |
129 | } |
130 | |
131 | private: |
132 | const Context & context; |
133 | }; |
134 | |
135 | |
136 | /// The currentRowPolicyIDs() function can be called with 0..2 arguments: |
137 | /// currentRowPolicyIDs() returns array of IDs of all the row policies applied for the current user; |
138 | /// currentRowPolicyIDs(table_name) is equivalent to currentRowPolicyIDs(currentDatabase(), table_name); |
139 | /// currentRowPolicyIDs(database, table_name) returns array of IDs of the row policies applied to a specific table and for the current user. |
140 | class FunctionCurrentRowPolicyIDs : public IFunction |
141 | { |
142 | public: |
143 | static constexpr auto name = "currentRowPolicyIDs" ; |
144 | |
145 | static FunctionPtr create(const Context & context_) { return std::make_shared<FunctionCurrentRowPolicyIDs>(context_); } |
146 | explicit FunctionCurrentRowPolicyIDs(const Context & context_) : context(context_) {} |
147 | |
148 | String getName() const override { return name; } |
149 | size_t getNumberOfArguments() const override { return 0; } |
150 | bool isVariadic() const override { return true; } |
151 | |
152 | void checkNumberOfArgumentsIfVariadic(size_t number_of_arguments) const override |
153 | { |
154 | if (number_of_arguments > 2) |
155 | throw Exception("Number of arguments for function " + String(name) + " doesn't match: passed " |
156 | + toString(number_of_arguments) + ", should be 0..2" , |
157 | ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH); |
158 | } |
159 | |
160 | DataTypePtr getReturnTypeImpl(const DataTypes & /* arguments */) const override |
161 | { |
162 | return std::make_shared<DataTypeArray>(std::make_shared<DataTypeUUID>()); |
163 | } |
164 | |
165 | bool isDeterministic() const override { return false; } |
166 | |
167 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result_pos, size_t input_rows_count) override |
168 | { |
169 | if (arguments.empty()) |
170 | { |
171 | auto policy_id_column = ColumnVector<UInt128>::create(); |
172 | for (const auto & policy_id : context.getRowPolicy()->getCurrentPolicyIDs()) |
173 | policy_id_column->insertValue(policy_id); |
174 | auto offset_column = ColumnArray::ColumnOffsets::create(); |
175 | offset_column->insertValue(policy_id_column->size()); |
176 | block.getByPosition(result_pos).column |
177 | = ColumnConst::create(ColumnArray::create(std::move(policy_id_column), std::move(offset_column)), input_rows_count); |
178 | return; |
179 | } |
180 | |
181 | const IColumn * database_column = nullptr; |
182 | if (arguments.size() == 2) |
183 | { |
184 | const auto & database_column_with_type = block.getByPosition(arguments[0]); |
185 | if (!isStringOrFixedString(database_column_with_type.type)) |
186 | throw Exception{"The first argument of function " + String(name) |
187 | + " should be a string containing database name, illegal type: " |
188 | + database_column_with_type.type->getName(), |
189 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; |
190 | database_column = database_column_with_type.column.get(); |
191 | } |
192 | |
193 | const auto & table_name_column_with_type = block.getByPosition(arguments[arguments.size() - 1]); |
194 | if (!isStringOrFixedString(table_name_column_with_type.type)) |
195 | throw Exception{"The" + String(database_column ? " last" : "" ) + " argument of function " + String(name) |
196 | + " should be a string containing table name, illegal type: " + table_name_column_with_type.type->getName(), |
197 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; |
198 | const IColumn * table_name_column = table_name_column_with_type.column.get(); |
199 | |
200 | auto policy_id_column = ColumnVector<UInt128>::create(); |
201 | auto offset_column = ColumnArray::ColumnOffsets::create(); |
202 | for (const auto i : ext::range(0, input_rows_count)) |
203 | { |
204 | String database = database_column ? database_column->getDataAt(i).toString() : context.getCurrentDatabase(); |
205 | String table_name = table_name_column->getDataAt(i).toString(); |
206 | for (const auto & policy_id : context.getRowPolicy()->getCurrentPolicyIDs(database, table_name)) |
207 | policy_id_column->insertValue(policy_id); |
208 | offset_column->insertValue(policy_id_column->size()); |
209 | } |
210 | |
211 | block.getByPosition(result_pos).column = ColumnArray::create(std::move(policy_id_column), std::move(offset_column)); |
212 | } |
213 | |
214 | private: |
215 | const Context & context; |
216 | }; |
217 | |
218 | |
219 | void registerFunctionCurrentRowPolicies(FunctionFactory & factory) |
220 | { |
221 | factory.registerFunction<FunctionCurrentRowPolicies>(); |
222 | factory.registerFunction<FunctionCurrentRowPolicyIDs>(); |
223 | } |
224 | |
225 | } |
226 | |