Skip to content

Commit a3d73ea

Browse files
authored
WM passed resource pool settings to user request context (#5734)
1 parent 5c18e12 commit a3d73ea

18 files changed

+189
-145
lines changed

ydb/core/kqp/common/events/workload_service.h

+4-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#pragma once
22

33
#include <ydb/core/kqp/common/simple/kqp_event_ids.h>
4+
#include <ydb/core/resource_pools/resource_pool_settings.h>
45

56
#include <ydb/library/aclib/aclib.h>
67
#include <ydb/library/actors/core/event_local.h>
@@ -24,12 +25,14 @@ struct TEvPlaceRequestIntoPool : public NActors::TEventLocal<TEvPlaceRequestInto
2425
};
2526

2627
struct TEvContinueRequest : public NActors::TEventLocal<TEvContinueRequest, TKqpWorkloadServiceEvents::EvContinueRequest> {
27-
explicit TEvContinueRequest(Ydb::StatusIds::StatusCode status, NYql::TIssues issues = {})
28+
TEvContinueRequest(Ydb::StatusIds::StatusCode status, const NResourcePool::TPoolSettings& poolConfig, NYql::TIssues issues = {})
2829
: Status(status)
30+
, PoolConfig(poolConfig)
2931
, Issues(std::move(issues))
3032
{}
3133

3234
const Ydb::StatusIds::StatusCode Status;
35+
const NResourcePool::TPoolSettings PoolConfig;
3336
const NYql::TIssues Issues;
3437
};
3538

ydb/core/kqp/common/events/ya.make

+1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ PEERDIR(
1515
ydb/core/grpc_services/cancelation
1616
ydb/core/kqp/common/shutdown
1717
ydb/core/kqp/common/compilation
18+
ydb/core/resource_pools
1819

1920
ydb/library/yql/dq/actors
2021
ydb/public/api/protos

ydb/core/kqp/common/kqp_user_request_context.h

+3
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
#include <util/generic/fwd.h>
55
#include <contrib/libs/protobuf/src/google/protobuf/map.h>
66

7+
#include <ydb/core/resource_pools/resource_pool_settings.h>
8+
79
namespace NKikimr::NKqp {
810

911
struct TUserRequestContext : public TAtomicRefCount<TUserRequestContext> {
@@ -13,6 +15,7 @@ namespace NKikimr::NKqp {
1315
TString CurrentExecutionId;
1416
TString CustomerSuppliedId;
1517
TString PoolId;
18+
NResourcePool::TPoolSettings PoolConfig;
1619

1720
TUserRequestContext() = default;
1821

ydb/core/kqp/session_actor/kqp_session_actor.cpp

+1
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
481481
}
482482

483483
LOG_D("continue request, pool id: " << poolId);
484+
QueryState->UserRequestContext->PoolConfig = ev->Get()->PoolConfig;
484485
CompileQuery();
485486
}
486487

ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp

+2-1
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
#include <ydb/core/kqp/ut/common/kqp_ut_common.h>
33
#include <ydb/core/kqp/ut/common/columnshard.h>
44
#include <ydb/core/kqp/workload_service/kqp_workload_service.h>
5+
#include <ydb/core/resource_pools/resource_pool_settings.h>
56
#include <ydb/core/testlib/common_helper.h>
67
#include <ydb/public/lib/ut_helpers/ut_helpers_query.h>
78
#include <ydb/public/sdk/cpp/client/ydb_operation/operation.h>
@@ -238,7 +239,7 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
238239

239240
Y_UNIT_TEST(ExecuteQueryWithWorkloadManager) {
240241
NWorkload::TWorkloadManagerConfig workloadManagerConfig;
241-
workloadManagerConfig.Pools.insert({"sample_pool_id", NWorkload::TWorkloadManagerConfig::TPoolConfig()});
242+
workloadManagerConfig.Pools.insert({"sample_pool_id", NResourcePool::TPoolSettings()});
242243
SetWorkloadManagerConfig(workloadManagerConfig);
243244

244245
NKikimrConfig::TAppConfig config;

ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include <ydb/core/kqp/counters/kqp_counters.h>
22
#include <ydb/core/kqp/ut/common/kqp_ut_common.h>
33
#include <ydb/core/kqp/workload_service/kqp_workload_service.h>
4+
#include <ydb/core/resource_pools/resource_pool_settings.h>
45
#include <ydb/public/lib/ut_helpers/ut_helpers_query.h>
56
#include <ydb/public/sdk/cpp/client/ydb_operation/operation.h>
67
#include <ydb/public/sdk/cpp/client/ydb_proto/accessor.h>
@@ -101,7 +102,7 @@ Y_UNIT_TEST_SUITE(KqpQueryServiceScripts) {
101102

102103
Y_UNIT_TEST(ExecuteScriptWithWorkloadManager) {
103104
NWorkload::TWorkloadManagerConfig workloadManagerConfig;
104-
workloadManagerConfig.Pools.insert({"sample_pool_id", NWorkload::TWorkloadManagerConfig::TPoolConfig()});
105+
workloadManagerConfig.Pools.insert({"sample_pool_id", NResourcePool::TPoolSettings()});
105106
SetWorkloadManagerConfig(workloadManagerConfig);
106107

107108
NKikimrConfig::TAppConfig config;

ydb/core/kqp/workload_service/kqp_workload_service.cpp

+1-6
Original file line numberDiff line numberDiff line change
@@ -118,11 +118,6 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
118118
return;
119119
}
120120

121-
if (!poolState->HasAccess(ev->Get()->UserToken)) {
122-
ReplyContinueError(workerActorId, Ydb::StatusIds::UNAUTHORIZED, TStringBuilder() << "You do not have access permissions for pool " << poolId);
123-
return;
124-
}
125-
126121
if (poolState->PlaceRequest(workerActorId, ev->Get()->SessionId) && poolState->TablesRequired()) {
127122
ScheduleLeaseUpdate();
128123
PrepareWorkloadServiceTables();
@@ -258,7 +253,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
258253

259254
void ReplyContinueError(const TActorId& replyActorId, Ydb::StatusIds::StatusCode status, const TString& message) const {
260255
LOG_W("Reply continue error " << status << " to " << replyActorId << ": " << message);
261-
Send(replyActorId, new TEvContinueRequest(status, {NYql::TIssue(message)}));
256+
Send(replyActorId, new TEvContinueRequest(status, {}, {NYql::TIssue(message)}));
262257
}
263258

264259
void ReplyCleanupError(const TActorId& replyActorId, Ydb::StatusIds::StatusCode status, const TString& message) const {

ydb/core/kqp/workload_service/kqp_workload_service.h

+2-9
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#pragma once
22

3+
#include <ydb/core/resource_pools/resource_pool_settings.h>
34
#include <ydb/library/actors/core/actor.h>
45

56

@@ -8,15 +9,7 @@ namespace NKikimr::NKqp {
89
namespace NWorkload {
910

1011
struct TWorkloadManagerConfig {
11-
struct TPoolConfig {
12-
ui64 ConcurrentQueryLimit = 0; // 0 = infinity
13-
ui64 QueryCountLimit = 0; // 0 = infinity
14-
TDuration QueryCancelAfter = TDuration::Days(1);
15-
16-
TString ACL = ""; // empty = full access for all users
17-
};
18-
19-
std::unordered_map<TString, TPoolConfig> Pools;
12+
std::unordered_map<TString, NResourcePool::TPoolSettings> Pools;
2013
};
2114

2215
void SetWorkloadManagerConfig(const TWorkloadManagerConfig& workloadManagerConfig);

ydb/core/kqp/workload_service/kqp_workload_service_impl.h

+2-4
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
#pragma once
22

3-
#include "kqp_workload_service.h"
43
#include "kqp_workload_service_tables_impl.h"
54

65
#include <ydb/core/kqp/common/events/events.h>
6+
#include <ydb/core/resource_pools/resource_pool_settings.h>
77

8-
#include <ydb/library/aclib/aclib.h>
98
#include <ydb/library/actors/core/actor.h>
109

1110

@@ -18,7 +17,6 @@ namespace NQueue {
1817
class IState : public TThrRefBase {
1918
public:
2019
virtual bool TablesRequired() const = 0;
21-
virtual bool HasAccess(const TIntrusiveConstPtr<NACLib::TUserToken>& userToken) const = 0;
2220
virtual ui64 GetLocalPoolSize() const = 0;
2321

2422
virtual void OnPreparingFinished(Ydb::StatusIds::StatusCode status, NYql::TIssues issues) = 0;
@@ -36,7 +34,7 @@ class IState : public TThrRefBase {
3634

3735
using TStatePtr = TIntrusivePtr<IState>;
3836

39-
TStatePtr CreateState(const NActors::TActorContext& actorContext, const TString& poolId, const TWorkloadManagerConfig::TPoolConfig& poolConfig, NMonitoring::TDynamicCounterPtr counters);
37+
TStatePtr CreateState(const NActors::TActorContext& actorContext, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters);
4038

4139
} // NQueue
4240

0 commit comments

Comments
 (0)