1 | #pragma once |
---|---|
2 | |
3 | #include <Access/QuotaContext.h> |
4 | #include <Access/IAccessStorage.h> |
5 | #include <memory> |
6 | #include <mutex> |
7 | #include <unordered_map> |
8 | #include <unordered_set> |
9 | |
10 | |
11 | namespace DB |
12 | { |
13 | class AccessControlManager; |
14 | |
15 | |
16 | /// Stores information how much amount of resources have been consumed and how much are left. |
17 | class QuotaContextFactory |
18 | { |
19 | public: |
20 | QuotaContextFactory(const AccessControlManager & access_control_manager_); |
21 | ~QuotaContextFactory(); |
22 | |
23 | QuotaContextPtr createContext(const String & user_name, const Poco::Net::IPAddress & address, const String & client_key); |
24 | std::vector<QuotaUsageInfo> getUsageInfo() const; |
25 | |
26 | private: |
27 | using Interval = QuotaContext::Interval; |
28 | using Intervals = QuotaContext::Intervals; |
29 | |
30 | struct QuotaInfo |
31 | { |
32 | QuotaInfo(const QuotaPtr & quota_, const UUID & quota_id_) { setQuota(quota_, quota_id_); } |
33 | void setQuota(const QuotaPtr & quota_, const UUID & quota_id_); |
34 | |
35 | bool canUseWithContext(const QuotaContext & context) const; |
36 | String calculateKey(const QuotaContext & context) const; |
37 | std::shared_ptr<const Intervals> getOrBuildIntervals(const String & key); |
38 | std::shared_ptr<const Intervals> rebuildIntervals(const String & key); |
39 | void rebuildAllIntervals(); |
40 | |
41 | QuotaPtr quota; |
42 | UUID quota_id; |
43 | std::unordered_set<String> roles; |
44 | bool all_roles = false; |
45 | std::unordered_set<String> except_roles; |
46 | std::unordered_map<String /* quota key */, std::shared_ptr<const Intervals>> key_to_intervals; |
47 | }; |
48 | |
49 | void ensureAllQuotasRead(); |
50 | void quotaAddedOrChanged(const UUID & quota_id, const std::shared_ptr<const Quota> & new_quota); |
51 | void quotaRemoved(const UUID & quota_id); |
52 | void chooseQuotaForAllContexts(); |
53 | void chooseQuotaForContext(const std::shared_ptr<QuotaContext> & context); |
54 | |
55 | const AccessControlManager & access_control_manager; |
56 | mutable std::mutex mutex; |
57 | std::unordered_map<UUID /* quota id */, QuotaInfo> all_quotas; |
58 | bool all_quotas_read = false; |
59 | IAccessStorage::SubscriptionPtr subscription; |
60 | std::vector<std::weak_ptr<QuotaContext>> contexts; |
61 | }; |
62 | } |
63 |