1 | #include <Functions/IFunctionImpl.h> |
2 | #include <Functions/FunctionFactory.h> |
3 | #include <Functions/FunctionHelpers.h> |
4 | #include <Columns/ColumnString.h> |
5 | #include <Columns/ColumnConst.h> |
6 | #include <DataTypes/DataTypesNumber.h> |
7 | #include <Storages/IStorage.h> |
8 | #include <Interpreters/Cluster.h> |
9 | #include <Interpreters/Context.h> |
10 | #include <Storages/getStructureOfRemoteTable.h> |
11 | |
12 | |
13 | namespace DB |
14 | { |
15 | |
16 | namespace ErrorCodes |
17 | { |
18 | extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; |
19 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
20 | } |
21 | |
22 | |
23 | /** Usage: |
24 | * hasColumnInTable(['hostname'[, 'username'[, 'password']],] 'database', 'table', 'column') |
25 | */ |
26 | class FunctionHasColumnInTable : public IFunction |
27 | { |
28 | public: |
29 | static constexpr auto name = "hasColumnInTable" ; |
30 | static FunctionPtr create(const Context & context) |
31 | { |
32 | return std::make_shared<FunctionHasColumnInTable>(context.getGlobalContext()); |
33 | } |
34 | |
35 | explicit FunctionHasColumnInTable(const Context & global_context_) : global_context(global_context_) |
36 | { |
37 | } |
38 | |
39 | bool isVariadic() const override |
40 | { |
41 | return true; |
42 | } |
43 | size_t getNumberOfArguments() const override |
44 | { |
45 | return 0; |
46 | } |
47 | |
48 | String getName() const override |
49 | { |
50 | return name; |
51 | } |
52 | |
53 | DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override; |
54 | |
55 | bool isDeterministic() const override { return false; } |
56 | |
57 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override; |
58 | |
59 | private: |
60 | const Context & global_context; |
61 | }; |
62 | |
63 | |
64 | DataTypePtr FunctionHasColumnInTable::getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const |
65 | { |
66 | if (arguments.size() < 3 || arguments.size() > 6) |
67 | throw Exception{"Invalid number of arguments for function " + getName(), |
68 | ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH}; |
69 | |
70 | static const std::string arg_pos_description[] = {"First" , "Second" , "Third" , "Fourth" , "Fifth" , "Sixth" }; |
71 | for (size_t i = 0; i < arguments.size(); ++i) |
72 | { |
73 | const ColumnWithTypeAndName & argument = arguments[i]; |
74 | |
75 | if (!checkColumnConst<ColumnString>(argument.column.get())) |
76 | { |
77 | throw Exception(arg_pos_description[i] + " argument for function " + getName() + " must be const String." , |
78 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
79 | } |
80 | } |
81 | |
82 | return std::make_shared<DataTypeUInt8>(); |
83 | } |
84 | |
85 | |
86 | void FunctionHasColumnInTable::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) |
87 | { |
88 | auto get_string_from_block = [&](size_t column_pos) -> String |
89 | { |
90 | ColumnPtr column = block.getByPosition(column_pos).column; |
91 | const ColumnConst * const_column = checkAndGetColumnConst<ColumnString>(column.get()); |
92 | return const_column->getValue<String>(); |
93 | }; |
94 | |
95 | size_t arg = 0; |
96 | String host_name; |
97 | String user_name; |
98 | String password; |
99 | |
100 | if (arguments.size() > 3) |
101 | host_name = get_string_from_block(arguments[arg++]); |
102 | |
103 | if (arguments.size() > 4) |
104 | user_name = get_string_from_block(arguments[arg++]); |
105 | |
106 | if (arguments.size() > 5) |
107 | password = get_string_from_block(arguments[arg++]); |
108 | |
109 | String database_name = get_string_from_block(arguments[arg++]); |
110 | String table_name = get_string_from_block(arguments[arg++]); |
111 | String column_name = get_string_from_block(arguments[arg++]); |
112 | |
113 | bool has_column; |
114 | if (host_name.empty()) |
115 | { |
116 | const StoragePtr & table = global_context.getTable(database_name, table_name); |
117 | has_column = table->hasColumn(column_name); |
118 | } |
119 | else |
120 | { |
121 | std::vector<std::vector<String>> host_names = {{ host_name }}; |
122 | |
123 | auto cluster = std::make_shared<Cluster>( |
124 | global_context.getSettings(), |
125 | host_names, |
126 | !user_name.empty() ? user_name : "default" , |
127 | password, |
128 | global_context.getTCPPort(), |
129 | false); |
130 | |
131 | auto remote_columns = getStructureOfRemoteTable(*cluster, database_name, table_name, global_context); |
132 | has_column = remote_columns.hasPhysical(column_name); |
133 | } |
134 | |
135 | block.getByPosition(result).column = DataTypeUInt8().createColumnConst(input_rows_count, Field(has_column)); |
136 | } |
137 | |
138 | |
139 | void registerFunctionHasColumnInTable(FunctionFactory & factory) |
140 | { |
141 | factory.registerFunction<FunctionHasColumnInTable>(); |
142 | } |
143 | |
144 | } |
145 | |