From ff117ecdee0ba94f12c9707f89e7dc6323b76798 Mon Sep 17 00:00:00 2001
From: Hor911 <hor911@ydb.tech>
Date: Mon, 25 Mar 2024 15:31:30 +0300
Subject: [PATCH] Restore in-progress stat (in not Ready replies) (#3111)

---
 ydb/core/fq/libs/compute/ydb/events/events.h   |  4 ++--
 .../libs/compute/ydb/status_tracker_actor.cpp  |  6 ------
 .../libs/compute/ydb/ydb_connector_actor.cpp   | 18 +++++++++---------
 .../proxy_service/kqp_script_executions.cpp    |  1 +
 4 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/ydb/core/fq/libs/compute/ydb/events/events.h b/ydb/core/fq/libs/compute/ydb/events/events.h
index b893fa3d2d50..3f19becf8203 100644
--- a/ydb/core/fq/libs/compute/ydb/events/events.h
+++ b/ydb/core/fq/libs/compute/ydb/events/events.h
@@ -125,14 +125,14 @@ struct TEvYdbCompute {
             , Ready(ready)
         {}
 
-        TEvGetOperationResponse(NYdb::NQuery::EExecStatus execStatus, Ydb::StatusIds::StatusCode statusCode, const TVector<Ydb::Query::ResultSetMeta>& resultSetsMeta, const Ydb::TableStats::QueryStats& queryStats, NYql::TIssues issues)
+        TEvGetOperationResponse(NYdb::NQuery::EExecStatus execStatus, Ydb::StatusIds::StatusCode statusCode, const TVector<Ydb::Query::ResultSetMeta>& resultSetsMeta, const Ydb::TableStats::QueryStats& queryStats, NYql::TIssues issues, bool ready = true)
             : ExecStatus(execStatus)
             , StatusCode(statusCode)
             , ResultSetsMeta(resultSetsMeta)
             , QueryStats(queryStats)
             , Issues(std::move(issues))
             , Status(NYdb::EStatus::SUCCESS)
-            , Ready(true)
+            , Ready(ready)
         {}
 
         NYdb::NQuery::EExecStatus ExecStatus = NYdb::NQuery::EExecStatus::Unspecified;
diff --git a/ydb/core/fq/libs/compute/ydb/status_tracker_actor.cpp b/ydb/core/fq/libs/compute/ydb/status_tracker_actor.cpp
index 6b5c7584d03d..c27323eb5748 100644
--- a/ydb/core/fq/libs/compute/ydb/status_tracker_actor.cpp
+++ b/ydb/core/fq/libs/compute/ydb/status_tracker_actor.cpp
@@ -123,12 +123,6 @@ class TStatusTrackerActor : public TBaseComputeActor<TStatusTrackerActor> {
     void Handle(const TEvYdbCompute::TEvGetOperationResponse::TPtr& ev) {
         const auto& response = *ev.Get()->Get();
 
-        if (response.Status == NYdb::EStatus::SUCCESS && !response.Ready) {
-            LOG_D("GetOperation IS NOT READY, repeating");
-            SendGetOperation(TDuration::MilliSeconds(BackoffTimer.NextBackoffMs()));
-            return;
-        }
-
         if (response.Status == NYdb::EStatus::NOT_FOUND) { // FAILING / ABORTING_BY_USER / ABORTING_BY_SYSTEM
             LOG_I("Operation has been already removed");
             Send(Parent, new TEvYdbCompute::TEvStatusTrackerResponse(response.Issues, response.Status, ExecStatus, ComputeStatus));
diff --git a/ydb/core/fq/libs/compute/ydb/ydb_connector_actor.cpp b/ydb/core/fq/libs/compute/ydb/ydb_connector_actor.cpp
index e0d7c520014a..a576908d779f 100644
--- a/ydb/core/fq/libs/compute/ydb/ydb_connector_actor.cpp
+++ b/ydb/core/fq/libs/compute/ydb/ydb_connector_actor.cpp
@@ -86,17 +86,17 @@ class TYdbConnectorActor : public NActors::TActorBootstrapped<TYdbConnectorActor
             .Apply([actorSystem = NActors::TActivationContext::ActorSystem(), recipient = ev->Sender, cookie = ev->Cookie, database = ComputeConnection.database()](auto future) {
                 try {
                     auto response = future.ExtractValueSync();
-                    if (!response.Ready()) {
+                    if (response.Id().GetKind() != Ydb::TOperationId::UNUSED) {
                         actorSystem->Send(
-                            recipient,
-                            MakeResponse<TEvYdbCompute::TEvGetOperationResponse>(
-                                database,
-                                response.Status().GetIssues(),
-                                response.Status().GetStatus(), 
-                                false),
+                            recipient, 
+                            new TEvYdbCompute::TEvGetOperationResponse(
+                                response.Metadata().ExecStatus,
+                                static_cast<Ydb::StatusIds::StatusCode>(response.Status().GetStatus()),
+                                response.Metadata().ResultSetsMeta,
+                                response.Metadata().ExecStats,
+                                RemoveDatabaseFromIssues(response.Status().GetIssues(), database),
+                                response.Ready()),
                             0, cookie);
-                    } else if (response.Id().GetKind() != Ydb::TOperationId::UNUSED) {
-                        actorSystem->Send(recipient, new TEvYdbCompute::TEvGetOperationResponse(response.Metadata().ExecStatus, static_cast<Ydb::StatusIds::StatusCode>(response.Status().GetStatus()), response.Metadata().ResultSetsMeta, response.Metadata().ExecStats, RemoveDatabaseFromIssues(response.Status().GetIssues(), database)), 0, cookie);
                     } else {
                         actorSystem->Send(
                             recipient,
diff --git a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
index 6d9bd00f6307..ef428b975d83 100644
--- a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
+++ b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
@@ -1289,6 +1289,7 @@ class TGetScriptExecutionOperationActor : public TCheckLeaseStatusActorBase {
             Response->Get()->Ready = false;
             Response->Get()->Status = Ydb::StatusIds::SUCCESS;
             Response->Get()->Issues.Clear();
+            Response->Get()->Metadata.set_exec_status(Ydb::Query::ExecStatus::EXEC_STATUS_UNSPECIFIED);
         } else {
             Response->Get()->Ready = true;
             Response->Get()->Status = GetOperationStatus();