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
12namespace DB
13{
14BlockIO 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
53void 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