From f3de9bd606519a9274acaf50aaf3b4e5ddfb725a Mon Sep 17 00:00:00 2001
From: Polina Volosnikova <vpolka@ydb.tech>
Date: Tue, 28 May 2024 09:43:47 +0000
Subject: [PATCH] fix test

---
 ydb/core/kqp/session_actor/kqp_query_state.cpp   | 9 ++++++---
 ydb/core/kqp/session_actor/kqp_query_state.h     | 1 +
 ydb/core/kqp/session_actor/kqp_session_actor.cpp | 2 +-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/ydb/core/kqp/session_actor/kqp_query_state.cpp b/ydb/core/kqp/session_actor/kqp_query_state.cpp
index 4f5955265083..d4bb4b2cd5da 100644
--- a/ydb/core/kqp/session_actor/kqp_query_state.cpp
+++ b/ydb/core/kqp/session_actor/kqp_query_state.cpp
@@ -323,6 +323,10 @@ std::unique_ptr<TEvKqp::TEvCompileRequest> TKqpQueryState::BuildCompileSplittedR
         false, SplittedCtx.Get(), SplittedExprs.at(NextSplittedExpr));
 }
 
+bool TKqpQueryState::ProcessingLastStatementPart() {
+    return SplittedExprs.empty() || (NextSplittedExpr + 1 >= static_cast<int>(SplittedExprs.size()));
+}
+
 bool TKqpQueryState::PrepareNextStatementPart() {
     QueryData = {};
     PreparedQuery = {};
@@ -335,9 +339,7 @@ bool TKqpQueryState::PrepareNextStatementPart() {
     TopicOperations = {};
     ReplayMessage = {};
 
-    ++NextSplittedExpr;
-    
-    if (NextSplittedExpr >= static_cast<int>(SplittedExprs.size()) || SplittedExprs.empty()) {
+    if (ProcessingLastStatementPart()) {
         SplittedWorld.Reset();
         SplittedExprs.clear();
         SplittedCtx.Reset();
@@ -345,6 +347,7 @@ bool TKqpQueryState::PrepareNextStatementPart() {
         return false;
     }
 
+    ++NextSplittedExpr;
     return true;
 }
 
diff --git a/ydb/core/kqp/session_actor/kqp_query_state.h b/ydb/core/kqp/session_actor/kqp_query_state.h
index ab7f8792295b..7cf13a2a14ee 100644
--- a/ydb/core/kqp/session_actor/kqp_query_state.h
+++ b/ydb/core/kqp/session_actor/kqp_query_state.h
@@ -481,6 +481,7 @@ class TKqpQueryState : public TNonCopyable {
     std::unique_ptr<TEvKqp::TEvCompileRequest> BuildSplitRequest(std::shared_ptr<std::atomic<bool>> cookie, const TGUCSettings::TPtr& gUCSettingsPtr);
     std::unique_ptr<TEvKqp::TEvCompileRequest> BuildCompileSplittedRequest(std::shared_ptr<std::atomic<bool>> cookie, const TGUCSettings::TPtr& gUCSettingsPtr);
 
+    bool ProcessingLastStatementPart();
     bool PrepareNextStatementPart();
 
     const ::google::protobuf::Map<TProtoStringType, ::Ydb::TypedValue>& GetYdbParameters() const {
diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
index 8fd99af4948a..d99c6416ccaf 100644
--- a/ydb/core/kqp/session_actor/kqp_session_actor.cpp
+++ b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -725,7 +725,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             QueryState->ImplicitTxId = QueryState->TxId;
         }
         Ydb::Table::TransactionControl control;
-        control.set_commit_tx(QueryState->ProcessingLastStatement());
+        control.set_commit_tx(QueryState->ProcessingLastStatement() && QueryState->ProcessingLastStatementPart());
         control.set_tx_id(QueryState->ImplicitTxId->GetValue().GetHumanStr());
         return control;
     }