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
11namespace DB
12{
13class AccessControlManager;
14
15
16/// Stores information how much amount of resources have been consumed and how much are left.
17class QuotaContextFactory
18{
19public:
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
26private:
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