1 | /** |
2 | * Copyright 2021 Alibaba, Inc. and its affiliates. All Rights Reserved. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | |
16 | * \author Haichao.chc |
17 | * \date Oct 2020 |
18 | * \brief Client tool which can connect & operate remote |
19 | * collections of proxima be |
20 | */ |
21 | |
22 | #include <iostream> |
23 | #include <map> |
24 | #include <string> |
25 | #include <ailego/utility/time_helper.h> |
26 | #include <brpc/channel.h> |
27 | #include <gflags/gflags.h> |
28 | #include <google/protobuf/util/json_util.h> |
29 | #include "common/logger.h" |
30 | #include "common/version.h" |
31 | #include "proto/proxima_be.pb.h" |
32 | |
33 | DEFINE_string(command, "" , "Command type: create | drop" ); |
34 | DEFINE_string(host, "" , "The host of proxima be" ); |
35 | DEFINE_string(collection, "" , "Collection name" ); |
36 | DEFINE_string(schema, "" , "Collection schema" ); |
37 | |
38 | static bool ValidateNotEmpty(const char *flagname, const std::string &value) { |
39 | return !value.empty(); |
40 | } |
41 | |
42 | DEFINE_validator(command, ValidateNotEmpty); |
43 | DEFINE_validator(host, ValidateNotEmpty); |
44 | DEFINE_validator(collection, ValidateNotEmpty); |
45 | |
46 | static inline void PrintUsage() { |
47 | std::cout << "Usage:" << std::endl; |
48 | std::cout << " admin_client <args>" << std::endl << std::endl; |
49 | std::cout << "Args: " << std::endl; |
50 | std::cout << " --command Command type: create | drop" << std::endl; |
51 | std::cout << " --host The host of proxima be" << std::endl; |
52 | std::cout << " --collection Specify collection name" << std::endl; |
53 | std::cout << " --schema Specify collection schema format" << std::endl; |
54 | std::cout << " --help, -h Display help info" << std::endl; |
55 | std::cout << " --version, -v Dipslay version info" << std::endl; |
56 | } |
57 | |
58 | static brpc::Channel g_client_channel; |
59 | |
60 | static bool InitClientChannel() { |
61 | brpc::ChannelOptions options; |
62 | options.protocol = "http" ; |
63 | options.timeout_ms = 60000; |
64 | options.max_retry = 2; |
65 | |
66 | if (g_client_channel.Init(FLAGS_host.c_str(), "" , &options) != 0) { |
67 | LOG_ERROR("Init client channel failed." ); |
68 | return false; |
69 | } |
70 | return true; |
71 | } |
72 | |
73 | static void CreateCollection() { |
74 | if (FLAGS_schema.empty()) { |
75 | LOG_ERROR("Input schema can't be empty" ); |
76 | exit(1); |
77 | } |
78 | |
79 | char url[2048] = {0}; |
80 | snprintf(url, sizeof(url), "%s/v1/collection/%s" , FLAGS_host.c_str(), |
81 | FLAGS_collection.c_str()); |
82 | ailego::ElapsedTime timer; |
83 | brpc::Controller cntl; |
84 | cntl.http_request().uri() = url; |
85 | cntl.http_request().set_method(brpc::HTTP_METHOD_POST); |
86 | cntl.request_attachment().append(FLAGS_schema); |
87 | g_client_channel.CallMethod(nullptr, &cntl, nullptr, nullptr, nullptr); |
88 | |
89 | if (!cntl.Failed()) { |
90 | google::protobuf::util::JsonParseOptions options; |
91 | options.ignore_unknown_fields = true; |
92 | proxima::be::proto::Status response; |
93 | std::string resp_body = cntl.response_attachment().to_string(); |
94 | auto status = google::protobuf::util::JsonStringToMessage( |
95 | resp_body, &response, options); |
96 | if (status.ok() && response.code() == 0) { |
97 | LOG_INFO("Create collection success. collection[%s] rt[%zums]" , |
98 | FLAGS_collection.c_str(), (size_t)timer.milli_seconds()); |
99 | } else { |
100 | LOG_ERROR("Create collection error. code[%d] reason[%s]" , response.code(), |
101 | response.reason().c_str()); |
102 | } |
103 | } else { |
104 | LOG_ERROR("Create collection error. error_msg[%s]" , |
105 | cntl.ErrorText().c_str()); |
106 | } |
107 | } |
108 | |
109 | static void DropCollection() { |
110 | char url[2048] = {0}; |
111 | snprintf(url, sizeof(url), "%s/v1/collection/%s" , FLAGS_host.c_str(), |
112 | FLAGS_collection.c_str()); |
113 | ailego::ElapsedTime timer; |
114 | brpc::Controller cntl; |
115 | cntl.http_request().uri() = url; |
116 | cntl.http_request().set_method(brpc::HTTP_METHOD_DELETE); |
117 | g_client_channel.CallMethod(nullptr, &cntl, nullptr, nullptr, nullptr); |
118 | |
119 | if (!cntl.Failed()) { |
120 | google::protobuf::util::JsonParseOptions options; |
121 | options.ignore_unknown_fields = true; |
122 | proxima::be::proto::Status response; |
123 | std::string resp_body = cntl.response_attachment().to_string(); |
124 | auto status = google::protobuf::util::JsonStringToMessage( |
125 | resp_body, &response, options); |
126 | if (status.ok() && response.code() == 0) { |
127 | LOG_INFO("Drop collection success. collection[%s] rt[%zums]" , |
128 | FLAGS_collection.c_str(), (size_t)timer.milli_seconds()); |
129 | } else { |
130 | LOG_ERROR("Drop collection error. code[%d] reason[%s]" , response.code(), |
131 | response.reason().c_str()); |
132 | } |
133 | } else { |
134 | LOG_ERROR("Drop collection error. error_msg[%s]" , cntl.ErrorText().c_str()); |
135 | } |
136 | } |
137 | |
138 | int main(int argc, char **argv) { |
139 | // Parse arguments |
140 | for (int i = 1; i < argc; ++i) { |
141 | const char *arg = argv[i]; |
142 | if (!strcmp(arg, "-help" ) || !strcmp(arg, "--help" ) || !strcmp(arg, "-h" )) { |
143 | PrintUsage(); |
144 | exit(0); |
145 | } else if (!strcmp(arg, "-version" ) || !strcmp(arg, "--version" ) || |
146 | !strcmp(arg, "-v" )) { |
147 | std::cout << proxima::be::Version::Details() << std::endl; |
148 | exit(0); |
149 | } |
150 | } |
151 | gflags::ParseCommandLineNonHelpFlags(&argc, &argv, false); |
152 | |
153 | // Init client channel |
154 | if (!InitClientChannel()) { |
155 | LOG_ERROR("Init client channel failed. host[%s]" , FLAGS_host.c_str()); |
156 | exit(1); |
157 | } |
158 | |
159 | // Register commands |
160 | std::map<std::string, std::function<void(void)>> collection_ops = { |
161 | {"create" , CreateCollection}, {"drop" , DropCollection}}; |
162 | if (collection_ops.find(FLAGS_command) != collection_ops.end()) { |
163 | collection_ops[FLAGS_command](); |
164 | } else { |
165 | LOG_ERROR("Unsupported command type: %s" , FLAGS_command.c_str()); |
166 | exit(1); |
167 | } |
168 | |
169 | return 0; |
170 | } |
171 | |