Fix hang of ExecuteScript method

2 files changed, 24 insertions, 8 deletions

diff --git a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
index 83939ba124..d199206af0 100644
--- a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
+++ b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
@@ -537,15 +537,17 @@ struct TCreateScriptExecutionActor : public TActorBootstrapped<TCreateScriptExec
         Become(&TCreateScriptExecutionActor::StateFunc);
 
         // Start request
-        const NActors::TActorId actorId = Register(CreateRunScriptActor(Event->Get()->Record, Event->Get()->Record.GetRequest().GetDatabase(), 1));
-        TString executionId = actorId.ToString();
+        RunScriptActorId = Register(CreateRunScriptActor(Event->Get()->Record, Event->Get()->Record.GetRequest().GetDatabase(), 1));
+        TString executionId = RunScriptActorId.ToString();
         Register(new TCreateScriptOperationQuery(executionId, Event->Get()->Record));
     }
 
     void Handle(TEvPrivate::TEvCreateScriptOperationResponse::TPtr& ev) {
         if (ev->Get()->Status == Ydb::StatusIds::SUCCESS) {
+            Send(RunScriptActorId, new NActors::TEvents::TEvWakeup());
             Send(Event->Sender, new TEvKqp::TEvScriptResponse(ScriptExecutionOperationFromExecutionId(ev->Get()->ExecutionId), ev->Get()->ExecutionId, Ydb::Query::EXEC_STATUS_STARTING, Ydb::Query::EXEC_MODE_EXECUTE));
         } else {
+            Send(RunScriptActorId, new NActors::TEvents::TEvPoison());
             Send(Event->Sender, new TEvKqp::TEvScriptResponse(ev->Get()->Status, std::move(ev->Get()->Issues)));
         }
     }
@@ -556,6 +558,7 @@ struct TCreateScriptExecutionActor : public TActorBootstrapped<TCreateScriptExec
 
 private:
     TEvKqp::TEvScriptRequest::TPtr Event;
+    NActors::TActorId RunScriptActorId;
 };
 
 class TScriptExecutionFinisher : public TQueryBase {
diff --git a/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp b/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp
index 5288e6b67a..f58613524f 100644
--- a/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp
+++ b/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp
@@ -39,7 +39,18 @@ public:
 
     void Bootstrap() {
         Become(&TRunScriptActor::StateFunc);
+    }
 
+private:
+    STRICT_STFUNC(StateFunc,
+        hFunc(NActors::TEvents::TEvWakeup, Handle);
+        hFunc(NActors::TEvents::TEvPoison, Handle);
+        hFunc(NKqp::TEvKqpExecuter::TEvStreamData, Handle);
+        hFunc(NKqp::TEvKqp::TEvQueryResponse, Handle);
+        hFunc(NKqp::TEvKqp::TEvFetchScriptResultsRequest, Handle);
+    )
+
+    void Start() {
         auto ev = MakeHolder<NKqp::TEvKqp::TEvQueryRequest>();
         ev->Record = Request;
 
@@ -51,12 +62,14 @@ public:
         }
     }
 
-private:
-    STRICT_STFUNC(StateFunc,
-        hFunc(NKqp::TEvKqpExecuter::TEvStreamData, Handle);
-        hFunc(NKqp::TEvKqp::TEvQueryResponse, Handle);
-        hFunc(NKqp::TEvKqp::TEvFetchScriptResultsRequest, Handle);
-    )
+    // TODO: remove this after there will be a normal way to store results and generate execution id
+    void Handle(NActors::TEvents::TEvWakeup::TPtr&) {
+        Start();
+    }
+
+    void Handle(NActors::TEvents::TEvPoison::TPtr&) {
+        PassAway();
+    }
 
     void Handle(NKqp::TEvKqpExecuter::TEvStreamData::TPtr& ev) {
         auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();