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
10namespace DB
11{
12
13class FunctionCurrentQuota : public IFunction
14{
15 const String quota_name;
16
17public:
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
51class FunctionCurrentQuotaId : public IFunction
52{
53 const UUID quota_id;
54
55public:
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
89class FunctionCurrentQuotaKey : public IFunction
90{
91 const String quota_key;
92
93public:
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
127void registerFunctionCurrentQuota(FunctionFactory & factory)
128{
129 factory.registerFunction<FunctionCurrentQuota>();
130 factory.registerFunction<FunctionCurrentQuotaId>();
131 factory.registerFunction<FunctionCurrentQuotaKey>();
132}
133
134}
135