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
33DEFINE_string(command, "", "Command type: create | drop");
34DEFINE_string(host, "", "The host of proxima be");
35DEFINE_string(collection, "", "Collection name");
36DEFINE_string(schema, "", "Collection schema");
37
38static bool ValidateNotEmpty(const char *flagname, const std::string &value) {
39 return !value.empty();
40}
41
42DEFINE_validator(command, ValidateNotEmpty);
43DEFINE_validator(host, ValidateNotEmpty);
44DEFINE_validator(collection, ValidateNotEmpty);
45
46static 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
58static brpc::Channel g_client_channel;
59
60static 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
73static 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
109static 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
138int 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