Add enum locks type instead of bools. (KIKIMR-18677)

5 files changed, 23 insertions, 21 deletions

diff --git a/ydb/core/kqp/executer_actor/kqp_data_executer.cpp b/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
index 797bcf0d6e..b97955b4f3 100644
--- a/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
+++ b/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
@@ -137,7 +137,7 @@ public:
 
         YQL_ENSURE(Request.IsolationLevel != NKikimrKqp::ISOLATION_LEVEL_UNDEFINED);
 
-        if (Request.AcquireLocksTxId || Request.ValidateLocks || Request.EraseLocks) {
+        if (Request.AcquireLocksTxId || Request.LocksOp == ELocksOp::Commit || Request.LocksOp == ELocksOp::Rollback) {
             YQL_ENSURE(Request.IsolationLevel == NKikimrKqp::ISOLATION_LEVEL_SERIALIZABLE);
         }
 
@@ -1306,7 +1306,7 @@ private:
             return false;
         }
 
-        if (Request.ValidateLocks && Request.EraseLocks) {
+        if (Request.LocksOp == ELocksOp::Commit) {
             YQL_ENSURE(!Request.UseImmediateEffects);
             return false;
         }
@@ -1969,7 +1969,7 @@ private:
 
         Request.TopicOperations.BuildTopicTxs(topicTxs);
 
-        const bool needRollback = Request.EraseLocks && !Request.ValidateLocks;
+        const bool needRollback = Request.LocksOp == ELocksOp::Rollback;
 
         VolatileTx = (
             // We want to use volatile transactions only when the feature is enabled
@@ -2004,11 +2004,9 @@ private:
             VolatileTx ||
             Request.TopicOperations.HasReadOperations())
         {
-            YQL_ENSURE(Request.ValidateLocks || Request.EraseLocks || VolatileTx);
+            YQL_ENSURE(Request.LocksOp == ELocksOp::Commit || Request.LocksOp == ELocksOp::Rollback || VolatileTx);
 
-            bool needCommit = (
-                (Request.ValidateLocks && Request.EraseLocks) ||
-                VolatileTx);
+            bool needCommit = Request.LocksOp == ELocksOp::Commit || VolatileTx;
 
             auto locksOp = needCommit
                 ? NKikimrTxDataShard::TKqpLocks::Commit
@@ -2077,7 +2075,7 @@ private:
                 }
             }
 
-            if (Request.ValidateLocks || needCommit) {
+            if (needCommit) {
                 NProtoBuf::RepeatedField<ui64> sendingShards(sendingShardsSet.begin(), sendingShardsSet.end());
                 NProtoBuf::RepeatedField<ui64> receivingShards(receivingShardsSet.begin(), receivingShardsSet.end());
 
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp b/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp
index f4e3444a58..18265decf2 100644
--- a/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp
+++ b/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp
@@ -96,7 +96,7 @@ IActor* CreateKqpExecuter(IKqpGateway::TExecPhysicalRequest&& request, const TSt
 {
     if (request.Transactions.empty()) {
         // commit-only or rollback-only data transaction
-        YQL_ENSURE(request.EraseLocks);
+        YQL_ENSURE(request.LocksOp == ELocksOp::Commit || request.LocksOp == ELocksOp::Rollback);
         return CreateKqpDataExecuter(std::move(request), database, userToken, counters, false, executerRetriesConfig, std::move(asyncIoFactory), creator);
     }
 
diff --git a/ydb/core/kqp/executer_actor/kqp_scan_executer.cpp b/ydb/core/kqp/executer_actor/kqp_scan_executer.cpp
index 00e42094a5..575c2fdfa9 100644
--- a/ydb/core/kqp/executer_actor/kqp_scan_executer.cpp
+++ b/ydb/core/kqp/executer_actor/kqp_scan_executer.cpp
@@ -66,8 +66,7 @@ public:
     {
         YQL_ENSURE(Request.Transactions.size() == 1);
         YQL_ENSURE(Request.DataShardLocks.empty());
-        YQL_ENSURE(!Request.ValidateLocks);
-        YQL_ENSURE(!Request.EraseLocks);
+        YQL_ENSURE(Request.LocksOp == ELocksOp::Unspecified);
         YQL_ENSURE(Request.IsolationLevel == NKikimrKqp::ISOLATION_LEVEL_UNDEFINED);
         YQL_ENSURE(Request.Snapshot.IsValid());
 
diff --git a/ydb/core/kqp/gateway/kqp_gateway.h b/ydb/core/kqp/gateway/kqp_gateway.h
index 1e10f4b7f6..c4be1e7fa2 100644
--- a/ydb/core/kqp/gateway/kqp_gateway.h
+++ b/ydb/core/kqp/gateway/kqp_gateway.h
@@ -61,6 +61,12 @@ struct TModuleResolverState : public TThrRefBase {
 
 void ApplyServiceConfig(NYql::TKikimrConfiguration& kqpConfig, const NKikimrConfig::TTableServiceConfig& serviceConfig);
 
+enum class ELocksOp {
+    Unspecified = 0,
+    Commit,
+    Rollback
+};
+
 class IKqpGateway : public NYql::IKikimrGateway {
 public:
     struct TPhysicalTxData : private TMoveOnly {
@@ -124,8 +130,7 @@ public:
         TVector<TPhysicalTxData> Transactions;
         TMap<ui64, TVector<NKikimrTxDataShard::TLock>> DataShardLocks;
         NKikimr::NKqp::TTxAllocatorState::TPtr TxAlloc;
-        bool ValidateLocks = false;
-        bool EraseLocks = false;
+        ELocksOp LocksOp = ELocksOp::Unspecified;
         TMaybe<ui64> AcquireLocksTxId;
         TDuration Timeout;
         TMaybe<TDuration> CancelAfter;
diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
index af5dd3041b..6b4509ba2b 100644
--- a/ydb/core/kqp/session_actor/kqp_session_actor.cpp
+++ b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -976,12 +976,13 @@ public:
             }
 
             if (txCtx.Locks.HasLocks() || txCtx.TopicOperations.HasReadOperations()) {
-                request.ValidateLocks = !txCtx.GetSnapshot().IsValid() || txCtx.TxHasEffects() ||
-                    txCtx.TopicOperations.HasReadOperations();
-                request.EraseLocks = true;
-
-                LOG_D("TExecPhysicalRequest, tx has locks, ValidateLocks: " << request.ValidateLocks
-                        << " EraseLocks: " << request.EraseLocks);
+                if (!txCtx.GetSnapshot().IsValid() || txCtx.TxHasEffects() || txCtx.TopicOperations.HasReadOperations()) {
+                    LOG_D("TExecPhysicalRequest, tx has commit locks");
+                    request.LocksOp = ELocksOp::Commit;
+                } else {
+                    LOG_D("TExecPhysicalRequest, tx has rollback locks");
+                    request.LocksOp = ELocksOp::Rollback;
+                }
 
                 for (auto& [lockId, lock] : txCtx.Locks.LocksMap) {
                     auto dsLock = ExtractLock(lock.GetValueRef(txCtx.Locks.LockType));
@@ -1588,8 +1589,7 @@ public:
             AppData()->TimeProvider, AppData()->RandomProvider);
         auto request = PreparePhysicalRequest(nullptr, allocPtr);
 
-        request.EraseLocks = true;
-        request.ValidateLocks = false;
+        request.LocksOp = ELocksOp::Rollback;
 
         // Should tx with empty LocksMap be aborted?
         for (auto& [lockId, lock] : txCtx->Locks.LocksMap) {