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
17namespace DB
18{
19namespace 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.
30class FunctionCurrentRowPolicies : public IFunction
31{
32public:
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
131private:
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.
140class FunctionCurrentRowPolicyIDs : public IFunction
141{
142public:
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
214private:
215 const Context & context;
216};
217
218
219void registerFunctionCurrentRowPolicies(FunctionFactory & factory)
220{
221 factory.registerFunction<FunctionCurrentRowPolicies>();
222 factory.registerFunction<FunctionCurrentRowPolicyIDs>();
223}
224
225}
226