Skip to content

Commit 47e79c0

Browse files
GrigoriyPAGrigoriyPA
authored
Merge 49151bc into 9e26730 
2 parents 9e26730 + 49151bc commit 47e79c0

10 files changed

+160
-35
lines changed

‎ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,11 @@ NMetadata::NModifications::TOperationParsingResult TResourcePoolClassifierManage
5353
} catch (...) {
5454
throw yexception() << "Failed to parse property " << property << ": " << CurrentExceptionMessage();
5555
}
56-
} else if (!featuresExtractor.ExtractResetFeature(property)) {
56+
} else if (featuresExtractor.ExtractResetFeature(property)) {
57+
if (property == "resource_pool") {
58+
ythrow yexception() << "Cannot reset required property resource_pool";
59+
}
60+
} else {
5761
continue;
5862
}
5963

@@ -65,6 +69,19 @@ NMetadata::NModifications::TOperationParsingResult TResourcePoolClassifierManage
6569
}
6670
}
6771

72+
if (context.GetActivityType() == EActivityType::Create) {
73+
if (!configJson.GetMap().contains("resource_pool")) {
74+
ythrow yexception() << "Missing required property resource_pool";
75+
}
76+
77+
static const TString extraPathSymbolsAllowed = "!\"#$%&'()*+,-.:;<=>?@[\\]^_`{|}~";
78+
const auto& name = settings.GetObjectId();
79+
if (const auto brokenAt = PathPartBrokenAt(name, extraPathSymbolsAllowed); brokenAt != name.end()) {
80+
ythrow yexception() << "Symbol '" << *brokenAt << "'" << " is not allowed in the resource pool classifier name '" << name << "'";
81+
}
82+
}
83+
resourcePoolClassifierSettings.Validate();
84+
6885
NJsonWriter::TBuf writer;
6986
writer.WriteJsonValue(&configJson);
7087
result.SetColumn(TResourcePoolClassifierConfig::TDecoder::ConfigJson, NMetadata::NInternal::TYDBValue::Utf8(writer.Str()));

‎ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -422,17 +422,19 @@ class TResourcePoolsCache {
422422
struct TClassifierInfo {
423423
const TString Membername;
424424
const TString PoolId;
425+
const i64 Rank;
425426

426427
TClassifierInfo(const NResourcePool::TClassifierSettings& classifierSettings)
427428
: Membername(classifierSettings.Membername)
428429
, PoolId(classifierSettings.ResourcePool)
430+
, Rank(classifierSettings.Rank)
429431
{}
430432
};
431433

432434
struct TDatabaseInfo {
433435
std::unordered_map<TString, TResourcePoolClassifierConfig> ResourcePoolsClassifiers = {};
434436
std::map<i64, TClassifierInfo> RankToClassifierInfo = {};
435-
std::unordered_map<TString, TString> UserToResourcePool = {};
437+
std::unordered_map<TString, std::pair<TString, i64>> UserToResourcePool = {};
436438
bool Serverless = false;
437439
};
438440

@@ -462,16 +464,16 @@ class TResourcePoolsCache {
462464
}
463465

464466
TDatabaseInfo& databaseInfo = *GetOrCreateDatabaseInfo(database);
465-
if (const auto& poolId = GetPoolIdFromClassifiers(database, userToken->GetUserSID(), databaseInfo, userToken, actorContext)) {
466-
return poolId;
467-
}
467+
auto [resultPoolId, resultRank] = GetPoolIdFromClassifiers(database, userToken->GetUserSID(), databaseInfo, userToken, actorContext);
468468
for (const auto& userSID : userToken->GetGroupSIDs()) {
469-
if (const auto& poolId = GetPoolIdFromClassifiers(database, userSID, databaseInfo, userToken, actorContext)) {
470-
return poolId;
469+
const auto& [poolId, rank] = GetPoolIdFromClassifiers(database, userSID, databaseInfo, userToken, actorContext);
470+
if (poolId && (!resultPoolId || resultRank > rank)) {
471+
resultPoolId = poolId;
472+
resultRank = rank;
471473
}
472474
}
473475

474-
return NResourcePool::DEFAULT_POOL_ID;
476+
return resultPoolId ? resultPoolId : NResourcePool::DEFAULT_POOL_ID;
475477
}
476478

477479
std::optional<TPoolInfo> GetPoolInfo(const TString& database, const TString& poolId, TActorContext actorContext) const {
@@ -582,13 +584,14 @@ class TResourcePoolsCache {
582584
}
583585
}
584586

585-
TString GetPoolIdFromClassifiers(const TString& database, const TString& userSID, TDatabaseInfo& databaseInfo, const TIntrusiveConstPtr<NACLib::TUserToken>& userToken, TActorContext actorContext) const {
587+
std::pair<TString, i64> GetPoolIdFromClassifiers(const TString& database, const TString& userSID, TDatabaseInfo& databaseInfo, const TIntrusiveConstPtr<NACLib::TUserToken>& userToken, TActorContext actorContext) const {
586588
auto& usersMap = databaseInfo.UserToResourcePool;
587589
if (const auto it = usersMap.find(userSID); it != usersMap.end()) {
588590
return it->second;
589591
}
590592

591593
TString poolId = "";
594+
i64 rank = -1;
592595
for (const auto& [_, classifier] : databaseInfo.RankToClassifierInfo) {
593596
if (classifier.Membername != userSID) {
594597
continue;
@@ -605,11 +608,12 @@ class TResourcePoolsCache {
605608
}
606609

607610
poolId = classifier.PoolId;
611+
rank = classifier.Rank;
608612
break;
609613
}
610614

611-
usersMap[userSID] = poolId;
612-
return poolId;
615+
usersMap[userSID] = {poolId, rank};
616+
return {poolId, rank};
613617
}
614618

615619
TDatabaseInfo* GetOrCreateDatabaseInfo(const TString& database) {

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

Lines changed: 51 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6877,8 +6877,8 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
68776877
// ALTER RESOURCE POOL CLASSIFIER
68786878
checkDisabled(R"(
68796879
ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier
6880-
SET (RANK = 1, MEMBERNAME = "test@user"),
6881-
RESET (RESOURCE_POOL);
6880+
SET (RANK = 1, RESOURCE_POOL = "test"),
6881+
RESET (MEMBERNAME);
68826882
)");
68836883

68846884
// DROP RESOURCE POOL CLASSIFIER
@@ -6911,8 +6911,8 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
69116911

69126912
const auto& alterSql = R"(
69136913
ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier
6914-
SET (RANK = 1, MEMBERNAME = "test@user"),
6915-
RESET (RESOURCE_POOL);
6914+
SET (RANK = 1, RESOURCE_POOL = "test"),
6915+
RESET (MEMBERNAME);
69166916
)";
69176917

69186918
const auto& dropSql = "DROP RESOURCE POOL CLASSIFIER MyResourcePoolClassifier;";
@@ -6951,6 +6951,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
69516951

69526952
auto result = session.ExecuteSchemeQuery(R"(
69536953
CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
6954+
RESOURCE_POOL="test",
69546955
ANOTHER_PROPERTY=20
69556956
);)").GetValueSync();
69566957
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
@@ -6966,10 +6967,40 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
69666967

69676968
result = session.ExecuteSchemeQuery(R"(
69686969
CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
6970+
RESOURCE_POOL="test",
69696971
RANK="StringValue"
69706972
);)").GetValueSync();
69716973
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
69726974
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Failed to parse property rank:");
6975+
6976+
result = session.ExecuteSchemeQuery(R"(
6977+
CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
6978+
RANK="0"
6979+
);)").GetValueSync();
6980+
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
6981+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Missing required property resource_pool");
6982+
6983+
result = session.ExecuteSchemeQuery(R"(
6984+
ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier
6985+
RESET (RESOURCE_POOL);
6986+
)").GetValueSync();
6987+
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
6988+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Cannot reset required property resource_pool");
6989+
6990+
result = session.ExecuteSchemeQuery(R"(
6991+
CREATE RESOURCE POOL CLASSIFIER `MyResource/PoolClassifier` WITH (
6992+
RESOURCE_POOL="test"
6993+
);)").GetValueSync();
6994+
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
6995+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Symbol '/' is not allowed in the resource pool classifier name 'MyResource/PoolClassifier'");
6996+
6997+
result = session.ExecuteSchemeQuery(TStringBuilder() << R"(
6998+
CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
6999+
RESOURCE_POOL="test",
7000+
MEMBERNAME=")" << BUILTIN_ACL_METADATA << R"("
7001+
);)").GetValueSync();
7002+
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
7003+
UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), TStringBuilder() << "Invalid resource pool classifier configuration, cannot create classifier for system user " << BUILTIN_ACL_METADATA);
69737004
}
69747005

69757006
Y_UNIT_TEST(ResourcePoolClassifiersRankValidation) {
@@ -6986,13 +7017,15 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
69867017
// Create with sample rank
69877018
auto result = session.ExecuteSchemeQuery(R"(
69887019
CREATE RESOURCE POOL CLASSIFIER ClassifierRank42 WITH (
7020+
RESOURCE_POOL="test_pool",
69897021
RANK=42
69907022
);)").GetValueSync();
69917023
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToOneLineString());
69927024

69937025
// Try to create with same rank
69947026
result = session.ExecuteSchemeQuery(R"(
69957027
CREATE RESOURCE POOL CLASSIFIER AnotherClassifierRank42 WITH (
7028+
RESOURCE_POOL="test_pool",
69967029
RANK=42
69977030
);)").GetValueSync();
69987031
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
@@ -7001,13 +7034,15 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
70017034
// Create with high rank
70027035
result = session.ExecuteSchemeQuery(R"(
70037036
CREATE RESOURCE POOL CLASSIFIER `ClassifierRank2^63` WITH (
7037+
RESOURCE_POOL="test_pool",
70047038
RANK=9223372036854775807
70057039
);)").GetValueSync();
70067040
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToOneLineString());
70077041

70087042
// Try to create with auto rank
70097043
result = session.ExecuteSchemeQuery(R"(
70107044
CREATE RESOURCE POOL CLASSIFIER ClassifierRankAuto WITH (
7045+
RESOURCE_POOL="test_pool",
70117046
MEMBERNAME="test@user"
70127047
);)").GetValueSync();
70137048
UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
@@ -7065,11 +7100,12 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
70657100
// Auto rank
70667101
query = R"(
70677102
CREATE RESOURCE POOL CLASSIFIER AnotherResourcePoolClassifier WITH (
7103+
RESOURCE_POOL="test_pool",
70687104
MEMBERNAME="another@user"
70697105
);)";
70707106
result = session.ExecuteSchemeQuery(query).GetValueSync();
70717107
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
7072-
UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"membername\":\"test@user\",\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"},{\"rank\":1020,\"name\":\"AnotherResourcePoolClassifier\",\"config\":{\"membername\":\"another@user\"},\"database\":\"\\/Root\"}]}");
7108+
UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"membername\":\"test@user\",\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"},{\"rank\":1020,\"name\":\"AnotherResourcePoolClassifier\",\"config\":{\"membername\":\"another@user\",\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"}]}");
70737109
}
70747110

70757111
Y_UNIT_TEST(DoubleCreateResourcePoolClassifier) {
@@ -7086,6 +7122,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
70867122
{
70877123
auto query = R"(
70887124
CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
7125+
RESOURCE_POOL="test_pool",
70897126
RANK=20
70907127
);)";
70917128
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
@@ -7095,6 +7132,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
70957132
{
70967133
auto query = R"(
70977134
CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
7135+
RESOURCE_POOL="test_pool",
70987136
RANK=1
70997137
);)";
71007138
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
@@ -7141,22 +7179,23 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
71417179
{
71427180
auto query = R"(
71437181
CREATE RESOURCE POOL CLASSIFIER AnotherResourcePoolClassifier WITH (
7182+
RESOURCE_POOL="test_pool",
71447183
RANK=42
71457184
);)";
71467185
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
71477186
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
7148-
UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"membername\":\"test@user\",\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"},{\"rank\":42,\"name\":\"AnotherResourcePoolClassifier\",\"config\":{},\"database\":\"\\/Root\"}]}");
7187+
UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"membername\":\"test@user\",\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"},{\"rank\":42,\"name\":\"AnotherResourcePoolClassifier\",\"config\":{\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"}]}");
71497188
}
71507189

71517190
// Test reset
71527191
{
71537192
auto query = R"(
71547193
ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier
7155-
RESET (RANK, RESOURCE_POOL);
7194+
RESET (RANK, MEMBERNAME);
71567195
)";
71577196
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
71587197
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
7159-
UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":1042,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"membername\":\"test@user\",\"resource_pool\":\"default\"},\"database\":\"\\/Root\"},{\"rank\":42,\"name\":\"AnotherResourcePoolClassifier\",\"config\":{},\"database\":\"\\/Root\"}]}");
7198+
UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":1042,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"membername\":\"\",\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"},{\"rank\":42,\"name\":\"AnotherResourcePoolClassifier\",\"config\":{\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"}]}");
71607199
}
71617200
}
71627201

@@ -7173,8 +7212,8 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
71737212

71747213
auto query = R"(
71757214
ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier
7176-
SET (MEMBERNAME = "test@user", RANK = 100),
7177-
RESET (RESOURCE_POOL);
7215+
SET (RESOURCE_POOL = "test", RANK = 100),
7216+
RESET (MEMBERNAME);
71787217
)";
71797218
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
71807219
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString());
@@ -7195,11 +7234,12 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
71957234
{
71967235
auto query = R"(
71977236
CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
7237+
RESOURCE_POOL="test_pool",
71987238
RANK=20
71997239
);)";
72007240
auto result = session.ExecuteSchemeQuery(query).GetValueSync();
72017241
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
7202-
UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{},\"database\":\"\\/Root\"}]}");
7242+
UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"}]}");
72037243
}
72047244

72057245
{

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -518,7 +518,7 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
518518
UNIT_ASSERT_C(settings.PoolId_, "Query pool id is not specified");
519519

520520
auto event = std::make_unique<TEvKqp::TEvQueryRequest>();
521-
event->Record.SetUserToken(NACLib::TUserToken("", settings.UserSID_, {}).SerializeAsString());
521+
event->Record.SetUserToken(NACLib::TUserToken("", settings.UserSID_, settings.GroupSIDs_).SerializeAsString());
522522

523523
auto request = event->Record.MutableRequest();
524524
request->SetQuery(query);

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

Lines changed: 1 addition & 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(TVector<TString>, GroupSIDs, {});
2930
FLUENT_SETTING_DEFAULT(TString, Database, "");
3031

3132
// Runner settings

0 commit comments

Comments
 (0)