1 | #include <Interpreters/InterpreterCreateQuotaQuery.h> |
2 | #include <Parsers/ASTCreateQuotaQuery.h> |
3 | #include <Parsers/ASTRoleList.h> |
4 | #include <Interpreters/Context.h> |
5 | #include <Access/AccessControlManager.h> |
6 | #include <ext/range.h> |
7 | #include <boost/range/algorithm/find_if.hpp> |
8 | #include <boost/range/algorithm/upper_bound.hpp> |
9 | #include <boost/range/algorithm/sort.hpp> |
10 | |
11 | |
12 | namespace DB |
13 | { |
14 | BlockIO InterpreterCreateQuotaQuery::execute() |
15 | { |
16 | context.checkQuotaManagementIsAllowed(); |
17 | const auto & query = query_ptr->as<const ASTCreateQuotaQuery &>(); |
18 | auto & access_control = context.getAccessControlManager(); |
19 | |
20 | if (query.alter) |
21 | { |
22 | auto update_func = [&](const AccessEntityPtr & entity) -> AccessEntityPtr |
23 | { |
24 | auto updated_quota = typeid_cast<std::shared_ptr<Quota>>(entity->clone()); |
25 | updateQuotaFromQuery(*updated_quota, query); |
26 | return updated_quota; |
27 | }; |
28 | if (query.if_exists) |
29 | { |
30 | if (auto id = access_control.find<Quota>(query.name)) |
31 | access_control.tryUpdate(*id, update_func); |
32 | } |
33 | else |
34 | access_control.update(access_control.getID<Quota>(query.name), update_func); |
35 | } |
36 | else |
37 | { |
38 | auto new_quota = std::make_shared<Quota>(); |
39 | updateQuotaFromQuery(*new_quota, query); |
40 | |
41 | if (query.if_not_exists) |
42 | access_control.tryInsert(new_quota); |
43 | else if (query.or_replace) |
44 | access_control.insertOrReplace(new_quota); |
45 | else |
46 | access_control.insert(new_quota); |
47 | } |
48 | |
49 | return {}; |
50 | } |
51 | |
52 | |
53 | void InterpreterCreateQuotaQuery::updateQuotaFromQuery(Quota & quota, const ASTCreateQuotaQuery & query) |
54 | { |
55 | if (query.alter) |
56 | { |
57 | if (!query.new_name.empty()) |
58 | quota.setName(query.new_name); |
59 | } |
60 | else |
61 | quota.setName(query.name); |
62 | |
63 | if (query.key_type) |
64 | quota.key_type = *query.key_type; |
65 | |
66 | auto & quota_all_limits = quota.all_limits; |
67 | for (const auto & query_limits : query.all_limits) |
68 | { |
69 | auto duration = query_limits.duration; |
70 | |
71 | auto it = boost::range::find_if(quota_all_limits, [&](const Quota::Limits & x) { return x.duration == duration; }); |
72 | if (query_limits.unset_tracking) |
73 | { |
74 | if (it != quota_all_limits.end()) |
75 | quota_all_limits.erase(it); |
76 | continue; |
77 | } |
78 | |
79 | if (it == quota_all_limits.end()) |
80 | { |
81 | /// We keep `all_limits` sorted by duration. |
82 | it = quota_all_limits.insert( |
83 | boost::range::upper_bound( |
84 | quota_all_limits, |
85 | duration, |
86 | [](const std::chrono::seconds & lhs, const Quota::Limits & rhs) { return lhs < rhs.duration; }), |
87 | Quota::Limits{}); |
88 | it->duration = duration; |
89 | } |
90 | |
91 | auto & quota_limits = *it; |
92 | quota_limits.randomize_interval = query_limits.randomize_interval; |
93 | for (auto resource_type : ext::range(Quota::MAX_RESOURCE_TYPE)) |
94 | { |
95 | if (query_limits.max[resource_type]) |
96 | quota_limits.max[resource_type] = *query_limits.max[resource_type]; |
97 | } |
98 | } |
99 | |
100 | if (query.roles) |
101 | { |
102 | const auto & query_roles = *query.roles; |
103 | |
104 | /// We keep `roles` sorted. |
105 | quota.roles = query_roles.roles; |
106 | if (query_roles.current_user) |
107 | quota.roles.push_back(context.getClientInfo().current_user); |
108 | boost::range::sort(quota.roles); |
109 | quota.roles.erase(std::unique(quota.roles.begin(), quota.roles.end()), quota.roles.end()); |
110 | |
111 | quota.all_roles = query_roles.all_roles; |
112 | |
113 | /// We keep `except_roles` sorted. |
114 | quota.except_roles = query_roles.except_roles; |
115 | if (query_roles.except_current_user) |
116 | quota.except_roles.push_back(context.getClientInfo().current_user); |
117 | boost::range::sort(quota.except_roles); |
118 | quota.except_roles.erase(std::unique(quota.except_roles.begin(), quota.except_roles.end()), quota.except_roles.end()); |
119 | } |
120 | } |
121 | } |
122 | |