Events in YQ and DQ for rate limiting

Fix build

Events in DQ

Events in YQ

9 files changed, 101 insertions, 0 deletions

diff --git a/CMakeLists.darwin.txt b/CMakeLists.darwin.txt
index 4a495cabf16..e8822457113 100644
--- a/CMakeLists.darwin.txt
+++ b/CMakeLists.darwin.txt
@@ -1308,6 +1308,7 @@ add_subdirectory(ydb/core/yq/libs/checkpointing/events)
 add_subdirectory(ydb/core/yq/libs/common/ut)
 add_subdirectory(ydb/core/yq/libs/control_plane_proxy/ut)
 add_subdirectory(ydb/core/yq/libs/hmac/ut)
+add_subdirectory(ydb/core/yq/libs/rate_limiter/events)
 add_subdirectory(ydb/core/yq/libs/result_formatter/ut)
 add_subdirectory(ydb/core/yq/libs/signer/ut)
 add_subdirectory(ydb/core/yq/libs/test_connection/ut)
diff --git a/CMakeLists.linux.txt b/CMakeLists.linux.txt
index 1e8c6079cc3..5e3a848703b 100644
--- a/CMakeLists.linux.txt
+++ b/CMakeLists.linux.txt
@@ -1330,6 +1330,7 @@ add_subdirectory(ydb/core/yq/libs/checkpointing/events)
 add_subdirectory(ydb/core/yq/libs/common/ut)
 add_subdirectory(ydb/core/yq/libs/control_plane_proxy/ut)
 add_subdirectory(ydb/core/yq/libs/hmac/ut)
+add_subdirectory(ydb/core/yq/libs/rate_limiter/events)
 add_subdirectory(ydb/core/yq/libs/result_formatter/ut)
 add_subdirectory(ydb/core/yq/libs/signer/ut)
 add_subdirectory(ydb/core/yq/libs/test_connection/ut)
diff --git a/ydb/core/yq/libs/events/event_subspace.h b/ydb/core/yq/libs/events/event_subspace.h
index 46aef906293..e283d642539 100644
--- a/ydb/core/yq/libs/events/event_subspace.h
+++ b/ydb/core/yq/libs/events/event_subspace.h
@@ -28,6 +28,7 @@ struct TYqEventSubspace {
         TestConnection,
         InternalService,
         QuotaService,
+        RateLimiter,
 
         SubspacesEnd,
     };
diff --git a/ydb/core/yq/libs/protos/yq_private.proto b/ydb/core/yq/libs/protos/yq_private.proto
index a67c0db9ab1..92ef72ce935 100644
--- a/ydb/core/yq/libs/protos/yq_private.proto
+++ b/ydb/core/yq/libs/protos/yq_private.proto
@@ -77,6 +77,9 @@ message GetTaskResult {
         google.protobuf.Timestamp deadline = 24;
         YandexQuery.StreamingDisposition disposition = 25;
         uint64 result_limit = 26;
+
+        YandexQuery.Limits limits = 27;
+        string rate_limiter = 28; // Kesus path // If empty, rate limiting is off.
     }
     repeated Task tasks = 1;
 }
@@ -85,6 +88,12 @@ message GetTaskResponse {
     Ydb.Operations.Operation operation = 1; // GetTaskResult
 }
 
+message RateLimiterResources {
+    string rate_limiter = 1; // Kesus path
+    string resource_path = 2;
+    bool delete_resource = 3; // Delete resource when query is finished. Otherwise create
+}
+
 message PingTaskRequest {
     string owner_id = 1;
     SignedIdentity query_id = 2;
@@ -107,6 +116,7 @@ message PingTaskRequest {
     YandexQuery.StateLoadMode state_load_mode = 18;
     YandexQuery.StreamingDisposition disposition = 19;
     Ydb.Operations.OperationParams operation_params = 15;
+    RateLimiterResources rate_limiter_resources = 22;
     string scope = 100;
     string tenant = 104;
     google.protobuf.Timestamp started_at = 101;
diff --git a/ydb/core/yq/libs/rate_limiter/events/CMakeLists.txt b/ydb/core/yq/libs/rate_limiter/events/CMakeLists.txt
new file mode 100644
index 00000000000..010549509ff
--- /dev/null
+++ b/ydb/core/yq/libs/rate_limiter/events/CMakeLists.txt
@@ -0,0 +1,19 @@
+
+# This file was gererated by the build system used internally in the Yandex monorepo.
+# Only simple modifications are allowed (adding source-files to targets, adding simple properties
+# like target_include_directories). These modifications will be ported to original
+# ya.make files by maintainers. Any complex modifications which can't be ported back to the
+# original buildsystem will not be accepted.
+
+
+
+add_library(libs-rate_limiter-events)
+target_link_libraries(libs-rate_limiter-events PUBLIC
+  contrib-libs-cxxsupp
+  yutil
+  cpp-actors-core
+  yq-libs-events
+)
+target_sources(libs-rate_limiter-events PRIVATE
+  ${CMAKE_SOURCE_DIR}/ydb/core/yq/libs/rate_limiter/events/events.cpp
+)
diff --git a/ydb/core/yq/libs/rate_limiter/events/events.cpp b/ydb/core/yq/libs/rate_limiter/events/events.cpp
new file mode 100644
index 00000000000..6c3d2603e7e
--- /dev/null
+++ b/ydb/core/yq/libs/rate_limiter/events/events.cpp
@@ -0,0 +1 @@
+#include "events.h"
diff --git a/ydb/core/yq/libs/rate_limiter/events/events.h b/ydb/core/yq/libs/rate_limiter/events/events.h
new file mode 100644
index 00000000000..7d2c360f594
--- /dev/null
+++ b/ydb/core/yq/libs/rate_limiter/events/events.h
@@ -0,0 +1,64 @@
+#pragma once
+#include <ydb/core/yq/libs/events/event_subspace.h>
+#include <ydb/library/yql/public/issue/yql_issue.h>
+
+#include <library/cpp/actors/core/event_local.h>
+#include <library/cpp/actors/core/events.h>
+
+namespace NYq {
+
+struct TEvRateLimiter {
+    // Event ids.
+    enum EEv : ui32 {
+        EvGetRateLimiterPath = YqEventSubspaceBegin(NYq::TYqEventSubspace::RateLimiter),
+        EvRateLimiterPath,
+        EvCreateResource,
+        EvCreateResourceResponse,
+        EvEnd,
+    };
+
+    static_assert(EvEnd <= YqEventSubspaceEnd(NYq::TYqEventSubspace::RateLimiter), "All events must be in their subspace");
+
+private:
+    struct TEvGetRateLimiterPath : NActors::TEventLocal<TEvGetRateLimiterPath, EEv::EvGetRateLimiterPath> {
+        explicit TEvGetRateLimiterPath(const TString& cloudId)
+            : CloudId(cloudId)
+        {
+        }
+
+        const TString CloudId;
+    };
+
+    struct TEvRateLimiterPath : NActors::TEventLocal<TEvRateLimiterPath, EEv::EvRateLimiterPath> {
+        explicit TEvRateLimiterPath(const TString& path)
+            : Path(path)
+        {
+        }
+
+        const TString Path;
+    };
+
+    struct TEvCreateResource : NActors::TEventLocal<TEvCreateResource, EEv::EvCreateResource> {
+        explicit TEvCreateResource(const TString& rateLimiterPath, const TString& resourcePath)
+            : RateLimiterPath(rateLimiterPath)
+            , ResourcePath(resourcePath)
+        {
+        }
+
+        const TString RateLimiterPath;
+        const TString ResourcePath;
+    };
+
+    struct TEvCreateResourceResponse : NActors::TEventLocal<TEvCreateResourceResponse, EEv::EvCreateResourceResponse> {
+        explicit TEvCreateResourceResponse(bool success, const NYql::TIssues& issues = {})
+            : Success(success)
+            , Issues(issues)
+        {
+        }
+
+        const bool Success;
+        const NYql::TIssues Issues;
+    };
+};
+
+} // namespace NYq
diff --git a/ydb/library/yql/dq/proto/dq_tasks.proto b/ydb/library/yql/dq/proto/dq_tasks.proto
index 5ae888c1143..5ac9c3fe850 100644
--- a/ydb/library/yql/dq/proto/dq_tasks.proto
+++ b/ydb/library/yql/dq/proto/dq_tasks.proto
@@ -151,4 +151,6 @@ message TDqTask {
     repeated TTaskOutput Outputs = 6;
     google.protobuf.Any Meta = 7;
     bool CreateSuspended = 8;
+    string RateLimiter = 10;
+    string RateLimiterResource = 11;
 }
diff --git a/ydb/library/yql/providers/dq/api/protos/service.proto b/ydb/library/yql/providers/dq/api/protos/service.proto
index b91351e5d12..c3d001ff954 100644
--- a/ydb/library/yql/providers/dq/api/protos/service.proto
+++ b/ydb/library/yql/providers/dq/api/protos/service.proto
@@ -81,6 +81,8 @@ message ExecuteGraphRequest {
     map<string, string> SecureParams = 9;
     bool Discard = 12;
     map<string, string> GraphParams = 13;
+    string RateLimiter = 14;
+    string RateLimiterResource = 15;
 }
 
 message ExecuteGraphResponse {