1 | #pragma once |
2 | |
3 | #include <Access/IAccessEntity.h> |
4 | #include <chrono> |
5 | |
6 | |
7 | |
8 | namespace DB |
9 | { |
10 | /** Quota for resources consumption for specific interval. |
11 | * Used to limit resource usage by user. |
12 | * Quota is applied "softly" - could be slightly exceed, because it is checked usually only on each block of processed data. |
13 | * Accumulated values are not persisted and are lost on server restart. |
14 | * Quota is local to server, |
15 | * but for distributed queries, accumulated values for read rows and bytes |
16 | * are collected from all participating servers and accumulated locally. |
17 | */ |
18 | struct Quota : public IAccessEntity |
19 | { |
20 | enum ResourceType |
21 | { |
22 | QUERIES, /// Number of queries. |
23 | ERRORS, /// Number of queries with exceptions. |
24 | RESULT_ROWS, /// Number of rows returned as result. |
25 | RESULT_BYTES, /// Number of bytes returned as result. |
26 | READ_ROWS, /// Number of rows read from tables. |
27 | READ_BYTES, /// Number of bytes read from tables. |
28 | EXECUTION_TIME, /// Total amount of query execution time in nanoseconds. |
29 | }; |
30 | static constexpr size_t MAX_RESOURCE_TYPE = 7; |
31 | |
32 | using ResourceAmount = UInt64; |
33 | static constexpr ResourceAmount UNLIMITED = 0; /// 0 means unlimited. |
34 | |
35 | /// Amount of resources available to consume for each duration. |
36 | struct Limits |
37 | { |
38 | ResourceAmount max[MAX_RESOURCE_TYPE]; |
39 | std::chrono::seconds duration = std::chrono::seconds::zero(); |
40 | |
41 | /// Intervals can be randomized (to avoid DoS if intervals for many users end at one time). |
42 | bool randomize_interval = false; |
43 | |
44 | Limits(); |
45 | friend bool operator ==(const Limits & lhs, const Limits & rhs); |
46 | friend bool operator !=(const Limits & lhs, const Limits & rhs) { return !(lhs == rhs); } |
47 | }; |
48 | |
49 | std::vector<Limits> all_limits; |
50 | |
51 | /// Key to share quota consumption. |
52 | /// Users with the same key share the same amount of resource. |
53 | enum class KeyType |
54 | { |
55 | NONE, /// All users share the same quota. |
56 | USER_NAME, /// Connections with the same user name share the same quota. |
57 | IP_ADDRESS, /// Connections from the same IP share the same quota. |
58 | CLIENT_KEY, /// Client should explicitly supply a key to use. |
59 | CLIENT_KEY_OR_USER_NAME, /// Same as CLIENT_KEY, but use USER_NAME if the client doesn't supply a key. |
60 | CLIENT_KEY_OR_IP_ADDRESS, /// Same as CLIENT_KEY, but use IP_ADDRESS if the client doesn't supply a key. |
61 | }; |
62 | static constexpr size_t MAX_KEY_TYPE = 6; |
63 | KeyType key_type = KeyType::NONE; |
64 | |
65 | /// Which roles or users should use this quota. |
66 | Strings roles; |
67 | bool all_roles = false; |
68 | Strings except_roles; |
69 | |
70 | bool equal(const IAccessEntity & other) const override; |
71 | std::shared_ptr<IAccessEntity> clone() const override { return cloneImpl<Quota>(); } |
72 | |
73 | static const char * getNameOfResourceType(ResourceType resource_type); |
74 | static const char * resourceTypeToKeyword(ResourceType resource_type); |
75 | static const char * resourceTypeToColumnName(ResourceType resource_type); |
76 | static const char * getNameOfKeyType(KeyType key_type); |
77 | static double executionTimeToSeconds(ResourceAmount ns); |
78 | static ResourceAmount secondsToExecutionTime(double s); |
79 | }; |
80 | |
81 | |
82 | inline const char * Quota::getNameOfResourceType(ResourceType resource_type) |
83 | { |
84 | switch (resource_type) |
85 | { |
86 | case Quota::QUERIES: return "queries" ; |
87 | case Quota::ERRORS: return "errors" ; |
88 | case Quota::RESULT_ROWS: return "result rows" ; |
89 | case Quota::RESULT_BYTES: return "result bytes" ; |
90 | case Quota::READ_ROWS: return "read rows" ; |
91 | case Quota::READ_BYTES: return "read bytes" ; |
92 | case Quota::EXECUTION_TIME: return "execution time" ; |
93 | } |
94 | __builtin_unreachable(); |
95 | } |
96 | |
97 | |
98 | inline const char * Quota::resourceTypeToKeyword(ResourceType resource_type) |
99 | { |
100 | switch (resource_type) |
101 | { |
102 | case Quota::QUERIES: return "QUERIES" ; |
103 | case Quota::ERRORS: return "ERRORS" ; |
104 | case Quota::RESULT_ROWS: return "RESULT ROWS" ; |
105 | case Quota::RESULT_BYTES: return "RESULT BYTES" ; |
106 | case Quota::READ_ROWS: return "READ ROWS" ; |
107 | case Quota::READ_BYTES: return "READ BYTES" ; |
108 | case Quota::EXECUTION_TIME: return "EXECUTION TIME" ; |
109 | } |
110 | __builtin_unreachable(); |
111 | } |
112 | |
113 | |
114 | inline const char * Quota::getNameOfKeyType(KeyType key_type) |
115 | { |
116 | switch (key_type) |
117 | { |
118 | case KeyType::NONE: return "none" ; |
119 | case KeyType::USER_NAME: return "user name" ; |
120 | case KeyType::IP_ADDRESS: return "ip address" ; |
121 | case KeyType::CLIENT_KEY: return "client key" ; |
122 | case KeyType::CLIENT_KEY_OR_USER_NAME: return "client key or user name" ; |
123 | case KeyType::CLIENT_KEY_OR_IP_ADDRESS: return "client key or ip address" ; |
124 | } |
125 | __builtin_unreachable(); |
126 | } |
127 | |
128 | |
129 | inline double Quota::executionTimeToSeconds(ResourceAmount ns) |
130 | { |
131 | return std::chrono::duration_cast<std::chrono::duration<double>>(std::chrono::nanoseconds{ns}).count(); |
132 | } |
133 | |
134 | inline Quota::ResourceAmount Quota::secondsToExecutionTime(double s) |
135 | { |
136 | return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::duration<double>(s)).count(); |
137 | } |
138 | |
139 | |
140 | using QuotaPtr = std::shared_ptr<const Quota>; |
141 | } |
142 | |