YQ kqprun added cancel after for script requests (#8768)

5 files changed, 30 insertions, 0 deletions

diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp
index 646c7ae493f..0b8031da80b 100644
--- a/ydb/tests/tools/kqprun/kqprun.cpp
+++ b/ydb/tests/tools/kqprun/kqprun.cpp
@@ -493,6 +493,10 @@ protected:
                 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 5966d278581..980752ef65e 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 34e319ea70d..dda287b372d 100644
--- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp
+++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp
@@ -247,6 +247,15 @@ private:
                 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 6a1f1cf7570..df82b94ca9b 100644
--- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp
+++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
@@ -287,6 +287,13 @@ public:
         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 67ca5c26c7b..0bd68d29793 100644
--- a/ydb/tests/tools/kqprun/src/ydb_setup.h
+++ b/ydb/tests/tools/kqprun/src/ydb_setup.h
@@ -62,6 +62,8 @@ public:
 
     TRequestResult ForgetScriptExecutionOperationRequest(const TString& operation) const;
 
+    TRequestResult CancelScriptExecutionOperationRequest(const TString& operation) const;
+
     void QueryRequestAsync(const TRequestOptions& query) const;
 
     void WaitAsyncQueries() const;