1 | #include <Functions/IFunctionImpl.h> |
2 | #include <Functions/FunctionFactory.h> |
3 | #include <Interpreters/Context.h> |
4 | #include <DataTypes/DataTypeString.h> |
5 | #include <DataTypes/DataTypeUUID.h> |
6 | #include <Access/QuotaContext.h> |
7 | #include <Core/Field.h> |
8 | |
9 | |
10 | namespace DB |
11 | { |
12 | |
13 | class FunctionCurrentQuota : public IFunction |
14 | { |
15 | const String quota_name; |
16 | |
17 | public: |
18 | static constexpr auto name = "currentQuota" ; |
19 | static FunctionPtr create(const Context & context) |
20 | { |
21 | return std::make_shared<FunctionCurrentQuota>(context.getQuota()->getUsageInfo().quota_name); |
22 | } |
23 | |
24 | explicit FunctionCurrentQuota(const String & quota_name_) : quota_name{quota_name_} |
25 | { |
26 | } |
27 | |
28 | String getName() const override |
29 | { |
30 | return name; |
31 | } |
32 | size_t getNumberOfArguments() const override |
33 | { |
34 | return 0; |
35 | } |
36 | |
37 | DataTypePtr getReturnTypeImpl(const DataTypes & /*arguments*/) const override |
38 | { |
39 | return std::make_shared<DataTypeString>(); |
40 | } |
41 | |
42 | bool isDeterministic() const override { return false; } |
43 | |
44 | void executeImpl(Block & block, const ColumnNumbers &, size_t result, size_t input_rows_count) override |
45 | { |
46 | block.getByPosition(result).column = DataTypeString().createColumnConst(input_rows_count, quota_name); |
47 | } |
48 | }; |
49 | |
50 | |
51 | class FunctionCurrentQuotaId : public IFunction |
52 | { |
53 | const UUID quota_id; |
54 | |
55 | public: |
56 | static constexpr auto name = "currentQuotaID" ; |
57 | static FunctionPtr create(const Context & context) |
58 | { |
59 | return std::make_shared<FunctionCurrentQuotaId>(context.getQuota()->getUsageInfo().quota_id); |
60 | } |
61 | |
62 | explicit FunctionCurrentQuotaId(const UUID quota_id_) : quota_id{quota_id_} |
63 | { |
64 | } |
65 | |
66 | String getName() const override |
67 | { |
68 | return name; |
69 | } |
70 | size_t getNumberOfArguments() const override |
71 | { |
72 | return 0; |
73 | } |
74 | |
75 | DataTypePtr getReturnTypeImpl(const DataTypes & /*arguments*/) const override |
76 | { |
77 | return std::make_shared<DataTypeUUID>(); |
78 | } |
79 | |
80 | bool isDeterministic() const override { return false; } |
81 | |
82 | void executeImpl(Block & block, const ColumnNumbers &, size_t result, size_t input_rows_count) override |
83 | { |
84 | block.getByPosition(result).column = DataTypeUUID().createColumnConst(input_rows_count, quota_id); |
85 | } |
86 | }; |
87 | |
88 | |
89 | class FunctionCurrentQuotaKey : public IFunction |
90 | { |
91 | const String quota_key; |
92 | |
93 | public: |
94 | static constexpr auto name = "currentQuotaKey" ; |
95 | static FunctionPtr create(const Context & context) |
96 | { |
97 | return std::make_shared<FunctionCurrentQuotaKey>(context.getQuota()->getUsageInfo().quota_key); |
98 | } |
99 | |
100 | explicit FunctionCurrentQuotaKey(const String & quota_key_) : quota_key{quota_key_} |
101 | { |
102 | } |
103 | |
104 | String getName() const override |
105 | { |
106 | return name; |
107 | } |
108 | size_t getNumberOfArguments() const override |
109 | { |
110 | return 0; |
111 | } |
112 | |
113 | DataTypePtr getReturnTypeImpl(const DataTypes & /*arguments*/) const override |
114 | { |
115 | return std::make_shared<DataTypeString>(); |
116 | } |
117 | |
118 | bool isDeterministic() const override { return false; } |
119 | |
120 | void executeImpl(Block & block, const ColumnNumbers &, size_t result, size_t input_rows_count) override |
121 | { |
122 | block.getByPosition(result).column = DataTypeString().createColumnConst(input_rows_count, quota_key); |
123 | } |
124 | }; |
125 | |
126 | |
127 | void registerFunctionCurrentQuota(FunctionFactory & factory) |
128 | { |
129 | factory.registerFunction<FunctionCurrentQuota>(); |
130 | factory.registerFunction<FunctionCurrentQuotaId>(); |
131 | factory.registerFunction<FunctionCurrentQuotaKey>(); |
132 | } |
133 | |
134 | } |
135 | |