ClickHouse

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

14

15
namespace DB
16
{
17

18
namespace
19
{
20
    enum class Kind
21
    {
22
        CURRENT_ROLES,
23
        ENABLED_ROLES,
24
        DEFAULT_ROLES,
25
    };
26

27
    template <Kind kind>
28
    class FunctionCurrentRoles : public IFunction
29
    {
30
    public:
31
        static constexpr auto name = (kind == Kind::CURRENT_ROLES) ? "currentRoles" : ((kind == Kind::ENABLED_ROLES) ? "enabledRoles" : "defaultRoles");
32
        static FunctionPtr create(const ContextPtr & context) { return std::make_shared<FunctionCurrentRoles>(context); }
33

34
        bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
35

36
        String getName() const override { return name; }
37

38
        explicit FunctionCurrentRoles(const ContextPtr & context_)
39
            : context(context_)
40
        {}
41

42
        size_t getNumberOfArguments() const override { return 0; }
43
        bool isDeterministic() const override { return false; }
44

45
        DataTypePtr getReturnTypeImpl(const DataTypes & /*arguments*/) const override
46
        {
47
            return std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>());
48
        }
49

50
        ColumnPtr executeImpl(const ColumnsWithTypeAndName &, const DataTypePtr &, size_t input_rows_count) const override
51
        {
52
            std::call_once(initialized_flag, [&]{ initialize(); });
53

54
            auto col_res = ColumnArray::create(ColumnString::create());
55
            ColumnString & res_strings = typeid_cast<ColumnString &>(col_res->getData());
56
            ColumnArray::Offsets & res_offsets = col_res->getOffsets();
57
            for (const String & role_name : role_names)
58
                res_strings.insertData(role_name.data(), role_name.length());
59
            res_offsets.push_back(res_strings.size());
60
            return ColumnConst::create(std::move(col_res), input_rows_count);
61
        }
62

63
    private:
64
        void initialize() const
65
        {
66
            if constexpr (kind == Kind::CURRENT_ROLES)
67
            {
68
                role_names = context->getRolesInfo()->getCurrentRolesNames();
69
            }
70
            else if constexpr (kind == Kind::ENABLED_ROLES)
71
            {
72
                role_names = context->getRolesInfo()->getEnabledRolesNames();
73
            }
74
            else
75
            {
76
                static_assert(kind == Kind::DEFAULT_ROLES);
77
                const auto & manager = context->getAccessControl();
78
                auto user = context->getUser();
79
                role_names = manager.tryReadNames(user->granted_roles.findGranted(user->default_roles));
80
            }
81

82
            /// We sort the names because the result of the function should not depend on the order of UUIDs.
83
            ::sort(role_names.begin(), role_names.end());
84
        }
85

86
        mutable std::once_flag initialized_flag;
87
        ContextPtr context;
88
        mutable Strings role_names;
89
    };
90
}
91

92
REGISTER_FUNCTION(CurrentRoles)
93
{
94
    factory.registerFunction<FunctionCurrentRoles<Kind::CURRENT_ROLES>>();
95
    factory.registerFunction<FunctionCurrentRoles<Kind::ENABLED_ROLES>>();
96
    factory.registerFunction<FunctionCurrentRoles<Kind::DEFAULT_ROLES>>();
97
}
98

99
}
100