Skip to content

Commit 6f6b4cb

Browse files
GrigoriyPAGrigoriyPA
committed
YQ WM added unit tests for sls disable (#8687)
1 parent c013f50 commit 6f6b4cb

File tree

8 files changed

+273
-6
lines changed

8 files changed

+273
-6
lines changed

‎ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include <ydb/core/kqp/gateway/behaviour/resource_pool_classifier/fetcher.h>
22
#include <ydb/core/kqp/ut/common/kqp_ut_common.h>
33
#include <ydb/core/kqp/ut/common/columnshard.h>
4+
#include <ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h>
45
#include <ydb/core/tx/columnshard/hooks/testing/controller.h>
56
#include <ydb/core/formats/arrow/arrow_helpers.h>
67
#include <ydb/core/tx/tx_proxy/proxy.h>
@@ -6211,6 +6212,57 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
62116212
"Path does not exist");
62126213
}
62136214

6215+
Y_UNIT_TEST(DisableResourcePoolsOnServerless) {
6216+
auto ydb = NWorkload::TYdbSetupSettings()
6217+
.CreateSampleTenants(true)
6218+
.EnableResourcePoolsOnServerless(false)
6219+
.Create();
6220+
6221+
auto checkDisabled = [](const auto& result) {
6222+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::GENERIC_ERROR, result.GetIssues().ToString());
6223+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Resource pools are disabled for serverless domains. Please contact your system administrator to enable it");
6224+
};
6225+
6226+
auto checkNotFound = [](const auto& result) {
6227+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::GENERIC_ERROR, result.GetIssues().ToString());
6228+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Path does not exist");
6229+
};
6230+
6231+
const auto& createSql = R"(
6232+
CREATE RESOURCE POOL MyResourcePool WITH (
6233+
CONCURRENT_QUERY_LIMIT=20,
6234+
QUEUE_SIZE=1000
6235+
);)";
6236+
6237+
const auto& alterSql = R"(
6238+
ALTER RESOURCE POOL MyResourcePool
6239+
SET (CONCURRENT_QUERY_LIMIT = 30, QUEUE_SIZE = 100),
6240+
RESET (QUERY_MEMORY_LIMIT_PERCENT_PER_NODE);
6241+
)";
6242+
6243+
const auto& dropSql = "DROP RESOURCE POOL MyResourcePool;";
6244+
6245+
auto settings = NWorkload::TQueryRunnerSettings().PoolId("");
6246+
6247+
// Dedicated, enabled
6248+
settings.Database(ydb->GetSettings().GetDedicatedTenantName()).NodeIndex(1);
6249+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(createSql, settings));
6250+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(alterSql, settings));
6251+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(dropSql, settings));
6252+
6253+
// Shared, enabled
6254+
settings.Database(ydb->GetSettings().GetSharedTenantName()).NodeIndex(2);
6255+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(createSql, settings));
6256+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(alterSql, settings));
6257+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(dropSql, settings));
6258+
6259+
// Serverless, disabled
6260+
settings.Database(ydb->GetSettings().GetServerlessTenantName()).NodeIndex(2);
6261+
checkDisabled(ydb->ExecuteQuery(createSql, settings));
6262+
checkNotFound(ydb->ExecuteQuery(alterSql, settings));
6263+
checkNotFound(ydb->ExecuteQuery(dropSql, settings));
6264+
}
6265+
62146266
Y_UNIT_TEST(ResourcePoolsValidation) {
62156267
NKikimrConfig::TAppConfig config;
62166268
config.MutableFeatureFlags()->SetEnableResourcePools(true);
@@ -6480,6 +6532,57 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
64806532
"Classifier with name MyResourcePoolClassifier not found in database /Root");
64816533
}
64826534

6535+
Y_UNIT_TEST(DisableResourcePoolClassifiersOnServerless) {
6536+
auto ydb = NWorkload::TYdbSetupSettings()
6537+
.CreateSampleTenants(true)
6538+
.EnableResourcePoolsOnServerless(false)
6539+
.Create();
6540+
6541+
auto checkDisabled = [](const auto& result) {
6542+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::GENERIC_ERROR, result.GetIssues().ToString());
6543+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Resource pool classifiers are disabled for serverless domains. Please contact your system administrator to enable it");
6544+
};
6545+
6546+
auto checkNotFound = [](const auto& result) {
6547+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::GENERIC_ERROR, result.GetIssues().ToString());
6548+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Classifier with name MyResourcePoolClassifier not found in database");
6549+
};
6550+
6551+
const auto& createSql = R"(
6552+
CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
6553+
RANK=20,
6554+
RESOURCE_POOL="test_pool"
6555+
);)";
6556+
6557+
const auto& alterSql = R"(
6558+
ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier
6559+
SET (RANK = 1, MEMBERNAME = "test@user"),
6560+
RESET (RESOURCE_POOL);
6561+
)";
6562+
6563+
const auto& dropSql = "DROP RESOURCE POOL CLASSIFIER MyResourcePoolClassifier;";
6564+
6565+
auto settings = NWorkload::TQueryRunnerSettings().PoolId("");
6566+
6567+
// Dedicated, enabled
6568+
settings.Database(ydb->GetSettings().GetDedicatedTenantName()).NodeIndex(1);
6569+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(createSql, settings));
6570+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(alterSql, settings));
6571+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(dropSql, settings));
6572+
6573+
// Shared, enabled
6574+
settings.Database(ydb->GetSettings().GetSharedTenantName()).NodeIndex(2);
6575+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(createSql, settings));
6576+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(alterSql, settings));
6577+
NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(dropSql, settings));
6578+
6579+
// Serverless, disabled
6580+
settings.Database(ydb->GetSettings().GetServerlessTenantName()).NodeIndex(2);
6581+
checkDisabled(ydb->ExecuteQuery(createSql, settings));
6582+
checkDisabled(ydb->ExecuteQuery(alterSql, settings));
6583+
checkNotFound(ydb->ExecuteQuery(dropSql, settings));
6584+
}
6585+
64836586
Y_UNIT_TEST(ResourcePoolClassifiersValidation) {
64846587
NKikimrConfig::TAppConfig config;
64856588
config.MutableFeatureFlags()->SetEnableResourcePools(true);

‎ydb/core/kqp/ut/scheme/ya.make

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ PEERDIR(
2222
library/cpp/threading/local_executor
2323
ydb/core/kqp
2424
ydb/core/kqp/ut/common
25+
ydb/core/kqp/workload_service/ut/common
2526
ydb/core/tx/columnshard/hooks/testing
2627
ydb/library/yql/sql/pg
2728
ydb/library/yql/parser/pg_wrapper

‎ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp

Lines changed: 58 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,7 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
230230
TAppConfig GetAppConfig() const {
231231
TAppConfig appConfig;
232232
appConfig.MutableFeatureFlags()->SetEnableResourcePools(Settings_.EnableResourcePools_);
233+
appConfig.MutableFeatureFlags()->SetEnableResourcePoolsOnServerless(Settings_.EnableResourcePoolsOnServerless_);
233234
appConfig.MutableFeatureFlags()->SetEnableResourcePoolsCounters(true);
234235

235236
return appConfig;
@@ -238,7 +239,7 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
238239
void SetLoggerSettings(TServerSettings& serverSettings) const {
239240
auto loggerInitializer = [](TTestActorRuntime& runtime) {
240241
runtime.SetLogPriority(NKikimrServices::KQP_WORKLOAD_SERVICE, NLog::EPriority::PRI_TRACE);
241-
runtime.SetLogPriority(NKikimrServices::KQP_SESSION, NLog::EPriority::PRI_DEBUG);
242+
runtime.SetLogPriority(NKikimrServices::KQP_SESSION, NLog::EPriority::PRI_TRACE);
242243
};
243244

244245
serverSettings.SetLoggerInitializer(loggerInitializer);
@@ -255,16 +256,50 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
255256
.SetAppConfig(appConfig)
256257
.SetFeatureFlags(appConfig.GetFeatureFlags());
257258

259+
if (Settings_.CreateSampleTenants_) {
260+
serverSettings
261+
.SetDynamicNodeCount(2)
262+
.AddStoragePoolType(Settings_.GetDedicatedTenantName())
263+
.AddStoragePoolType(Settings_.GetSharedTenantName());
264+
}
265+
258266
SetLoggerSettings(serverSettings);
259267

260268
return serverSettings;
261269
}
262270

271+
void SetupResourcesTenant(Ydb::Cms::CreateDatabaseRequest& request, Ydb::Cms::StorageUnits* storage, const TString& name) {
272+
request.set_path(name);
273+
storage->set_unit_kind(name);
274+
storage->set_count(1);
275+
}
276+
277+
void CreateTenants() {
278+
{ // Dedicated
279+
Ydb::Cms::CreateDatabaseRequest request;
280+
SetupResourcesTenant(request, request.mutable_resources()->add_storage_units(), Settings_.GetDedicatedTenantName());
281+
Tenants_->CreateTenant(std::move(request));
282+
}
283+
284+
{ // Shared
285+
Ydb::Cms::CreateDatabaseRequest request;
286+
SetupResourcesTenant(request, request.mutable_shared_resources()->add_storage_units(), Settings_.GetSharedTenantName());
287+
Tenants_->CreateTenant(std::move(request));
288+
}
289+
290+
{ // Serverless
291+
Ydb::Cms::CreateDatabaseRequest request;
292+
request.set_path(Settings_.GetServerlessTenantName());
293+
request.mutable_serverless_resources()->set_shared_database_path(Settings_.GetSharedTenantName());
294+
Tenants_->CreateTenant(std::move(request));
295+
}
296+
}
297+
263298
void InitializeServer() {
264299
ui32 grpcPort = PortManager_.GetPort();
265300
TServerSettings serverSettings = GetServerSettings(grpcPort);
266301

267-
Server_ = std::make_unique<TServer>(serverSettings);
302+
Server_ = MakeIntrusive<TServer>(serverSettings);
268303
Server_->EnableGRpc(grpcPort);
269304
GetRuntime()->SetDispatchTimeout(FUTURE_WAIT_TIMEOUT);
270305

@@ -277,10 +312,15 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
277312

278313
TableClient_ = std::make_unique<NYdb::NTable::TTableClient>(*YdbDriver_, NYdb::NTable::TClientSettings().AuthToken("user@" BUILTIN_SYSTEM_DOMAIN));
279314
TableClientSession_ = std::make_unique<NYdb::NTable::TSession>(TableClient_->CreateSession().GetValueSync().GetSession());
315+
316+
Tenants_ = std::make_unique<TTenants>(Server_);
317+
if (Settings_.CreateSampleTenants_) {
318+
CreateTenants();
319+
}
280320
}
281321

282322
void CreateSamplePool() const {
283-
if (!Settings_.EnableResourcePools_) {
323+
if (!Settings_.EnableResourcePools_ || Settings_.CreateSampleTenants_) {
284324
return;
285325
}
286326

@@ -483,7 +523,7 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
483523
request->SetQuery(query);
484524
request->SetType(NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY);
485525
request->SetAction(NKikimrKqp::QUERY_ACTION_EXECUTE);
486-
request->SetDatabase(Settings_.DomainName_);
526+
request->SetDatabase(settings.Database_ ? settings.Database_ : Settings_.DomainName_);
487527
request->SetPoolId(*settings.PoolId_);
488528

489529
return event;
@@ -525,9 +565,10 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
525565
const TYdbSetupSettings Settings_;
526566

527567
TPortManager PortManager_;
528-
std::unique_ptr<TServer> Server_;
568+
TServer::TPtr Server_;
529569
std::unique_ptr<TClient> Client_;
530570
std::unique_ptr<TDriver> YdbDriver_;
571+
std::unique_ptr<TTenants> Tenants_;
531572

532573
std::unique_ptr<NYdb::NTable::TTableClient> TableClient_;
533574
std::unique_ptr<NYdb::NTable::TSession> TableClientSession_;
@@ -586,6 +627,18 @@ TIntrusivePtr<IYdbSetup> TYdbSetupSettings::Create() const {
586627
return MakeIntrusive<TWorkloadServiceYdbSetup>(*this);
587628
}
588629

630+
TString TYdbSetupSettings::GetDedicatedTenantName() const {
631+
return TStringBuilder() << CanonizePath(DomainName_) << "/test-dedicated";
632+
}
633+
634+
TString TYdbSetupSettings::GetSharedTenantName() const {
635+
return TStringBuilder() << CanonizePath(DomainName_) << "/test-shared";
636+
}
637+
638+
TString TYdbSetupSettings::GetServerlessTenantName() const {
639+
return TStringBuilder() << CanonizePath(DomainName_) << "/test-serverless";
640+
}
641+
589642
//// IYdbSetup
590643

591644
void IYdbSetup::WaitFor(TDuration timeout, TString description, std::function<bool(TString&)> callback) {

‎ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ struct TQueryRunnerSettings {
2626
FLUENT_SETTING_DEFAULT(ui32, NodeIndex, 0);
2727
FLUENT_SETTING_DEFAULT(std::optional<TString>, PoolId, std::nullopt);
2828
FLUENT_SETTING_DEFAULT(TString, UserSID, "user@" BUILTIN_SYSTEM_DOMAIN);
29+
FLUENT_SETTING_DEFAULT(TString, Database, "");
2930

3031
// Runner settings
3132
FLUENT_SETTING_DEFAULT(bool, HangUpDuringExecution, false);
@@ -66,7 +67,9 @@ struct TYdbSetupSettings {
6667
// Cluster settings
6768
FLUENT_SETTING_DEFAULT(ui32, NodeCount, 1);
6869
FLUENT_SETTING_DEFAULT(TString, DomainName, "Root");
70+
FLUENT_SETTING_DEFAULT(bool, CreateSampleTenants, false);
6971
FLUENT_SETTING_DEFAULT(bool, EnableResourcePools, true);
72+
FLUENT_SETTING_DEFAULT(bool, EnableResourcePoolsOnServerless, false);
7073

7174
// Default pool settings
7275
FLUENT_SETTING_DEFAULT(TString, PoolId, "sample_pool_id");
@@ -78,6 +81,10 @@ struct TYdbSetupSettings {
7881

7982
NResourcePool::TPoolSettings GetDefaultPoolSettings() const;
8083
TIntrusivePtr<IYdbSetup> Create() const;
84+
85+
TString GetDedicatedTenantName() const;
86+
TString GetSharedTenantName() const;
87+
TString GetServerlessTenantName() const;
8188
};
8289

8390
class IYdbSetup : public TThrRefBase {
@@ -127,6 +134,12 @@ struct TSampleQueries {
127134
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Request timeout exceeded, cancelling after");
128135
}
129136

137+
template <typename TResult>
138+
static void CheckNotFound(const TResult& result, const TString& poolId) {
139+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::NOT_FOUND, result.GetIssues().ToString());
140+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), TStringBuilder() << "Resource pool " << poolId << " not found or you don't have access permissions");
141+
}
142+
130143
struct TSelect42 {
131144
static constexpr char Query[] = "SELECT 42;";
132145

‎ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include <ydb/core/base/appdata_fwd.h>
2+
#include <ydb/core/base/path.h>
23

34
#include <ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h>
45

@@ -47,6 +48,56 @@ Y_UNIT_TEST_SUITE(KqpWorkloadService) {
4748
TSampleQueries::TSelect42::CheckResult(ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, TQueryRunnerSettings().PoolId("another_pool_id")));
4849
}
4950

51+
Y_UNIT_TEST(WorkloadServiceDisabledByFeatureFlagOnServerless) {
52+
auto ydb = TYdbSetupSettings()
53+
.CreateSampleTenants(true)
54+
.EnableResourcePoolsOnServerless(false)
55+
.Create();
56+
57+
const TString& poolId = "another_pool_id";
58+
auto settings = TQueryRunnerSettings().PoolId(poolId);
59+
60+
// Dedicated, enabled
61+
TSampleQueries::CheckNotFound(ydb->ExecuteQuery(
62+
TSampleQueries::TSelect42::Query,
63+
settings.Database(ydb->GetSettings().GetDedicatedTenantName()).NodeIndex(1)
64+
), poolId);
65+
66+
// Shared, enabled
67+
TSampleQueries::CheckNotFound(ydb->ExecuteQuery(
68+
TSampleQueries::TSelect42::Query,
69+
settings.Database(ydb->GetSettings().GetSharedTenantName()).NodeIndex(2)
70+
), poolId);
71+
72+
// Serverless, disabled
73+
TSampleQueries::TSelect42::CheckResult(ydb->ExecuteQuery(
74+
TSampleQueries::TSelect42::Query,
75+
settings.Database(ydb->GetSettings().GetServerlessTenantName()).NodeIndex(2)
76+
));
77+
}
78+
79+
Y_UNIT_TEST(WorkloadServiceDisabledByInvalidDatabasePath) {
80+
auto ydb = TYdbSetupSettings().Create();
81+
82+
const TString& poolId = "another_pool_id";
83+
auto settings = TQueryRunnerSettings().PoolId(poolId);
84+
85+
TSampleQueries::CheckNotFound(ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, settings), poolId);
86+
87+
const TString& tabmleName = "sub_path";
88+
ydb->ExecuteSchemeQuery(TStringBuilder() << R"(
89+
CREATE TABLE )" << tabmleName << R"( (
90+
Key Int32,
91+
PRIMARY KEY (Key)
92+
);
93+
)");
94+
95+
TSampleQueries::TSelect42::CheckResult(ydb->ExecuteQuery(
96+
TSampleQueries::TSelect42::Query,
97+
settings.Database(TStringBuilder() << CanonizePath(ydb->GetSettings().DomainName_) << "/" << tabmleName)
98+
));
99+
}
100+
50101
TQueryRunnerResultAsync StartQueueSizeCheckRequests(TIntrusivePtr<IYdbSetup> ydb, const TQueryRunnerSettings& settings) {
51102
// One of these requests should be rejected by QueueSize
52103
auto firstRequest = ydb->ExecuteQueryAsync(TSampleQueries::TSelect42::Query, settings);

0 commit comments

Comments
 (0)