Skip to content

Commit 2b267c4

Browse files
pixccpixcc
committed
Add cluster & database restore to YDB CLI (#14481)
1 parent fa0a65b commit 2b267c4

File tree

6 files changed

+203
-74
lines changed

6 files changed

+203
-74
lines changed

‎include/ydb-cpp-sdk/client/cms/cms.h

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -133,14 +133,14 @@ struct TTargetTrackingPolicy {
133133
TTargetTrackingPolicy() = default;
134134
TTargetTrackingPolicy(const Ydb::Cms::ScaleRecommenderPolicies_ScaleRecommenderPolicy_TargetTrackingPolicy& proto);
135135

136-
std::variant<TAverageCpuUtilizationPercent> Target;
136+
std::variant<std::monostate, TAverageCpuUtilizationPercent> Target;
137137
};
138138

139139
struct TScaleRecommenderPolicy {
140140
TScaleRecommenderPolicy() = default;
141141
TScaleRecommenderPolicy(const Ydb::Cms::ScaleRecommenderPolicies_ScaleRecommenderPolicy& proto);
142142

143-
std::variant<TTargetTrackingPolicy> Policy;
143+
std::variant<std::monostate, TTargetTrackingPolicy> Policy;
144144
};
145145

146146
struct TScaleRecommenderPolicies {
@@ -150,13 +150,15 @@ struct TScaleRecommenderPolicies {
150150
std::vector<TScaleRecommenderPolicy> Policies;
151151
};
152152

153+
using TResourcesKind = std::variant<std::monostate, TResources, TSharedResources, TServerlessResources>;
154+
153155
class TGetDatabaseStatusResult : public TStatus {
154156
public:
155157
TGetDatabaseStatusResult(TStatus&& status, const Ydb::Cms::GetDatabaseStatusResult& proto);
156158

157159
const std::string& GetPath() const;
158160
EState GetState() const;
159-
const std::variant<TResources, TSharedResources, TServerlessResources>& GetResourcesKind() const;
161+
const TResourcesKind& GetResourcesKind() const;
160162
const TResources& GetAllocatedResources() const;
161163
const std::vector<TAllocatedComputationalUnit>& GetRegisteredResources() const;
162164
std::uint64_t GetGeneration() const;
@@ -170,7 +172,7 @@ class TGetDatabaseStatusResult : public TStatus {
170172
private:
171173
std::string Path_;
172174
EState State_;
173-
std::variant<TResources, TSharedResources, TServerlessResources> ResourcesKind_;
175+
TResourcesKind ResourcesKind_;
174176
TResources AllocatedResources_;
175177
std::vector<TAllocatedComputationalUnit> RegisteredResources_;
176178
std::uint64_t Generation_;
@@ -181,14 +183,28 @@ class TGetDatabaseStatusResult : public TStatus {
181183

182184
using TAsyncGetDatabaseStatusResult = NThreading::TFuture<TGetDatabaseStatusResult>;
183185

186+
struct TCreateDatabaseSettings : public TOperationRequestSettings<TCreateDatabaseSettings> {
187+
TCreateDatabaseSettings() = default;
188+
explicit TCreateDatabaseSettings(const Ydb::Cms::CreateDatabaseRequest& request);
189+
190+
// Fills CreateDatabaseRequest proto from this settings
191+
void SerializeTo(Ydb::Cms::CreateDatabaseRequest& request) const;
192+
193+
FLUENT_SETTING(TResourcesKind, ResourcesKind);
194+
FLUENT_SETTING(TSchemaOperationQuotas, SchemaOperationQuotas);
195+
FLUENT_SETTING(TDatabaseQuotas, DatabaseQuotas);
196+
FLUENT_SETTING(TScaleRecommenderPolicies, ScaleRecommenderPolicies);
197+
};
198+
184199
class TCmsClient {
185200
public:
186201
explicit TCmsClient(const TDriver& driver, const TCommonClientSettings& settings = TCommonClientSettings());
187202

188203
TAsyncListDatabasesResult ListDatabases(const TListDatabasesSettings& settings = TListDatabasesSettings());
189204
TAsyncGetDatabaseStatusResult GetDatabaseStatus(const std::string& path,
190205
const TGetDatabaseStatusSettings& settings = TGetDatabaseStatusSettings());
191-
206+
TAsyncStatus CreateDatabase(const std::string& path,
207+
const TCreateDatabaseSettings& settings = TCreateDatabaseSettings());
192208
private:
193209
class TImpl;
194210
std::shared_ptr<TImpl> Impl_;

‎include/ydb-cpp-sdk/client/driver/driver.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ class TDriver {
142142
template<typename TExtension>
143143
void AddExtension(typename TExtension::TParams params = typename TExtension::TParams());
144144

145+
TDriverConfig GetConfig() const;
145146
private:
146147
std::shared_ptr<TGRpcConnectionsImpl> Impl_;
147148
};

‎src/client/cms/cms.cpp

Lines changed: 129 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
#include <ydb-cpp-sdk/client/cms/cms.h>
22

3+
#include <src/client/common_client/impl/client.h>
34
#include <src/api/grpc/ydb_cms_v1.grpc.pb.h>
45
#include <src/api/protos/ydb_cms.pb.h>
5-
#include <src/client/common_client/impl/client.h>
6+
7+
#define INCLUDE_YDB_INTERNAL_H
8+
#include <src/client/impl/ydb_internal/make_request/make.h>
9+
#undef INCLUDE_YDB_INTERNAL_H
610

711
namespace NYdb::inline V3::NCms {
812

@@ -25,6 +29,82 @@ namespace {
2529
return EState::StateUnspecified;
2630
}
2731
}
32+
33+
void SerializeToImpl(
34+
const TResourcesKind& resourcesKind,
35+
const TSchemaOperationQuotas& schemaQuotas,
36+
const TDatabaseQuotas& dbQuotas,
37+
const TScaleRecommenderPolicies& scaleRecommenderPolicies,
38+
Ydb::Cms::CreateDatabaseRequest& out)
39+
{
40+
if (std::holds_alternative<NCms::TResources>(resourcesKind)) {
41+
const auto& resources = std::get<NCms::TResources>(resourcesKind);
42+
for (const auto& storageUnit : resources.StorageUnits) {
43+
auto* protoUnit = out.mutable_resources()->add_storage_units();
44+
protoUnit->set_unit_kind(storageUnit.UnitKind);
45+
protoUnit->set_count(storageUnit.Count);
46+
}
47+
for (const auto& computationalUnit : resources.ComputationalUnits) {
48+
auto* protoUnit = out.mutable_resources()->add_computational_units();
49+
protoUnit->set_unit_kind(computationalUnit.UnitKind);
50+
protoUnit->set_count(computationalUnit.Count);
51+
protoUnit->set_availability_zone(computationalUnit.AvailabilityZone);
52+
}
53+
} else if (std::holds_alternative<NCms::TSharedResources>(resourcesKind)) {
54+
const auto& resources = std::get<NCms::TSharedResources>(resourcesKind);
55+
for (const auto& storageUnit : resources.StorageUnits) {
56+
auto* protoUnit = out.mutable_shared_resources()->add_storage_units();
57+
protoUnit->set_unit_kind(storageUnit.UnitKind);
58+
protoUnit->set_count(storageUnit.Count);
59+
}
60+
for (const auto& computationalUnit : resources.ComputationalUnits) {
61+
auto* protoUnit = out.mutable_shared_resources()->add_computational_units();
62+
protoUnit->set_unit_kind(computationalUnit.UnitKind);
63+
protoUnit->set_count(computationalUnit.Count);
64+
protoUnit->set_availability_zone(computationalUnit.AvailabilityZone);
65+
}
66+
} else if (std::holds_alternative<NCms::TServerlessResources>(resourcesKind)) {
67+
const auto& resources = std::get<NCms::TServerlessResources>(resourcesKind);
68+
out.mutable_serverless_resources()->set_shared_database_path(resources.SharedDatabasePath);
69+
} else if (std::holds_alternative<std::monostate>(resourcesKind)) {
70+
out.clear_resources_kind();
71+
}
72+
73+
for (const auto& quota : schemaQuotas.LeakyBucketQuotas) {
74+
auto protoQuota = out.mutable_schema_operation_quotas()->add_leaky_bucket_quotas();
75+
protoQuota->set_bucket_seconds(quota.BucketSeconds);
76+
protoQuota->set_bucket_size(quota.BucketSize);
77+
}
78+
79+
out.mutable_database_quotas()->set_data_size_hard_quota(dbQuotas.DataSizeHardQuota);
80+
out.mutable_database_quotas()->set_data_size_soft_quota(dbQuotas.DataSizeSoftQuota);
81+
out.mutable_database_quotas()->set_data_stream_shards_quota(dbQuotas.DataStreamShardsQuota);
82+
out.mutable_database_quotas()->set_data_stream_reserved_storage_quota(dbQuotas.DataStreamReservedStorageQuota);
83+
out.mutable_database_quotas()->set_ttl_min_run_internal_seconds(dbQuotas.TtlMinRunInternalSeconds);
84+
85+
for (const auto& quota : dbQuotas.StorageQuotas) {
86+
auto protoQuota = out.mutable_database_quotas()->add_storage_quotas();
87+
protoQuota->set_unit_kind(quota.UnitKind);
88+
protoQuota->set_data_size_hard_quota(quota.DataSizeHardQuota);
89+
protoQuota->set_data_size_soft_quota(quota.DataSizeSoftQuota);
90+
}
91+
92+
for (const auto& policy : scaleRecommenderPolicies.Policies) {
93+
auto* protoPolicy = out.mutable_scale_recommender_policies()->add_policies();
94+
if (std::holds_alternative<NCms::TTargetTrackingPolicy>(policy.Policy)) {
95+
const auto& targetTracking = std::get<NCms::TTargetTrackingPolicy>(policy.Policy);
96+
auto* protoTargetTracking = protoPolicy->mutable_target_tracking_policy();
97+
if (std::holds_alternative<NCms::TTargetTrackingPolicy::TAverageCpuUtilizationPercent>(targetTracking.Target)) {
98+
const auto& target = std::get<NCms::TTargetTrackingPolicy::TAverageCpuUtilizationPercent>(targetTracking.Target);
99+
protoTargetTracking->set_average_cpu_utilization_percent(target);
100+
} else if (std::holds_alternative<std::monostate>(targetTracking.Target)) {
101+
protoTargetTracking->clear_target();
102+
}
103+
} else if (std::holds_alternative<std::monostate>(policy.Policy)) {
104+
protoPolicy->clear_policy();
105+
}
106+
}
107+
}
28108
} // anonymous namespace
29109

30110
TListDatabasesResult::TListDatabasesResult(TStatus&& status, const Ydb::Cms::ListDatabasesResult& proto)
@@ -93,6 +173,7 @@ TTargetTrackingPolicy::TTargetTrackingPolicy(const Ydb::Cms::ScaleRecommenderPol
93173
Target = proto.average_cpu_utilization_percent();
94174
break;
95175
case Ydb::Cms::ScaleRecommenderPolicies_ScaleRecommenderPolicy_TargetTrackingPolicy::TARGET_NOT_SET:
176+
Target = std::monostate();
96177
break;
97178
}
98179
}
@@ -104,6 +185,7 @@ TScaleRecommenderPolicy::TScaleRecommenderPolicy(const Ydb::Cms::ScaleRecommende
104185
Policy = proto.target_tracking_policy();
105186
break;
106187
case Ydb::Cms::ScaleRecommenderPolicies_ScaleRecommenderPolicy::POLICY_NOT_SET:
188+
Policy = std::monostate();
107189
break;
108190
}
109191
}
@@ -134,6 +216,7 @@ TGetDatabaseStatusResult::TGetDatabaseStatusResult(TStatus&& status, const Ydb::
134216
ResourcesKind_ = proto.serverless_resources();
135217
break;
136218
case Ydb::Cms::GetDatabaseStatusResult::RESOURCES_KIND_NOT_SET:
219+
ResourcesKind_ = std::monostate();
137220
break;
138221
}
139222
}
@@ -146,7 +229,7 @@ EState TGetDatabaseStatusResult::GetState() const {
146229
return State_;
147230
}
148231

149-
const std::variant<TResources, TSharedResources, TServerlessResources>& TGetDatabaseStatusResult::GetResourcesKind() const {
232+
const TResourcesKind& TGetDatabaseStatusResult::GetResourcesKind() const {
150233
return ResourcesKind_;
151234
}
152235

@@ -176,67 +259,32 @@ const TScaleRecommenderPolicies& TGetDatabaseStatusResult::GetScaleRecommenderPo
176259

177260
void TGetDatabaseStatusResult::SerializeTo(Ydb::Cms::CreateDatabaseRequest& request) const {
178261
request.set_path(Path_);
179-
if (std::holds_alternative<NCms::TResources>(ResourcesKind_)) {
180-
const auto& resources = std::get<NCms::TResources>(ResourcesKind_);
181-
for (const auto& storageUnit : resources.StorageUnits) {
182-
auto* protoUnit = request.mutable_resources()->add_storage_units();
183-
protoUnit->set_unit_kind(storageUnit.UnitKind);
184-
protoUnit->set_count(storageUnit.Count);
185-
}
186-
for (const auto& computationalUnit : resources.ComputationalUnits) {
187-
auto* protoUnit = request.mutable_resources()->add_computational_units();
188-
protoUnit->set_unit_kind(computationalUnit.UnitKind);
189-
protoUnit->set_count(computationalUnit.Count);
190-
protoUnit->set_availability_zone(computationalUnit.AvailabilityZone);
191-
}
192-
} else if (std::holds_alternative<NCms::TSharedResources>(ResourcesKind_)) {
193-
const auto& resources = std::get<NCms::TSharedResources>(ResourcesKind_);
194-
for (const auto& storageUnit : resources.StorageUnits) {
195-
auto* protoUnit = request.mutable_shared_resources()->add_storage_units();
196-
protoUnit->set_unit_kind(storageUnit.UnitKind);
197-
protoUnit->set_count(storageUnit.Count);
198-
}
199-
for (const auto& computationalUnit : resources.ComputationalUnits) {
200-
auto* protoUnit = request.mutable_shared_resources()->add_computational_units();
201-
protoUnit->set_unit_kind(computationalUnit.UnitKind);
202-
protoUnit->set_count(computationalUnit.Count);
203-
protoUnit->set_availability_zone(computationalUnit.AvailabilityZone);
204-
}
205-
} else if (std::holds_alternative<NCms::TServerlessResources>(ResourcesKind_)) {
206-
const auto& resources = std::get<NCms::TServerlessResources>(ResourcesKind_);
207-
request.mutable_serverless_resources()->set_shared_database_path(resources.SharedDatabasePath);
208-
}
209-
210-
for (const auto& quota : SchemaOperationQuotas_.LeakyBucketQuotas) {
211-
auto protoQuota = request.mutable_schema_operation_quotas()->add_leaky_bucket_quotas();
212-
protoQuota->set_bucket_seconds(quota.BucketSeconds);
213-
protoQuota->set_bucket_size(quota.BucketSize);
214-
}
262+
SerializeToImpl(ResourcesKind_, SchemaOperationQuotas_, DatabaseQuotas_, ScaleRecommenderPolicies_, request);
263+
}
215264

216-
request.mutable_database_quotas()->set_data_size_hard_quota(DatabaseQuotas_.DataSizeHardQuota);
217-
request.mutable_database_quotas()->set_data_size_soft_quota(DatabaseQuotas_.DataSizeSoftQuota);
218-
request.mutable_database_quotas()->set_data_stream_shards_quota(DatabaseQuotas_.DataStreamShardsQuota);
219-
request.mutable_database_quotas()->set_data_stream_reserved_storage_quota(DatabaseQuotas_.DataStreamReservedStorageQuota);
220-
request.mutable_database_quotas()->set_ttl_min_run_internal_seconds(DatabaseQuotas_.TtlMinRunInternalSeconds);
221-
222-
for (const auto& quota : DatabaseQuotas_.StorageQuotas) {
223-
auto protoQuota = request.mutable_database_quotas()->add_storage_quotas();
224-
protoQuota->set_unit_kind(quota.UnitKind);
225-
protoQuota->set_data_size_hard_quota(quota.DataSizeHardQuota);
226-
protoQuota->set_data_size_soft_quota(quota.DataSizeSoftQuota);
227-
}
265+
TCreateDatabaseSettings::TCreateDatabaseSettings(const Ydb::Cms::CreateDatabaseRequest& request)
266+
: SchemaOperationQuotas_(request.schema_operation_quotas())
267+
, DatabaseQuotas_(request.database_quotas())
268+
, ScaleRecommenderPolicies_(request.scale_recommender_policies())
269+
{
270+
switch (request.resources_kind_case()) {
271+
case Ydb::Cms::CreateDatabaseRequest::kResources:
272+
ResourcesKind_ = TResources(request.resources());
273+
break;
274+
case Ydb::Cms::CreateDatabaseRequest::kSharedResources:
275+
ResourcesKind_ = TSharedResources(request.shared_resources());
276+
break;
277+
case Ydb::Cms::CreateDatabaseRequest::kServerlessResources:
278+
ResourcesKind_ = request.serverless_resources();
279+
break;
280+
case Ydb::Cms::CreateDatabaseRequest::RESOURCES_KIND_NOT_SET:
281+
ResourcesKind_ = std::monostate();
282+
break;
283+
}
284+
}
228285

229-
for (const auto& policy : ScaleRecommenderPolicies_.Policies) {
230-
auto* protoPolicy = request.mutable_scale_recommender_policies()->add_policies();
231-
if (std::holds_alternative<NCms::TTargetTrackingPolicy>(policy.Policy)) {
232-
const auto& targetTracking = std::get<NCms::TTargetTrackingPolicy>(policy.Policy);
233-
auto* protoTargetTracking = protoPolicy->mutable_target_tracking_policy();
234-
if (std::holds_alternative<NCms::TTargetTrackingPolicy::TAverageCpuUtilizationPercent>(targetTracking.Target)) {
235-
const auto& target = std::get<NCms::TTargetTrackingPolicy::TAverageCpuUtilizationPercent>(targetTracking.Target);
236-
protoTargetTracking->set_average_cpu_utilization_percent(target);
237-
}
238-
}
239-
}
286+
void TCreateDatabaseSettings::SerializeTo(Ydb::Cms::CreateDatabaseRequest& request) const {
287+
SerializeToImpl(ResourcesKind_, SchemaOperationQuotas_, DatabaseQuotas_, ScaleRecommenderPolicies_, request);
240288
}
241289

242290
class TCmsClient::TImpl : public TClientImplCommon<TCmsClient::TImpl> {
@@ -246,7 +294,7 @@ class TCmsClient::TImpl : public TClientImplCommon<TCmsClient::TImpl> {
246294
{ }
247295

248296
TAsyncListDatabasesResult ListDatabases(const TListDatabasesSettings& settings) {
249-
Ydb::Cms::ListDatabasesRequest request;
297+
auto request = MakeOperationRequest<Ydb::Cms::ListDatabasesRequest>(settings);
250298

251299
auto promise = NThreading::NewPromise<TListDatabasesResult>();
252300

@@ -272,7 +320,7 @@ class TCmsClient::TImpl : public TClientImplCommon<TCmsClient::TImpl> {
272320
}
273321

274322
TAsyncGetDatabaseStatusResult GetDatabaseStatus(const std::string& path, const TGetDatabaseStatusSettings& settings) {
275-
Ydb::Cms::GetDatabaseStatusRequest request;
323+
auto request = MakeOperationRequest<Ydb::Cms::GetDatabaseStatusRequest>(settings);
276324
request.set_path(path);
277325

278326
auto promise = NThreading::NewPromise<TGetDatabaseStatusResult>();
@@ -297,6 +345,17 @@ class TCmsClient::TImpl : public TClientImplCommon<TCmsClient::TImpl> {
297345

298346
return promise.GetFuture();
299347
}
348+
349+
TAsyncStatus CreateDatabase(const std::string& path, const TCreateDatabaseSettings& settings) {
350+
auto request = MakeOperationRequest<Ydb::Cms::CreateDatabaseRequest>(settings);
351+
request.set_path(path);
352+
settings.SerializeTo(request);
353+
354+
return RunSimple<Ydb::Cms::V1::CmsService, Ydb::Cms::CreateDatabaseRequest, Ydb::Cms::CreateDatabaseResponse>(
355+
std::move(request),
356+
&Ydb::Cms::V1::CmsService::Stub::AsyncCreateDatabase,
357+
TRpcRequestSettings::Make(settings));
358+
}
300359
};
301360

302361
TCmsClient::TCmsClient(const TDriver& driver, const TCommonClientSettings& settings)
@@ -314,4 +373,11 @@ TAsyncGetDatabaseStatusResult TCmsClient::GetDatabaseStatus(
314373
return Impl_->GetDatabaseStatus(path, settings);
315374
}
316375

376+
TAsyncStatus TCmsClient::CreateDatabase(
377+
const std::string& path,
378+
const TCreateDatabaseSettings& settings)
379+
{
380+
return Impl_->CreateDatabase(path, settings);
381+
}
382+
317383
} // namespace NYdb::NCms

‎src/client/driver/driver.cpp

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,4 +223,39 @@ void TDriver::Stop(bool wait) {
223223
Impl_->Stop(wait);
224224
}
225225

226+
TDriverConfig TDriver::GetConfig() const {
227+
TDriverConfig config;
228+
229+
config.SetEndpoint(Impl_->DefaultDiscoveryEndpoint_);
230+
config.SetNetworkThreadsNum(Impl_->NetworkThreadsNum_);
231+
config.SetClientThreadsNum(Impl_->ClientThreadsNum_);
232+
config.SetMaxClientQueueSize(Impl_->MaxQueuedResponses_);
233+
if (Impl_->SslCredentials_.IsEnabled) {
234+
config.UseSecureConnection(Impl_->SslCredentials_.CaCert);
235+
}
236+
config.UseClientCertificate(Impl_->SslCredentials_.Cert, Impl_->SslCredentials_.PrivateKey);
237+
config.SetCredentialsProviderFactory(Impl_->DefaultCredentialsProviderFactory_);
238+
config.SetDatabase(Impl_->DefaultDatabase_);
239+
config.SetDiscoveryMode(Impl_->DefaultDiscoveryMode_);
240+
config.SetMaxQueuedRequests(Impl_->MaxQueuedRequests_);
241+
config.SetGrpcMemoryQuota(Impl_->MemoryQuota_);
242+
config.SetTcpKeepAliveSettings(
243+
Impl_->TcpKeepAliveSettings_.Enabled,
244+
Impl_->TcpKeepAliveSettings_.Idle,
245+
Impl_->TcpKeepAliveSettings_.Count,
246+
Impl_->TcpKeepAliveSettings_.Interval
247+
);
248+
config.SetDrainOnDtors(Impl_->DrainOnDtors_);
249+
config.SetBalancingPolicy(Impl_->BalancingSettings_.Policy, Impl_->BalancingSettings_.PolicyParams);
250+
config.SetGRpcKeepAliveTimeout(Impl_->GRpcKeepAliveTimeout_);
251+
config.SetGRpcKeepAlivePermitWithoutCalls(Impl_->GRpcKeepAlivePermitWithoutCalls_);
252+
config.SetSocketIdleTimeout(Impl_->SocketIdleTimeout_);
253+
config.SetMaxInboundMessageSize(Impl_->MaxInboundMessageSize_);
254+
config.SetMaxOutboundMessageSize(Impl_->MaxOutboundMessageSize_);
255+
config.SetMaxMessageSize(Impl_->MaxMessageSize_);
256+
config.Impl_->Log = Impl_->Log;
257+
258+
return config;
259+
}
260+
226261
} // namespace NYdb

0 commit comments

Comments
 (0)