ClickHouse

1
#include <Functions/IFunction.h>
2
#include <Functions/FunctionFactory.h>
3
#include <Interpreters/Context.h>
4
#include <Access/AccessControl.h>
5
#include <Access/User.h>
6
#include <Columns/ColumnArray.h>
7
#include <Columns/ColumnConst.h>
8
#include <Columns/ColumnString.h>
9
#include <DataTypes/DataTypeString.h>
10
#include <DataTypes/DataTypeArray.h>
11

12

13
namespace DB
14
{
15

16
namespace
17
{
18
    enum class Kind
19
    {
20
        currentProfiles,
21
        enabledProfiles,
22
        defaultProfiles,
23
    };
24

25
    String toString(Kind kind)
26
    {
27
        switch (kind)
28
        {
29
            case Kind::currentProfiles: return "currentProfiles";
30
            case Kind::enabledProfiles: return "enabledProfiles";
31
            case Kind::defaultProfiles: return "defaultProfiles";
32
        }
33
    }
34

35
    class FunctionProfiles : public IFunction
36
    {
37
    public:
38
        bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override
39
        {
40
            return false;
41
        }
42

43
        String getName() const override
44
        {
45
            return toString(kind);
46
        }
47

48
        explicit FunctionProfiles(const ContextPtr & context_, Kind kind_)
49
            : kind(kind_)
50
            , context(context_)
51
        {}
52

53
        size_t getNumberOfArguments() const override { return 0; }
54
        bool isDeterministic() const override { return false; }
55

56
        DataTypePtr getReturnTypeImpl(const DataTypes & /*arguments*/) const override
57
        {
58
            return std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>());
59
        }
60

61
        ColumnPtr executeImpl(const ColumnsWithTypeAndName &, const DataTypePtr &, size_t input_rows_count) const override
62
        {
63
            std::call_once(initialized_flag, [&]{ initialize(); });
64

65
            auto col_res = ColumnArray::create(ColumnString::create());
66
            ColumnString & res_strings = typeid_cast<ColumnString &>(col_res->getData());
67
            ColumnArray::Offsets & res_offsets = col_res->getOffsets();
68
            for (const String & profile_name : profile_names)
69
                res_strings.insertData(profile_name.data(), profile_name.length());
70
            res_offsets.push_back(res_strings.size());
71
            return ColumnConst::create(std::move(col_res), input_rows_count);
72
        }
73

74
    private:
75
        void initialize() const
76
        {
77
            const auto & manager = context->getAccessControl();
78

79
            std::vector<UUID> profile_ids;
80

81
            switch (kind)
82
            {
83
                case Kind::currentProfiles: profile_ids = context->getCurrentProfiles(); break;
84
                case Kind::enabledProfiles: profile_ids = context->getEnabledProfiles(); break;
85
                case Kind::defaultProfiles: profile_ids = context->getUser()->settings.toProfileIDs(); break;
86
            }
87

88
            profile_names = manager.tryReadNames(profile_ids);
89
        }
90

91
        mutable std::once_flag initialized_flag;
92

93
        Kind kind;
94
        ContextPtr context;
95
        mutable Strings profile_names;
96
    };
97
}
98

99
REGISTER_FUNCTION(Profiles)
100
{
101
    factory.registerFunction("currentProfiles", [](ContextPtr context){ return std::make_shared<FunctionProfiles>(context, Kind::currentProfiles); });
102
    factory.registerFunction("enabledProfiles", [](ContextPtr context){ return std::make_shared<FunctionProfiles>(context, Kind::enabledProfiles); });
103
    factory.registerFunction("defaultProfiles", [](ContextPtr context){ return std::make_shared<FunctionProfiles>(context, Kind::defaultProfiles); });
104
}
105

106
}
107