From 93a30e6e6d1b8c910899dc0362f78cb1d16d82bd Mon Sep 17 00:00:00 2001 From: Grigoriy Pisarenko Date: Tue, 10 Sep 2024 18:23:19 +0000 Subject: [PATCH 1/2] Added classifiers validations --- .../resource_pool_classifier/manager.cpp | 19 +++++- .../proxy_service/kqp_proxy_service_impl.h | 24 ++++--- ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp | 62 ++++++++++++++--- .../common/kqp_workload_service_ut_common.cpp | 2 +- .../common/kqp_workload_service_ut_common.h | 1 + .../ut/kqp_workload_service_ut.cpp | 67 +++++++++++++++---- .../resource_pool_classifier_settings.cpp | 9 +++ .../resource_pool_classifier_settings.h | 2 + .../resource_pool_classifier_settings_ut.cpp | 8 +++ ydb/core/resource_pools/ya.make | 1 + 10 files changed, 160 insertions(+), 35 deletions(-) diff --git a/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp b/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp index 723073954f28..dc884cd029ba 100644 --- a/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp +++ b/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp @@ -53,7 +53,11 @@ NMetadata::NModifications::TOperationParsingResult TResourcePoolClassifierManage } catch (...) { throw yexception() << "Failed to parse property " << property << ": " << CurrentExceptionMessage(); } - } else if (!featuresExtractor.ExtractResetFeature(property)) { + } else if (featuresExtractor.ExtractResetFeature(property)) { + if (property == "resource_pool") { + ythrow yexception() << "Cannot reset required property resource_pool"; + } + } else { continue; } @@ -65,6 +69,19 @@ NMetadata::NModifications::TOperationParsingResult TResourcePoolClassifierManage } } + if (context.GetActivityType() == EActivityType::Create) { + if (!configJson.GetMap().contains("resource_pool")) { + ythrow yexception() << "Missing required property resource_pool"; + } + + static const TString extraPathSymbolsAllowed = "!\"#$%&'()*+,-.:;<=>?@[\\]^_`{|}~"; + const auto& name = settings.GetObjectId(); + if (const auto brokenAt = PathPartBrokenAt(name, extraPathSymbolsAllowed); brokenAt != name.end()) { + ythrow yexception() << "Symbol '" << *brokenAt << "'" << " is not allowed in the resource pool classifier name '" << name << "'"; + } + } + resourcePoolClassifierSettings.Validate(); + NJsonWriter::TBuf writer; writer.WriteJsonValue(&configJson); result.SetColumn(TResourcePoolClassifierConfig::TDecoder::ConfigJson, NMetadata::NInternal::TYDBValue::Utf8(writer.Str())); diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h b/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h index 33b1d6e9172b..b55c8e67d2e5 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h +++ b/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h @@ -422,17 +422,19 @@ class TResourcePoolsCache { struct TClassifierInfo { const TString Membername; const TString PoolId; + const i64 Rank; TClassifierInfo(const NResourcePool::TClassifierSettings& classifierSettings) : Membername(classifierSettings.Membername) , PoolId(classifierSettings.ResourcePool) + , Rank(classifierSettings.Rank) {} }; struct TDatabaseInfo { std::unordered_map ResourcePoolsClassifiers = {}; std::map RankToClassifierInfo = {}; - std::unordered_map UserToResourcePool = {}; + std::unordered_map> UserToResourcePool = {}; bool Serverless = false; }; @@ -462,16 +464,16 @@ class TResourcePoolsCache { } TDatabaseInfo& databaseInfo = *GetOrCreateDatabaseInfo(database); - if (const auto& poolId = GetPoolIdFromClassifiers(database, userToken->GetUserSID(), databaseInfo, userToken, actorContext)) { - return poolId; - } + auto [resultPoolId, resultRank] = GetPoolIdFromClassifiers(database, userToken->GetUserSID(), databaseInfo, userToken, actorContext); for (const auto& userSID : userToken->GetGroupSIDs()) { - if (const auto& poolId = GetPoolIdFromClassifiers(database, userSID, databaseInfo, userToken, actorContext)) { - return poolId; + const auto& [poolId, rank] = GetPoolIdFromClassifiers(database, userSID, databaseInfo, userToken, actorContext); + if (poolId && (!resultPoolId || resultRank > rank)) { + resultPoolId = poolId; + resultRank = rank; } } - return NResourcePool::DEFAULT_POOL_ID; + return resultPoolId ? resultPoolId : NResourcePool::DEFAULT_POOL_ID; } std::optional GetPoolInfo(const TString& database, const TString& poolId, TActorContext actorContext) const { @@ -582,13 +584,14 @@ class TResourcePoolsCache { } } - TString GetPoolIdFromClassifiers(const TString& database, const TString& userSID, TDatabaseInfo& databaseInfo, const TIntrusiveConstPtr& userToken, TActorContext actorContext) const { + std::pair GetPoolIdFromClassifiers(const TString& database, const TString& userSID, TDatabaseInfo& databaseInfo, const TIntrusiveConstPtr& userToken, TActorContext actorContext) const { auto& usersMap = databaseInfo.UserToResourcePool; if (const auto it = usersMap.find(userSID); it != usersMap.end()) { return it->second; } TString poolId = ""; + i64 rank = -1; for (const auto& [_, classifier] : databaseInfo.RankToClassifierInfo) { if (classifier.Membername != userSID) { continue; @@ -605,11 +608,12 @@ class TResourcePoolsCache { } poolId = classifier.PoolId; + rank = classifier.Rank; break; } - usersMap[userSID] = poolId; - return poolId; + usersMap[userSID] = {poolId, rank}; + return {poolId, rank}; } TDatabaseInfo* GetOrCreateDatabaseInfo(const TString& database) { diff --git a/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp b/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp index c39d824c1279..3651546b3914 100644 --- a/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp +++ b/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp @@ -6877,8 +6877,8 @@ Y_UNIT_TEST_SUITE(KqpScheme) { // ALTER RESOURCE POOL CLASSIFIER checkDisabled(R"( ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier - SET (RANK = 1, MEMBERNAME = "test@user"), - RESET (RESOURCE_POOL); + SET (RANK = 1, RESOURCE_POOL = "test"), + RESET (MEMBERNAME); )"); // DROP RESOURCE POOL CLASSIFIER @@ -6911,8 +6911,8 @@ Y_UNIT_TEST_SUITE(KqpScheme) { const auto& alterSql = R"( ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier - SET (RANK = 1, MEMBERNAME = "test@user"), - RESET (RESOURCE_POOL); + SET (RANK = 1, RESOURCE_POOL = "test"), + RESET (MEMBERNAME); )"; const auto& dropSql = "DROP RESOURCE POOL CLASSIFIER MyResourcePoolClassifier;"; @@ -6951,6 +6951,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { auto result = session.ExecuteSchemeQuery(R"( CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RESOURCE_POOL="test", ANOTHER_PROPERTY=20 );)").GetValueSync(); UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR); @@ -6966,10 +6967,40 @@ Y_UNIT_TEST_SUITE(KqpScheme) { result = session.ExecuteSchemeQuery(R"( CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RESOURCE_POOL="test", RANK="StringValue" );)").GetValueSync(); UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR); UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Failed to parse property rank:"); + + result = session.ExecuteSchemeQuery(R"( + CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RANK="0" + );)").GetValueSync(); + UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR); + UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Missing required property resource_pool"); + + result = session.ExecuteSchemeQuery(R"( + ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier + RESET (RESOURCE_POOL); + )").GetValueSync(); + UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR); + UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Cannot reset required property resource_pool"); + + result = session.ExecuteSchemeQuery(R"( + CREATE RESOURCE POOL CLASSIFIER `MyResource/PoolClassifier` WITH ( + RESOURCE_POOL="test" + );)").GetValueSync(); + UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR); + UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Symbol '/' is not allowed in the resource pool classifier name 'MyResource/PoolClassifier'"); + + result = session.ExecuteSchemeQuery(TStringBuilder() << R"( + CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RESOURCE_POOL="test", + MEMBERNAME=")" << BUILTIN_ACL_METADATA << R"(" + );)").GetValueSync(); + UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR); + UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), TStringBuilder() << "Invalid resource pool classifier configuration, cannot create classifier for system user " << BUILTIN_ACL_METADATA); } Y_UNIT_TEST(ResourcePoolClassifiersRankValidation) { @@ -6986,6 +7017,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { // Create with sample rank auto result = session.ExecuteSchemeQuery(R"( CREATE RESOURCE POOL CLASSIFIER ClassifierRank42 WITH ( + RESOURCE_POOL="test_pool", RANK=42 );)").GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToOneLineString()); @@ -6993,6 +7025,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { // Try to create with same rank result = session.ExecuteSchemeQuery(R"( CREATE RESOURCE POOL CLASSIFIER AnotherClassifierRank42 WITH ( + RESOURCE_POOL="test_pool", RANK=42 );)").GetValueSync(); UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR); @@ -7001,6 +7034,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { // Create with high rank result = session.ExecuteSchemeQuery(R"( CREATE RESOURCE POOL CLASSIFIER `ClassifierRank2^63` WITH ( + RESOURCE_POOL="test_pool", RANK=9223372036854775807 );)").GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToOneLineString()); @@ -7008,6 +7042,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { // Try to create with auto rank result = session.ExecuteSchemeQuery(R"( CREATE RESOURCE POOL CLASSIFIER ClassifierRankAuto WITH ( + RESOURCE_POOL="test_pool", MEMBERNAME="test@user" );)").GetValueSync(); UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR); @@ -7065,11 +7100,12 @@ Y_UNIT_TEST_SUITE(KqpScheme) { // Auto rank query = R"( CREATE RESOURCE POOL CLASSIFIER AnotherResourcePoolClassifier WITH ( + RESOURCE_POOL="test_pool", MEMBERNAME="another@user" );)"; result = session.ExecuteSchemeQuery(query).GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); - 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\"}]}"); + 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\"}]}"); } Y_UNIT_TEST(DoubleCreateResourcePoolClassifier) { @@ -7086,6 +7122,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { { auto query = R"( CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RESOURCE_POOL="test_pool", RANK=20 );)"; auto result = session.ExecuteSchemeQuery(query).GetValueSync(); @@ -7095,6 +7132,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { { auto query = R"( CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RESOURCE_POOL="test_pool", RANK=1 );)"; auto result = session.ExecuteSchemeQuery(query).GetValueSync(); @@ -7141,22 +7179,23 @@ Y_UNIT_TEST_SUITE(KqpScheme) { { auto query = R"( CREATE RESOURCE POOL CLASSIFIER AnotherResourcePoolClassifier WITH ( + RESOURCE_POOL="test_pool", RANK=42 );)"; auto result = session.ExecuteSchemeQuery(query).GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); - 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\"}]}"); + 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\"}]}"); } // Test reset { auto query = R"( ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier - RESET (RANK, RESOURCE_POOL); + RESET (RANK, MEMBERNAME); )"; auto result = session.ExecuteSchemeQuery(query).GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); - 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\"}]}"); + 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\"}]}"); } } @@ -7173,8 +7212,8 @@ Y_UNIT_TEST_SUITE(KqpScheme) { auto query = R"( ALTER RESOURCE POOL CLASSIFIER MyResourcePoolClassifier - SET (MEMBERNAME = "test@user", RANK = 100), - RESET (RESOURCE_POOL); + SET (RESOURCE_POOL = "test", RANK = 100), + RESET (MEMBERNAME); )"; auto result = session.ExecuteSchemeQuery(query).GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString()); @@ -7195,11 +7234,12 @@ Y_UNIT_TEST_SUITE(KqpScheme) { { auto query = R"( CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RESOURCE_POOL="test_pool", RANK=20 );)"; auto result = session.ExecuteSchemeQuery(query).GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString()); - UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{},\"database\":\"\\/Root\"}]}"); + UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"}]}"); } { diff --git a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp index 412116c51b7d..7e06c54a9b0b 100644 --- a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp +++ b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp @@ -518,7 +518,7 @@ class TWorkloadServiceYdbSetup : public IYdbSetup { UNIT_ASSERT_C(settings.PoolId_, "Query pool id is not specified"); auto event = std::make_unique(); - event->Record.SetUserToken(NACLib::TUserToken("", settings.UserSID_, {}).SerializeAsString()); + event->Record.SetUserToken(NACLib::TUserToken("", settings.UserSID_, settings.GroupSIDs_).SerializeAsString()); auto request = event->Record.MutableRequest(); request->SetQuery(query); diff --git a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h index 037cd3e5373e..c38e61e45b76 100644 --- a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h +++ b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h @@ -26,6 +26,7 @@ struct TQueryRunnerSettings { FLUENT_SETTING_DEFAULT(ui32, NodeIndex, 0); FLUENT_SETTING_DEFAULT(std::optional, PoolId, std::nullopt); FLUENT_SETTING_DEFAULT(TString, UserSID, "user@" BUILTIN_SYSTEM_DOMAIN); + FLUENT_SETTING_DEFAULT(TVector, GroupSIDs, {}); FLUENT_SETTING_DEFAULT(TString, Database, ""); // Runner settings diff --git a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp index efdc3aa8cdfb..6e426a543d3f 100644 --- a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp +++ b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp @@ -597,8 +597,9 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { return result.GetStatus() == EStatus::GENERIC_ERROR && errorString.Contains("You don't have access permissions for database Root"); }); - auto createResult = ydb->ExecuteQuery(R"( + auto createResult = ydb->ExecuteQuery(TStringBuilder() << R"( CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RESOURCE_POOL=")" << NResourcePool::DEFAULT_POOL_ID << R"(", RANK=20 ); )", settings); @@ -614,18 +615,22 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { UNIT_ASSERT_STRING_CONTAINS(alterResult.GetIssues().ToOneLineString(), "You don't have access permissions for database Root"); } - TString CreateSampleResourcePoolClassifier(TIntrusivePtr ydb, const TQueryRunnerSettings& settings, const TString& poolId) { - const TString& classifierId = "my_pool_classifier"; + void CreateSampleResourcePoolClassifier(TIntrusivePtr ydb, const TString& classifierId, const TString& userSID, const TString& poolId) { ydb->ExecuteSchemeQuery(TStringBuilder() << R"( + GRANT ALL ON `/Root` TO `)" << userSID << R"(`; CREATE RESOURCE POOL )" << poolId << R"( WITH ( CONCURRENT_QUERY_LIMIT=0 ); CREATE RESOURCE POOL CLASSIFIER )" << classifierId << R"( WITH ( RESOURCE_POOL=")" << poolId << R"(", - MEMBERNAME=")" << settings.UserSID_ << R"(" + MEMBERNAME=")" << userSID << R"(" ); )"); + } + TString CreateSampleResourcePoolClassifier(TIntrusivePtr ydb, const TQueryRunnerSettings& settings, const TString& poolId) { + const TString& classifierId = "my_pool_classifier"; + CreateSampleResourcePoolClassifier(ydb, classifierId, settings.UserSID_, poolId); return classifierId; } @@ -650,7 +655,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { Y_UNIT_TEST(TestCreateResourcePoolClassifier) { auto ydb = TYdbSetupSettings().Create(); - auto settings = TQueryRunnerSettings().PoolId(""); + auto settings = TQueryRunnerSettings().PoolId("").UserSID("test@user"); const TString& poolId = "my_pool"; CreateSampleResourcePoolClassifier(ydb, settings, poolId); @@ -660,15 +665,15 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { Y_UNIT_TEST(TestAlterResourcePoolClassifier) { auto ydb = TYdbSetupSettings().Create(); - auto settings = TQueryRunnerSettings().PoolId(""); + auto settings = TQueryRunnerSettings().PoolId("").UserSID("test@user"); const TString& poolId = "my_pool"; const TString& classifierId = CreateSampleResourcePoolClassifier(ydb, settings, poolId); WaitForFail(ydb, settings, poolId); ydb->ExecuteSchemeQuery(TStringBuilder() << R"( - ALTER RESOURCE POOL CLASSIFIER )" << classifierId << R"( RESET ( - RESOURCE_POOL + ALTER RESOURCE POOL CLASSIFIER )" << classifierId << R"( SET ( + RESOURCE_POOL=")" << NResourcePool::DEFAULT_POOL_ID << R"(" ); )"); @@ -678,7 +683,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { Y_UNIT_TEST(TestDropResourcePoolClassifier) { auto ydb = TYdbSetupSettings().Create(); - auto settings = TQueryRunnerSettings().PoolId(""); + auto settings = TQueryRunnerSettings().PoolId("").UserSID("test@user"); const TString& poolId = "my_pool"; const TString& classifierId = CreateSampleResourcePoolClassifier(ydb, settings, poolId); @@ -694,7 +699,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { Y_UNIT_TEST(TestDropResourcePool) { auto ydb = TYdbSetupSettings().Create(); - auto settings = TQueryRunnerSettings().PoolId(""); + auto settings = TQueryRunnerSettings().PoolId("").UserSID("test@user"); const TString& poolId = "my_pool"; CreateSampleResourcePoolClassifier(ydb, settings, poolId); @@ -710,7 +715,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { Y_UNIT_TEST(TestResourcePoolClassifierRanks) { auto ydb = TYdbSetupSettings().Create(); - auto settings = TQueryRunnerSettings().PoolId(""); + auto settings = TQueryRunnerSettings().PoolId("").UserSID("test@user"); const TString& poolId = "my_pool"; CreateSampleResourcePoolClassifier(ydb, settings, poolId); @@ -720,6 +725,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { ydb->ExecuteSchemeQuery(TStringBuilder() << R"( CREATE RESOURCE POOL CLASSIFIER )" << classifierId << R"( WITH ( RANK="1", + RESOURCE_POOL=")" << NResourcePool::DEFAULT_POOL_ID << R"(", MEMBERNAME=")" << settings.UserSID_ << R"(" ); )"); @@ -738,13 +744,50 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { Y_UNIT_TEST(TestExplicitPoolId) { auto ydb = TYdbSetupSettings().Create(); - auto settings = TQueryRunnerSettings().PoolId(""); + auto settings = TQueryRunnerSettings().PoolId("").UserSID("test@user"); const TString& poolId = "my_pool"; CreateSampleResourcePoolClassifier(ydb, settings, poolId); WaitForFail(ydb, settings, poolId); TSampleQueries::TSelect42::CheckResult(ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, TQueryRunnerSettings().PoolId(NResourcePool::DEFAULT_POOL_ID))); } + + Y_UNIT_TEST(TestMultyGrupClassification) { + auto ydb = TYdbSetupSettings().Create(); + + auto settings = TQueryRunnerSettings().PoolId(""); + + const TString& poolId = "my_pool"; + const TString& firstSID = "first@user"; + const TString& secondSID = "second@user"; + ydb->ExecuteSchemeQuery(TStringBuilder() << R"( + CREATE RESOURCE POOL )" << poolId << R"( WITH ( + CONCURRENT_QUERY_LIMIT=0 + ); + CREATE RESOURCE POOL CLASSIFIER first_classifier WITH ( + RESOURCE_POOL=")" << poolId << R"(", + MEMBERNAME=")" << firstSID << R"(", + RANK=1 + ); + CREATE RESOURCE POOL CLASSIFIER second_classifier WITH ( + RESOURCE_POOL=")" << NResourcePool::DEFAULT_POOL_ID << R"(", + MEMBERNAME=")" << secondSID << R"(", + RANK=2 + ); + )"); + + WaitForFail(ydb, settings.GroupSIDs({firstSID}), poolId); + WaitForSuccess(ydb, settings.GroupSIDs({secondSID})); + WaitForFail(ydb, settings.GroupSIDs({firstSID, secondSID}), poolId); + + ydb->ExecuteSchemeQuery(TStringBuilder() << R"( + ALTER RESOURCE POOL CLASSIFIER second_classifier SET ( + RANK=0 + ); + )"); + + WaitForSuccess(ydb, settings.GroupSIDs({firstSID, secondSID})); + } } } // namespace NKikimr::NKqp diff --git a/ydb/core/resource_pools/resource_pool_classifier_settings.cpp b/ydb/core/resource_pools/resource_pool_classifier_settings.cpp index e8dd673c4c17..34ce0e30e0ed 100644 --- a/ydb/core/resource_pools/resource_pool_classifier_settings.cpp +++ b/ydb/core/resource_pools/resource_pool_classifier_settings.cpp @@ -1,5 +1,7 @@ #include "resource_pool_classifier_settings.h" +#include + namespace NKikimr::NResourcePool { @@ -37,4 +39,11 @@ std::unordered_map TClassifierSettings: return properties; } +void TClassifierSettings::Validate() const { + NACLib::TUserToken token(Membername, TVector{}); + if (token.IsSystemUser()) { + throw yexception() << "Invalid resource pool classifier configuration, cannot create classifier for system user " << Membername; + } +} + } // namespace NKikimr::NResourcePool diff --git a/ydb/core/resource_pools/resource_pool_classifier_settings.h b/ydb/core/resource_pools/resource_pool_classifier_settings.h index e65cd2f0e7de..2decfc1a581a 100644 --- a/ydb/core/resource_pools/resource_pool_classifier_settings.h +++ b/ydb/core/resource_pools/resource_pool_classifier_settings.h @@ -25,7 +25,9 @@ struct TClassifierSettings : public TSettingsBase { }; bool operator==(const TClassifierSettings& other) const = default; + std::unordered_map GetPropertiesMap(); + void Validate() const; i64 Rank = -1; // -1 = max rank + CLASSIFIER_RANK_OFFSET TString ResourcePool = DEFAULT_POOL_ID; diff --git a/ydb/core/resource_pools/resource_pool_classifier_settings_ut.cpp b/ydb/core/resource_pools/resource_pool_classifier_settings_ut.cpp index e110f1d6b587..a91933a96603 100644 --- a/ydb/core/resource_pools/resource_pool_classifier_settings_ut.cpp +++ b/ydb/core/resource_pools/resource_pool_classifier_settings_ut.cpp @@ -2,6 +2,8 @@ #include +#include + namespace NKikimr { @@ -47,6 +49,12 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifierTest) { UNIT_ASSERT_VALUES_EQUAL(std::visit(extractor, propertiesMap["resource_pool"]), "test_pool"); UNIT_ASSERT_VALUES_EQUAL(std::visit(extractor, propertiesMap["membername"]), "test@user"); } + + Y_UNIT_TEST(SettingsValidation) { + TClassifierSettings settings; + settings.Membername = BUILTIN_ACL_METADATA; + UNIT_ASSERT_EXCEPTION_CONTAINS(settings.Validate(), yexception, TStringBuilder() << "Invalid resource pool classifier configuration, cannot create classifier for system user " << settings.Membername); + } } } // namespace NKikimr diff --git a/ydb/core/resource_pools/ya.make b/ydb/core/resource_pools/ya.make index 8238fde6d0f6..5fd79c0bf4e4 100644 --- a/ydb/core/resource_pools/ya.make +++ b/ydb/core/resource_pools/ya.make @@ -8,6 +8,7 @@ SRCS( PEERDIR( contrib/libs/protobuf util + ydb/library/aclib ) END() From 49151bca4d3a7a58b9dce2ad91b731b9b886a180 Mon Sep 17 00:00:00 2001 From: Grigoriy Pisarenko Date: Tue, 10 Sep 2024 18:30:42 +0000 Subject: [PATCH 2/2] Fixed typo --- ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp index 6e426a543d3f..cad50ad143da 100644 --- a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp +++ b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp @@ -752,7 +752,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { TSampleQueries::TSelect42::CheckResult(ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, TQueryRunnerSettings().PoolId(NResourcePool::DEFAULT_POOL_ID))); } - Y_UNIT_TEST(TestMultyGrupClassification) { + Y_UNIT_TEST(TestMultiGroupClassification) { auto ydb = TYdbSetupSettings().Create(); auto settings = TQueryRunnerSettings().PoolId("");