fix multiple table updates in same tx (#2919)

2 files changed, 53 insertions, 18 deletions

diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp b/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp
index a5b05aeb415..5817fc44224 100644
--- a/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp
+++ b/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp
@@ -96,7 +96,9 @@ IActor* CreateKqpExecuter(IKqpGateway::TExecPhysicalRequest&& request, const TSt
     for (auto& tx : request.Transactions) {
         if (txsType) {
             YQL_ENSURE(*txsType == tx.Body->GetType(), "Mixed physical tx types in executer.");
-            YQL_ENSURE(*txsType == NKqpProto::TKqpPhyTx::TYPE_DATA, "Cannot execute multiple non-data physical txs.");
+            YQL_ENSURE((*txsType == NKqpProto::TKqpPhyTx::TYPE_DATA)
+                || (*txsType == NKqpProto::TKqpPhyTx::TYPE_GENERIC),
+                "Cannot execute multiple non-data physical txs.");
         } else {
             txsType = tx.Body->GetType();
         }
diff --git a/ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp b/ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp
index 85974d07551..f8a6b4666f6 100644
--- a/ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp
+++ b/ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp
@@ -242,27 +242,60 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
         UNIT_ASSERT_VALUES_EQUAL_C(sessionResult.GetStatus(), EStatus::SUCCESS, sessionResult.GetIssues().ToString());
         auto session = sessionResult.GetSession();
 
-        const TString query = "UPDATE TwoShard SET Value2 = 0";
-        auto result = session.ExecuteQuery(query, TTxControl::BeginTx()).ExtractValueSync();
-        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
-        auto transaction = result.GetTransaction();
-        UNIT_ASSERT(transaction->IsActive());
+        {
+            const TString query = "UPDATE TwoShard SET Value2 = 0";
+            auto result = session.ExecuteQuery(query, TTxControl::BeginTx()).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+            auto transaction = result.GetTransaction();
+            UNIT_ASSERT(transaction->IsActive());
 
-        auto checkResult = [&](TString expected) {
-            auto selectRes = db.ExecuteQuery(
-                "SELECT * FROM TwoShard ORDER BY Key",
-                TTxControl::BeginTx().CommitTx()
-            ).ExtractValueSync();
+            auto checkResult = [&](TString expected) {
+                auto selectRes = db.ExecuteQuery(
+                    "SELECT * FROM TwoShard ORDER BY Key",
+                    TTxControl::BeginTx().CommitTx()
+                ).ExtractValueSync();
 
-            UNIT_ASSERT_C(selectRes.IsSuccess(), selectRes.GetIssues().ToString());
-            CompareYson(expected, FormatResultSetYson(selectRes.GetResultSet(0)));
-        };
-        checkResult(R"([[[1u];["One"];[-1]];[[2u];["Two"];[0]];[[3u];["Three"];[1]];[[4000000001u];["BigOne"];[-1]];[[4000000002u];["BigTwo"];[0]];[[4000000003u];["BigThree"];[1]]])");
+                UNIT_ASSERT_C(selectRes.IsSuccess(), selectRes.GetIssues().ToString());
+                CompareYson(expected, FormatResultSetYson(selectRes.GetResultSet(0)));
+            };
+            checkResult(R"([[[1u];["One"];[-1]];[[2u];["Two"];[0]];[[3u];["Three"];[1]];[[4000000001u];["BigOne"];[-1]];[[4000000002u];["BigTwo"];[0]];[[4000000003u];["BigThree"];[1]]])");
 
-        auto txRes = transaction->Commit().GetValueSync();
-        UNIT_ASSERT_VALUES_EQUAL_C(txRes.GetStatus(), EStatus::SUCCESS, txRes.GetIssues().ToString());
+            auto txRes = transaction->Commit().GetValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(txRes.GetStatus(), EStatus::SUCCESS, txRes.GetIssues().ToString());
 
-        checkResult(R"([[[1u];["One"];[0]];[[2u];["Two"];[0]];[[3u];["Three"];[0]];[[4000000001u];["BigOne"];[0]];[[4000000002u];["BigTwo"];[0]];[[4000000003u];["BigThree"];[0]]])");
+            checkResult(R"([[[1u];["One"];[0]];[[2u];["Two"];[0]];[[3u];["Three"];[0]];[[4000000001u];["BigOne"];[0]];[[4000000002u];["BigTwo"];[0]];[[4000000003u];["BigThree"];[0]]])");
+        }
+
+        {
+            const TString query = "UPDATE TwoShard SET Value2 = 1";
+            auto result = session.ExecuteQuery(query, TTxControl::BeginTx()).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+            auto transaction = result.GetTransaction();
+            UNIT_ASSERT(transaction->IsActive());
+
+            const TString query2 = "UPDATE KeyValue SET Value = 'Vic'";
+            auto result2 = session.ExecuteQuery(query2, TTxControl::Tx(transaction->GetId())).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result2.GetStatus(), EStatus::SUCCESS, result2.GetIssues().ToString());
+            auto transaction2 = result2.GetTransaction();
+            UNIT_ASSERT(transaction2->IsActive());
+
+            auto checkResult = [&](TString table, TString expected) {
+                auto selectRes = db.ExecuteQuery(
+                    Sprintf("SELECT * FROM %s ORDER BY Key", table.data()),
+                    TTxControl::BeginTx().CommitTx()
+                ).ExtractValueSync();
+
+                UNIT_ASSERT_C(selectRes.IsSuccess(), selectRes.GetIssues().ToString());
+                CompareYson(expected, FormatResultSetYson(selectRes.GetResultSet(0)));
+            };
+            checkResult("TwoShard", R"([[[1u];["One"];[0]];[[2u];["Two"];[0]];[[3u];["Three"];[0]];[[4000000001u];["BigOne"];[0]];[[4000000002u];["BigTwo"];[0]];[[4000000003u];["BigThree"];[0]]])");
+            checkResult("KeyValue", R"([[[1u];["One"]];[[2u];["Two"]]])");
+            auto txRes = transaction->Commit().GetValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(txRes.GetStatus(), EStatus::SUCCESS, txRes.GetIssues().ToString());
+
+            checkResult("KeyValue", R"([[[1u];["Vic"]];[[2u];["Vic"]]])");
+            checkResult("TwoShard", R"([[[1u];["One"];[1]];[[2u];["Two"];[1]];[[3u];["Three"];[1]];[[4000000001u];["BigOne"];[1]];[[4000000002u];["BigTwo"];[1]];[[4000000003u];["BigThree"];[1]]])");
+        }
     }
 
     Y_UNIT_TEST(ExecuteQueryInteractiveTxCommitWithQuery) {