From c07ddf8417c28e058874d675003784c39b9cffdd Mon Sep 17 00:00:00 2001
From: Pisarenko Grigoriy <79596613+GrigoriyPA@users.noreply.github.com>
Date: Thu, 5 Sep 2024 11:28:27 +0300
Subject: [PATCH] YQ kqprun added cancel after for script requests (#8768)

---
 ydb/tests/tools/kqprun/kqprun.cpp         |  4 ++++
 ydb/tests/tools/kqprun/src/common.h       |  2 ++
 ydb/tests/tools/kqprun/src/kqp_runner.cpp |  9 +++++++++
 ydb/tests/tools/kqprun/src/ydb_setup.cpp  | 13 +++++++++++++
 ydb/tests/tools/kqprun/src/ydb_setup.h    |  2 ++
 5 files changed, 30 insertions(+)

diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp
index 646c7ae493f4..0b8031da80b0 100644
--- a/ydb/tests/tools/kqprun/kqprun.cpp
+++ b/ydb/tests/tools/kqprun/kqprun.cpp
@@ -493,6 +493,10 @@ class TMain : public TMainClassArgs {
                 ExecutionOptions.ScriptQueryActions.emplace_back(scriptAction(choice));
             });
 
+        options.AddLongOption("cancel-after", "Cancel script execution operation after specified delay in milliseconds")
+            .RequiredArgument("uint")
+            .StoreMappedResultT<ui64>(&RunnerOptions.ScriptCancelAfter, &TDuration::MilliSeconds<ui64>);
+
         options.AddLongOption('F', "forget", "Forget script execution operation after fetching results")
             .NoArgument()
             .SetFlag(&ExecutionOptions.ForgetExecution);
diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h
index 5966d2785818..980752ef65e3 100644
--- a/ydb/tests/tools/kqprun/src/common.h
+++ b/ydb/tests/tools/kqprun/src/common.h
@@ -71,6 +71,8 @@ struct TRunnerOptions {
     NYdb::NConsoleClient::EOutputFormat PlanOutputFormat = NYdb::NConsoleClient::EOutputFormat::Default;
     ETraceOptType TraceOptType = ETraceOptType::Disabled;
 
+    TDuration ScriptCancelAfter;
+
     TYdbSetupSettings YdbSettings;
 };
 
diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp
index 34e319ea70d5..dda287b372d4 100644
--- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp
+++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp
@@ -247,6 +247,15 @@ class TKqpRunner::TImpl {
                 return false;
             }
 
+            if (Options_.ScriptCancelAfter && TInstant::Now() - StartTime_ > Options_.ScriptCancelAfter) {
+                Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Cancelling script execution..." << CoutColors_.Default() << Endl;
+                TRequestResult cancelStatus = YdbSetup_.CancelScriptExecutionOperationRequest(ExecutionOperation_);
+                if (!cancelStatus.IsSuccess()) {
+                    Cerr << CerrColors_.Red() << "Failed to cancel script execution operation, reason:" << CerrColors_.Default() << Endl << cancelStatus.ToString() << Endl;
+                    return false;
+                }
+            }
+
             Sleep(getOperationPeriod);
         }
 
diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
index 6a1f1cf7570d..df82b94ca9b2 100644
--- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp
+++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
@@ -287,6 +287,13 @@ class TYdbSetup::TImpl {
         return RunKqpProxyRequest<NKikimr::NKqp::TEvForgetScriptExecutionOperation, NKikimr::NKqp::TEvForgetScriptExecutionOperationResponse>(std::move(event));
     }
 
+    NKikimr::NKqp::TEvCancelScriptExecutionOperationResponse::TPtr CancelScriptExecutionOperationRequest(const TString& operation) const {
+        NKikimr::NOperationId::TOperationId operationId(operation);
+        auto event = MakeHolder<NKikimr::NKqp::TEvCancelScriptExecutionOperation>(Settings_.DomainName, operationId);
+
+        return RunKqpProxyRequest<NKikimr::NKqp::TEvCancelScriptExecutionOperation, NKikimr::NKqp::TEvCancelScriptExecutionOperationResponse>(std::move(event));
+    }
+
     void QueryRequestAsync(const TRequestOptions& query) {
         if (!AsyncQueryRunnerActorId_) {
             AsyncQueryRunnerActorId_ = GetRuntime()->Register(CreateAsyncQueryRunnerActor(Settings_.AsyncQueriesSettings), 0, GetRuntime()->GetAppData().UserPoolId);
@@ -505,6 +512,12 @@ TRequestResult TYdbSetup::ForgetScriptExecutionOperationRequest(const TString& o
     return TRequestResult(forgetScriptExecutionOperationResponse->Get()->Status, forgetScriptExecutionOperationResponse->Get()->Issues);
 }
 
+TRequestResult TYdbSetup::CancelScriptExecutionOperationRequest(const TString& operation) const {
+    auto cancelScriptExecutionOperationResponse = Impl_->CancelScriptExecutionOperationRequest(operation);
+
+    return TRequestResult(cancelScriptExecutionOperationResponse->Get()->Status, cancelScriptExecutionOperationResponse->Get()->Issues);
+}
+
 void TYdbSetup::QueryRequestAsync(const TRequestOptions& query) const {
     Impl_->QueryRequestAsync(query);
 }
diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.h b/ydb/tests/tools/kqprun/src/ydb_setup.h
index 67ca5c26c7b5..0bd68d29793e 100644
--- a/ydb/tests/tools/kqprun/src/ydb_setup.h
+++ b/ydb/tests/tools/kqprun/src/ydb_setup.h
@@ -62,6 +62,8 @@ class TYdbSetup {
 
     TRequestResult ForgetScriptExecutionOperationRequest(const TString& operation) const;
 
+    TRequestResult CancelScriptExecutionOperationRequest(const TString& operation) const;
+
     void QueryRequestAsync(const TRequestOptions& query) const;
 
     void WaitAsyncQueries() const;