Skip to content

Commit dc4ea7f

Browse files
snaurysnaury
authored
Use the new tablet service in ydbd admin tablet commands (#8478)
1 parent e6f96d4 commit dc4ea7f

File tree

4 files changed

+126
-64
lines changed

4 files changed

+126
-64
lines changed

‎ydb/core/driver_lib/cli_base/cli_kicli.cpp

Lines changed: 24 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -64,13 +64,8 @@ int OnMessageBus(const TClientCommand::TConfig& config, const NMsgBusProxy::TBus
6464
return 0;
6565
}
6666

67-
int InvokeThroughKikimr(TClientCommand::TConfig& config, std::function<int(NClient::TKikimr&)> handler) {
68-
NClient::TKikimr kikimr(CommandConfig.ClientConfig);
69-
if (!config.SecurityToken.empty()) {
70-
kikimr.SetSecurityToken(config.SecurityToken);
71-
}
72-
73-
if (!config.StaticCredentials.User.empty()) {
67+
std::optional<TString> AcquireSecurityToken(TClientCommand::TConfig& config) {
68+
if (config.SecurityToken.empty() && !config.StaticCredentials.User.empty()) {
7469
NYdb::TDriverConfig driverConfig;
7570
driverConfig.SetEndpoint(TCommandConfig::ParseServerAddress(config.Address).Address);
7671
NYdb::TDriver connection(driverConfig);
@@ -79,14 +74,32 @@ int InvokeThroughKikimr(TClientCommand::TConfig& config, std::function<int(NClie
7974
auto credentialsProviderFactory = NYdb::CreateLoginCredentialsProviderFactory(config.StaticCredentials);
8075
auto loginProvider = credentialsProviderFactory->CreateProvider(client.GetCoreFacility());
8176
try {
77+
// Note: may throw exceptions
8278
config.SecurityToken = loginProvider->GetAuthInfo();
83-
} catch (yexception& ex) {
84-
Cerr << ex.what() << Endl;
79+
} catch (...) {
8580
connection.Stop();
86-
return 1;
81+
throw;
8782
}
8883
connection.Stop();
89-
kikimr.SetSecurityToken(config.SecurityToken);
84+
}
85+
86+
if (!config.SecurityToken.empty()) {
87+
return config.SecurityToken;
88+
} else {
89+
return std::nullopt;
90+
}
91+
}
92+
93+
int InvokeThroughKikimr(TClientCommand::TConfig& config, std::function<int(NClient::TKikimr&)> handler) {
94+
NClient::TKikimr kikimr(CommandConfig.ClientConfig);
95+
96+
try {
97+
if (auto token = AcquireSecurityToken(config)) {
98+
kikimr.SetSecurityToken(*token);
99+
}
100+
} catch (const std::exception& e) {
101+
Cerr << e.what() << Endl;
102+
return 1;
90103
}
91104

92105
return handler(kikimr);

‎ydb/core/driver_lib/cli_base/cli_kicli.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ int HandleResponse(const NThreading::TFuture<ResultType>& future, std::function<
4343
return callback(result);
4444
}
4545

46+
std::optional<TString> AcquireSecurityToken(TClientCommand::TConfig& config);
47+
4648
int InvokeThroughKikimr(TClientCommand::TConfig& config, std::function<int(NClient::TKikimr&)> handler);
4749

4850
template <typename RequestType>

‎ydb/core/driver_lib/cli_utils/cli_cmds_tablet.cpp

Lines changed: 95 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
#include "cli_cmds.h"
33

44
#include <ydb/core/protos/base.pb.h>
5+
#include <ydb/library/grpc/client/grpc_common.h>
6+
#include <ydb/public/api/grpc/draft/ydb_tablet_v1.grpc.pb.h>
57

68
namespace NKikimr {
79
namespace NDriverClient {
@@ -83,123 +85,163 @@ class TClientCommandKeyValue : public TClientCommandTree {
8385
}
8486
};
8587

86-
class TClientCommandTabletExec : public TClientCommandBase {
88+
template<class TRequest, class TResponse>
89+
class TClientCommandTabletCommon {
90+
protected:
91+
int RunTabletCommon(TClientCommand::TConfig& config) {
92+
grpc::ClientContext context;
93+
if (auto token = AcquireSecurityToken(config)) {
94+
context.AddMetadata("x-ydb-auth-ticket", *token);
95+
}
96+
auto channel = NYdbGrpc::CreateChannelInterface(CommandConfig.ClientConfig);
97+
auto stub = Ydb::Tablet::V1::TabletService::NewStub(channel);
98+
if (NClient::TKikimr::DUMP_REQUESTS) {
99+
Cerr << "<-- " << TypeName<TRequest>() << "\n" << Request.DebugString();
100+
}
101+
auto status = Send(context, stub);
102+
if (!status.ok()) {
103+
Cerr << "ERROR: " << int(status.error_code()) << " " << status.error_message() << Endl;
104+
return 1;
105+
}
106+
if (NClient::TKikimr::DUMP_REQUESTS) {
107+
Cerr << "--> " << TypeName<TResponse>() << "\n" << Response.DebugString();
108+
}
109+
if (Response.status() != Ydb::StatusIds::SUCCESS) {
110+
Cerr << "ERROR: " << Response.status() << Endl;
111+
for (const auto& issue : Response.issues()) {
112+
Cerr << issue.message() << Endl;
113+
}
114+
return 1;
115+
}
116+
return 0;
117+
}
118+
119+
virtual grpc::Status Send(
120+
grpc::ClientContext& context,
121+
const std::unique_ptr<Ydb::Tablet::V1::TabletService::Stub>& stub) = 0;
122+
123+
protected:
124+
TRequest Request;
125+
TResponse Response;
126+
};
127+
128+
class TClientCommandTabletExec
129+
: public TClientCommandBase
130+
, public TClientCommandTabletCommon<
131+
Ydb::Tablet::ExecuteTabletMiniKQLRequest,
132+
Ydb::Tablet::ExecuteTabletMiniKQLResponse>
133+
{
87134
public:
88135
TClientCommandTabletExec()
89136
: TClientCommandBase("execute", { "exec" })
90137
{
91138
}
92139

93-
TAutoPtr<NMsgBusProxy::TBusTabletLocalMKQL> Request;
94140
TString Program;
95-
TString Params;
96141

97142
virtual void Config(TConfig& config) override {
98143
TClientCommand::Config(config);
99-
config.Opts->AddLongOption("follower", "connect to follower").NoArgument();
100-
config.Opts->AddLongOption("json-ui64-as-string", "json output ui64 as string").NoArgument();
101-
config.Opts->AddLongOption("json-binary-as-base64", "json output binary data in base64").NoArgument();
102-
config.SetFreeArgsNum(1, 2);
144+
config.Opts->AddLongOption("dry-run", "test changes without applying").NoArgument();
145+
config.SetFreeArgsNum(1, 1);
103146
SetFreeArgTitle(0, "<PROGRAM>", "Program to execute");
104-
SetFreeArgTitle(1, "<PARAMS>", "Parameters of the program");
105147
}
106148

107149
virtual void Parse(TConfig& config) override {
108150
TClientCommand::Parse(config);
109151

110152
Program = GetMiniKQL(config.ParseResult->GetFreeArgs().at(0));
111-
if (config.ParseResult->GetFreeArgCount() > 1)
112-
Params = GetMiniKQL(config.ParseResult->GetFreeArgs().at(1));
113-
114-
Request = new NMsgBusProxy::TBusTabletLocalMKQL;
115-
Request->Record.SetTabletID(config.TabletId);
116-
auto* pgm = Request->Record.MutableProgram();
117-
if (IsMiniKQL(Program)) {
118-
pgm->MutableProgram()->SetText(Program);
119-
} else {
120-
pgm->MutableProgram()->SetBin(Program);
121-
}
122153

123-
if (!Params.empty()) {
124-
if (IsMiniKQL(Params)) {
125-
pgm->MutableParams()->SetText(Params);
126-
} else {
127-
pgm->MutableParams()->SetBin(Params);
128-
}
129-
}
130-
131-
Request->Record.SetConnectToFollower(config.ParseResult->Has("follower"));
132-
config.JsonUi64AsText = config.ParseResult->Has("json-ui64-as-string");
133-
config.JsonBinaryAsBase64 = config.ParseResult->Has("json-binary-as-base64");
154+
Request.set_tablet_id(config.TabletId);
155+
Request.set_program(Program);
156+
Request.set_dry_run(config.ParseResult->Has("dry-run"));
134157
}
135158

136159
virtual int Run(TConfig& config) override {
137-
return MessageBusCall(config, Request);
160+
return RunTabletCommon(config);
161+
}
162+
163+
virtual grpc::Status Send(
164+
grpc::ClientContext& context,
165+
const std::unique_ptr<Ydb::Tablet::V1::TabletService::Stub>& stub) override
166+
{
167+
return stub->ExecuteTabletMiniKQL(&context, Request, &Response);
138168
}
139169
};
140170

141-
class TClientCommandTabletKill : public TClientCommand {
171+
class TClientCommandTabletKill
172+
: public TClientCommand
173+
, public TClientCommandTabletCommon<
174+
Ydb::Tablet::RestartTabletRequest,
175+
Ydb::Tablet::RestartTabletResponse>
176+
{
142177
public:
143178
TClientCommandTabletKill()
144179
: TClientCommand("kill")
145180
{
146181
}
147182

148-
TAutoPtr<NMsgBusProxy::TBusTabletKillRequest> Request;
149-
150183
virtual void Config(TConfig& config) override {
151184
TClientCommand::Config(config);
152185
config.SetFreeArgsNum(0);
153186
}
154187

155188
virtual void Parse(TConfig& config) override {
156189
TClientCommand::Parse(config);
157-
Request = new NMsgBusProxy::TBusTabletKillRequest;
158-
Request->Record.SetTabletID(config.TabletId);
190+
Request.set_tablet_id(config.TabletId);
159191
}
160192

161193
virtual int Run(TConfig& config) override {
162-
return MessageBusCall(config, Request);
194+
return RunTabletCommon(config);
195+
}
196+
197+
virtual grpc::Status Send(
198+
grpc::ClientContext& context,
199+
const std::unique_ptr<Ydb::Tablet::V1::TabletService::Stub>& stub) override
200+
{
201+
return stub->RestartTablet(&context, Request, &Response);
163202
}
164203
};
165204

166-
class TClientCommandTabletSchemeTx : public TClientCommand {
205+
class TClientCommandTabletSchemeTx
206+
: public TClientCommand
207+
, public TClientCommandTabletCommon<
208+
Ydb::Tablet::ChangeTabletSchemaRequest,
209+
Ydb::Tablet::ChangeTabletSchemaResponse>
210+
{
167211
public:
168212
TClientCommandTabletSchemeTx()
169213
: TClientCommand("scheme-tx", { "scheme" })
170214
{
171215
}
172216

173-
TAutoPtr<NMsgBusProxy::TBusTabletLocalSchemeTx> Request;
174217
TString SchemeChanges;
175218

176219
virtual void Config(TConfig& config) override {
177220
TClientCommand::Config(config);
178-
config.Opts->AddLongOption("follower", "connect to follower");
179-
config.Opts->AddLongOption("dry-run", "test changes without applying");
221+
config.Opts->AddLongOption("dry-run", "test changes without applying").NoArgument();
180222
config.SetFreeArgsNum(1, 1);
181-
SetFreeArgTitle(0, "<SCHEME CHANGES>", "Scheme changes to apply");
223+
SetFreeArgTitle(0, "<SCHEME CHANGES>", "Scheme changes json to apply");
182224
}
183225

184226
virtual void Parse(TConfig& config) override {
185227
TClientCommand::Parse(config);
186228

187229
SchemeChanges = config.ParseResult->GetFreeArgs().at(0);
188230

189-
Request = new NMsgBusProxy::TBusTabletLocalSchemeTx;
190-
Request->Record.SetTabletID(config.TabletId);
191-
auto* schemeChanges = Request->Record.MutableSchemeChanges();
192-
if (!google::protobuf::TextFormat::ParseFromString(SchemeChanges, schemeChanges)) {
193-
ythrow TWithBackTrace<yexception>() << "Invalid scheme changes protobuf passed";
194-
}
195-
196-
if (config.ParseResult->Has("follower"))
197-
Request->Record.SetConnectToFollower(true);
198-
Request->Record.SetDryRun(config.ParseResult->Has("dry-run"));
231+
Request.set_tablet_id(config.TabletId);
232+
Request.set_schema_changes(SchemeChanges);
233+
Request.set_dry_run(config.ParseResult->Has("dry-run"));
199234
}
200235

201236
virtual int Run(TConfig& config) override {
202-
return MessageBusCall(config, Request);
237+
return RunTabletCommon(config);
238+
}
239+
240+
virtual grpc::Status Send(
241+
grpc::ClientContext& context,
242+
const std::unique_ptr<Ydb::Tablet::V1::TabletService::Stub>& stub) override
243+
{
244+
return stub->ChangeTabletSchema(&context, Request, &Response);
203245
}
204246
};
205247

‎ydb/core/driver_lib/run/run.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -659,6 +659,11 @@ void TKikimrRunner::InitializeGRpc(const TKikimrRunConfig& runConfig) {
659659
hasQueryService = true;
660660
}
661661

662+
if (hasLegacy) {
663+
// Enable new public services when the legacy service is enabled
664+
hasTabletService = true;
665+
}
666+
662667
// Enable RL for all services if enabled list is empty
663668
if (rlServicesEnabled.empty()) {
664669
for (auto& [name, cfg] : names) {

0 commit comments

Comments
 (0)