ClickHouse

1
#include <Columns/ColumnString.h>
2
#include <DataTypes/DataTypeString.h>
3
#include <DataTypes/DataTypeEnum.h>
4
#include <Functions/FunctionFactory.h>
5
#include <Functions/IFunction.h>
6
#include <Interpreters/parseColumnsListForTableFunction.h>
7
#include <Interpreters/Context.h>
8
#include <IO/WriteBufferFromVector.h>
9
#include <Formats/StructureToCapnProtoSchema.h>
10
#include <Formats/StructureToProtobufSchema.h>
11

12
#include <Common/randomSeed.h>
13

14

15
namespace DB
16
{
17

18
namespace ErrorCodes
19
{
20
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
21
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
22
}
23

24
template <class Impl>
25
class FunctionStructureToFormatSchema : public IFunction
26
{
27
public:
28

29
    static constexpr auto name = Impl::name;
30
    explicit FunctionStructureToFormatSchema(ContextPtr context_) : context(std::move(context_))
31
    {
32
    }
33

34
    static FunctionPtr create(ContextPtr ctx)
35
    {
36
        return std::make_shared<FunctionStructureToFormatSchema>(std::move(ctx));
37
    }
38

39
    String getName() const override { return name; }
40

41
    size_t getNumberOfArguments() const override { return 0; }
42
    bool isVariadic() const override { return true; }
43

44
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
45
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const  override { return {0, 1}; }
46
    bool useDefaultImplementationForConstants() const override { return false; }
47
    bool useDefaultImplementationForNulls() const override { return false; }
48

49
    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
50
    {
51
        if (arguments.empty() || arguments.size() > 2)
52
            throw Exception(
53
                ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
54
                "Number of arguments for function {} doesn't match: passed {}, expected 1 or 2",
55
                getName(), arguments.size());
56

57
        if (!isString(arguments[0]))
58
        {
59
            throw Exception(
60
                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
61
                "Illegal type {} of the first argument of function {}, expected constant string",
62
                arguments[0]->getName(),
63
                getName());
64
        }
65

66
        if (arguments.size() > 1 && !isString(arguments[1]))
67
        {
68
            throw Exception(
69
                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
70
                "Illegal type {} of the second argument of function {}, expected constant string",
71
                arguments[1]->getName(),
72
                getName());
73
        }
74

75
        return std::make_shared<DataTypeString>();
76
    }
77

78
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
79
    {
80
        if (arguments.empty() || arguments.size() > 2)
81
            throw Exception(
82
                ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
83
                "Number of arguments for function {} doesn't match: passed {}, expected 1 or 2",
84
                getName(), arguments.size());
85

86
        String structure = arguments[0].column->getDataAt(0).toString();
87
        String message_name = arguments.size() == 2 ? arguments[1].column->getDataAt(0).toString() : "Message";
88
        auto columns_list = parseColumnsListFromString(structure, context);
89
        auto col_res = ColumnString::create();
90
        auto & data = assert_cast<ColumnString &>(*col_res).getChars();
91
        WriteBufferFromVector buf(data);
92
        Impl::writeSchema(buf, message_name, columns_list.getAll());
93
        buf.finalize();
94
        auto & offsets = assert_cast<ColumnString &>(*col_res).getOffsets();
95
        offsets.push_back(data.size());
96
        return ColumnConst::create(std::move(col_res), input_rows_count);
97
    }
98

99
private:
100
    ContextPtr context;
101
};
102

103

104
REGISTER_FUNCTION(StructureToCapnProtoSchema)
105
{
106
    factory.registerFunction<FunctionStructureToFormatSchema<StructureToCapnProtoSchema>>(FunctionDocumentation
107
        {
108
            .description=R"(
109
Function that converts ClickHouse table structure to CapnProto format schema
110
)",
111
            .examples{
112
                {"random", "SELECT structureToCapnProtoSchema('s String, x UInt32', 'MessageName') format TSVRaw", "struct MessageName\n"
113
"{\n"
114
"    s @0 : Data;\n"
115
"    x @1 : UInt32;\n"
116
"}"},
117
            },
118
            .categories{"Other"}
119
        },
120
        FunctionFactory::CaseSensitive);
121
}
122

123

124
REGISTER_FUNCTION(StructureToProtobufSchema)
125
{
126
    factory.registerFunction<FunctionStructureToFormatSchema<StructureToProtobufSchema>>(FunctionDocumentation
127
        {
128
            .description=R"(
129
Function that converts ClickHouse table structure to Protobuf format schema
130
)",
131
            .examples{
132
                {"random", "SELECT structureToCapnProtoSchema('s String, x UInt32', 'MessageName') format TSVRaw", "syntax = \"proto3\";\n"
133
"\n"
134
"message MessageName\n"
135
"{\n"
136
"    bytes s = 1;\n"
137
"    uint32 x = 2;\n"
138
"}"},
139
            },
140
            .categories{"Other"}
141
        },
142
        FunctionFactory::CaseSensitive);
143
}
144

145
}
146