| 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 | |