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

Author: GrigoriyPA

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);

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;
 };
 

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);
         }
 

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);
 }

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;