KQP locks in EvWrite (#1551)

10 files changed, 473 insertions, 36 deletions

diff --git a/ydb/core/protos/data_events.proto b/ydb/core/protos/data_events.proto
index 25945777d2c..337b5091241 100644
--- a/ydb/core/protos/data_events.proto
+++ b/ydb/core/protos/data_events.proto
@@ -99,6 +99,7 @@ message TEvWriteResult {
         STATUS_CANCELLED = 6;        
         STATUS_BAD_REQUEST = 7;
         STATUS_SCHEME_CHANGED = 8;
+        STATUS_LOCKS_BROKEN = 9;
     }
 
     // Status
diff --git a/ydb/core/tx/data_events/events.h b/ydb/core/tx/data_events/events.h
index 8fab00c8699..df69ae04c87 100644
--- a/ydb/core/tx/data_events/events.h
+++ b/ydb/core/tx/data_events/events.h
@@ -107,6 +107,19 @@ struct TDataEvents {
             return result;
         }
 
+        void AddTxLock(ui64 lockId, ui64 shard, ui32 generation, ui64 counter, ui64 ssId, ui64 pathId, bool hasWrites) {
+            auto entry = Record.AddTxLocks();
+            entry->SetLockId(lockId);
+            entry->SetDataShard(shard);
+            entry->SetGeneration(generation);
+            entry->SetCounter(counter);
+            entry->SetSchemeShard(ssId);
+            entry->SetPathId(pathId);
+            if (hasWrites) {
+                entry->SetHasWrites(true);
+            }
+        }
+
         TString GetError() const {
             return TStringBuilder() << "Status: " << Record.GetStatus() << " Issues: " << Record.GetIssues();
         }
diff --git a/ydb/core/tx/datashard/datashard_pipeline.cpp b/ydb/core/tx/datashard/datashard_pipeline.cpp
index 1e0c3353c9a..bd70f187cbc 100644
--- a/ydb/core/tx/datashard/datashard_pipeline.cpp
+++ b/ydb/core/tx/datashard/datashard_pipeline.cpp
@@ -1595,6 +1595,11 @@ TOperation::TPtr TPipeline::BuildOperation(NEvents::TDataEvents::TEvWrite::TPtr&
 
     writeTx->ExtractKeys(true);
 
+    if (!writeTx->Ready()) {
+        badRequest(EvWrite::Convertor::ConvertErrCode(writeOp->GetWriteTx()->GetErrCode()), TStringBuilder() << "Cannot parse tx keys " << writeOp->GetTxId() << ". " << writeOp->GetWriteTx()->GetErrCode() << ": " << writeOp->GetWriteTx()->GetErrStr());
+        return writeOp;
+    }
+
     switch (rec.txmode()) {
         case NKikimrDataEvents::TEvWrite::MODE_PREPARE:
             break;
diff --git a/ydb/core/tx/datashard/datashard_user_db.h b/ydb/core/tx/datashard/datashard_user_db.h
index 04a856d4757..a7ac89eaca8 100644
--- a/ydb/core/tx/datashard/datashard_user_db.h
+++ b/ydb/core/tx/datashard/datashard_user_db.h
@@ -152,6 +152,7 @@ private:
     YDB_ACCESSOR_DEF(ui32, LockNodeId);
     YDB_ACCESSOR_DEF(ui64, VolatileTxId);
     YDB_ACCESSOR_DEF(bool, IsImmediateTx);
+    YDB_ACCESSOR_DEF(bool, IsWriteTx);
     YDB_ACCESSOR_DEF(bool, IsRepeatableSnapshot);
 
     YDB_ACCESSOR_DEF(TRowVersion, ReadVersion);
diff --git a/ydb/core/tx/datashard/datashard_ut_read_iterator.cpp b/ydb/core/tx/datashard/datashard_ut_read_iterator.cpp
index 044080f8d71..9863e8fd8b3 100644
--- a/ydb/core/tx/datashard/datashard_ut_read_iterator.cpp
+++ b/ydb/core/tx/datashard/datashard_ut_read_iterator.cpp
@@ -3,12 +3,14 @@
 #include "datashard_active_transaction.h"
 #include "read_iterator.h"
 
+#include <ydb/core/testlib/tablet_helpers.h>
 #include <ydb/core/formats/arrow/arrow_helpers.h>
 #include <ydb/core/formats/arrow/converter.h>
 #include <ydb/core/kqp/ut/common/kqp_ut_common.h>
 #include <ydb/core/tablet_flat/shared_cache_events.h>
 #include <ydb/core/tx/tx_proxy/proxy.h>
 #include <ydb/core/tx/tx_proxy/read_table.h>
+#include <ydb/core/tx/long_tx_service/public/lock_handle.h>
 
 #include <ydb/core/tx/data_events/events.h>
 #include <ydb/core/tx/data_events/payload_helper.h>
@@ -731,9 +733,23 @@ struct TTestHelper {
         UNIT_ASSERT_VALUES_EQUAL(rowsRead, Min(rowCount, limit));
     }
 
+    void TestReadOneKey(const TString& tableName, const std::vector<ui32>& keys, ui32 value)
+    {
+        auto readRequest = GetBaseReadRequest(tableName, 1);
+        AddKeyQuery(*readRequest, keys);
+
+        auto readResult = SendRead(tableName, readRequest.release());
+
+        std::vector<std::vector<ui32>> gold(1);
+        std::copy(keys.begin(), keys.end(), std::back_inserter(gold[0]));
+        gold[0].push_back(value);
+
+        CheckResult(Tables[tableName].UserTable, *readResult, gold);
+    }
+
     void WriteRowTwin(const TString& tableName, const TVector<ui32>& values, bool isEvWrite) {
         if(isEvWrite)
-            WriteRow(tableName, ++TxId, values);
+            WriteRow(tableName, values);
         else
             ExecSQL(Server, Sender, TStringBuilder() 
                 << "UPSERT INTO `/Root/" << tableName << "`\n"
@@ -741,7 +757,7 @@ struct TTestHelper {
                 << "VALUES\n(" << JoinSeq(",", values) << ");");
     }
 
-    NKikimrDataEvents::TEvWriteResult WriteRow(const TString& tableName, ui64 txId, const TVector<ui32>& values, NKikimrDataEvents::TEvWrite::ETxMode txMode = NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE) {
+    std::unique_ptr<NEvents::TDataEvents::TEvWrite> MakeWriteRequest(const TString& tableName, ui64 txId, const TVector<ui32>& values, NKikimrDataEvents::TEvWrite::ETxMode txMode = NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE) {
         const auto& table = Tables[tableName];
 
         auto opts = TShardedTableOptions().Columns(table.Columns);
@@ -762,7 +778,17 @@ struct TTestHelper {
         ui64 payloadIndex = NKikimr::NEvWrite::TPayloadHelper<NKikimr::NEvents::TDataEvents::TEvWrite>(*evWrite).AddDataToPayload(matrix.ReleaseBuffer());
         evWrite->AddOperation(NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPSERT, table.TableId, columnIds, payloadIndex, NKikimrDataEvents::FORMAT_CELLVEC);
 
-        return Write(*Server->GetRuntime(), Sender, table.TabletId, std::move(evWrite));
+        return evWrite;
+    }
+
+    NKikimrDataEvents::TEvWriteResult SendWrite(ui64 tabletId, std::unique_ptr<NEvents::TDataEvents::TEvWrite> writeRequest, NKikimrDataEvents::TEvWriteResult::EStatus expectedStatus = NKikimrDataEvents::TEvWriteResult::STATUS_UNSPECIFIED) {
+        return Write(*Server->GetRuntime(), Sender, tabletId, std::move(writeRequest), expectedStatus);
+    }
+
+    NKikimrDataEvents::TEvWriteResult WriteRow(const TString& tableName, const TVector<ui32>& values) {
+        auto writeRequest = MakeWriteRequest(tableName, ++TxId, values);
+
+        return SendWrite(Tables[tableName].TabletId, std::move(writeRequest));
     }
 
     struct THangedReturn {
@@ -3117,6 +3143,212 @@ Y_UNIT_TEST_SUITE(DataShardReadIterator) {
         UNIT_ASSERT_VALUES_EQUAL(readResults, 0);
     }
 
+    Y_UNIT_TEST(ShouldCommitLocksWhenReadWriteInOneTransaction) {
+        TTestHelper helper;
+
+        auto runtime = helper.Server->GetRuntime();
+
+        const ui64 lockTxId = 1011121314;
+        const TString tableName = "table-1";
+        const ui64 tabletId = helper.Tables["table-1"].TabletId;
+        const ui64 nodeId = runtime->GetNodeId();
+
+        NLongTxService::TLockHandle lockHandle(lockTxId, runtime->GetActorSystem(0));
+
+        // Read in a transaction.
+        auto readRequest1 = helper.GetBaseReadRequest(tableName, 1);
+        readRequest1->Record.SetLockTxId(lockTxId);
+        AddKeyQuery(*readRequest1, {1, 1, 1});
+
+        auto readResult1 = helper.SendRead(tableName, readRequest1.release());
+
+        CheckResult(helper.Tables[tableName].UserTable, *readResult1, { {1, 1, 1, 100} });
+
+        UNIT_ASSERT_VALUES_EQUAL(readResult1->Record.TxLocksSize(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(readResult1->Record.BrokenTxLocksSize(), 0);
+        const auto& readLock = readResult1->Record.GetTxLocks(0);
+
+        // Write in the same transaction.
+        {
+            auto writeRequest = helper.MakeWriteRequest(tableName, ++helper.TxId, {1, 1, 1, 101});
+            writeRequest->Record.SetLockTxId(lockTxId);
+            writeRequest->Record.SetLockNodeId(nodeId);
+
+            NKikimrDataEvents::TEvWriteResult writeResult = helper.SendWrite(tabletId, std::move(writeRequest));
+
+            UNIT_ASSERT_VALUES_EQUAL(writeResult.GetStatus(), NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+            UNIT_ASSERT(writeResult.TxLocksSize() == 1);
+            const auto& writeLock = writeResult.GetTxLocks(0);
+            UNIT_ASSERT_VALUES_EQUAL(writeLock.GetLockId(), readLock.GetLockId());
+            UNIT_ASSERT_VALUES_EQUAL(writeLock.GetDataShard(), readLock.GetDataShard());
+            UNIT_ASSERT_VALUES_EQUAL(writeLock.GetGeneration(), readLock.GetGeneration());
+            UNIT_ASSERT_VALUES_EQUAL(writeLock.GetCounter(), readLock.GetCounter());
+            UNIT_ASSERT_VALUES_EQUAL(writeLock.GetSchemeShard(), readLock.GetSchemeShard());
+            UNIT_ASSERT_VALUES_EQUAL(writeLock.GetPathId(), readLock.GetPathId());
+            UNIT_ASSERT_VALUES_UNEQUAL(writeLock.GetHasWrites(), readLock.GetHasWrites());
+        }
+
+        // Commit locks.
+        {
+            auto writeRequest = std::make_unique<NKikimr::NEvents::TDataEvents::TEvWrite>(++helper.TxId, NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE);
+            NKikimrDataEvents::TKqpLocks& kqpLocks = *writeRequest->Record.MutableLocks();
+            kqpLocks.MutableLocks()->CopyFrom(readResult1->Record.GetTxLocks());
+            kqpLocks.AddSendingShards(tabletId);
+            kqpLocks.AddReceivingShards(tabletId);
+            kqpLocks.SetOp(::NKikimrDataEvents::TKqpLocks::Commit);
+
+            NKikimrDataEvents::TEvWriteResult writeResult = helper.SendWrite(tabletId, std::move(writeRequest));
+
+            UNIT_ASSERT_VALUES_EQUAL(writeResult.GetStatus(), NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+        }
+
+        // Read written data.
+        helper.TestReadOneKey(tableName, {1, 1, 1}, 101);
+    }
+
+    Y_UNIT_TEST(ShouldCommitLocksWhenReadWriteInSeparateTransactions) {
+        TTestHelper helper;
+
+        auto runtime = helper.Server->GetRuntime();
+
+        const ui64 lockTxId = 1011121314;
+        const TString tableName = "table-1";
+        const ui64 tabletId = helper.Tables["table-1"].TabletId;
+        const ui64 nodeId = runtime->GetNodeId();
+
+        NLongTxService::TLockHandle lockHandle(lockTxId, runtime->GetActorSystem(0));
+
+        // Write in first transaction.
+        auto writeRequest = helper.MakeWriteRequest(tableName, ++helper.TxId, {1, 1, 1, 101});
+        writeRequest->Record.SetLockTxId(lockTxId);
+        writeRequest->Record.SetLockNodeId(nodeId);
+
+        NKikimrDataEvents::TEvWriteResult writeResult = helper.SendWrite(tabletId, std::move(writeRequest));
+
+        UNIT_ASSERT_VALUES_EQUAL(writeResult.GetStatus(), NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+        UNIT_ASSERT(writeResult.TxLocksSize() == 1);
+        const auto& writeLock = writeResult.GetTxLocks(0);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetLockId(), lockTxId);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetDataShard(), tabletId);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetGeneration(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetCounter(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetHasWrites(), true);
+
+        // Read in separate transaction. No dirty-read.
+        helper.TestReadOneKey(tableName, {1, 1, 1}, 100);
+
+        // Commit locks in first transaction. 
+        auto writeRequest2 = std::make_unique<NKikimr::NEvents::TDataEvents::TEvWrite>(++helper.TxId, NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE);
+        NKikimrDataEvents::TKqpLocks& kqpLocks2 = *writeRequest2->Record.MutableLocks();
+        kqpLocks2.MutableLocks()->CopyFrom(writeResult.GetTxLocks());
+        kqpLocks2.AddSendingShards(tabletId);
+        kqpLocks2.AddReceivingShards(tabletId);
+        kqpLocks2.SetOp(::NKikimrDataEvents::TKqpLocks::Commit);
+
+        NKikimrDataEvents::TEvWriteResult writeResult2 = helper.SendWrite(tabletId, std::move(writeRequest2));
+
+        UNIT_ASSERT_VALUES_EQUAL(writeResult.GetStatus(), NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+
+        // Read written data.
+        helper.TestReadOneKey(tableName, {1, 1, 1}, 101);
+    }
+
+    Y_UNIT_TEST(ShouldRollbackLocksWhenWrite) {
+        TTestHelper helper;
+
+        auto runtime = helper.Server->GetRuntime();
+
+        const ui64 lockTxId = 1011121314;
+        const TString tableName = "table-1";
+        const ui64 tabletId = helper.Tables["table-1"].TabletId;
+        const ui64 nodeId = runtime->GetNodeId();
+
+        NLongTxService::TLockHandle lockHandle(lockTxId, runtime->GetActorSystem(0));
+
+        // Write in transaction.
+        auto writeRequest = helper.MakeWriteRequest(tableName, ++helper.TxId, {1, 1, 1, 101});
+        writeRequest->Record.SetLockTxId(lockTxId);
+        writeRequest->Record.SetLockNodeId(nodeId);
+
+        NKikimrDataEvents::TEvWriteResult writeResult = helper.SendWrite(tabletId, std::move(writeRequest));
+
+        UNIT_ASSERT_VALUES_EQUAL(writeResult.GetStatus(), NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+        UNIT_ASSERT(writeResult.TxLocksSize() == 1);
+        const auto& writeLock = writeResult.GetTxLocks(0);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetLockId(), lockTxId);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetDataShard(), tabletId);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetGeneration(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetCounter(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetHasWrites(), true);
+
+        // Rollback locks in transaction.
+        auto writeRequest2 = std::make_unique<NKikimr::NEvents::TDataEvents::TEvWrite>(++helper.TxId, NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE);
+        NKikimrDataEvents::TKqpLocks& kqpLocks2 = *writeRequest2->Record.MutableLocks();
+        kqpLocks2.MutableLocks()->CopyFrom(writeResult.GetTxLocks());
+        kqpLocks2.SetOp(::NKikimrDataEvents::TKqpLocks::Rollback);
+
+        NKikimrDataEvents::TEvWriteResult writeResult2 = helper.SendWrite(tabletId, std::move(writeRequest2));
+
+        UNIT_ASSERT_VALUES_EQUAL(writeResult2.GetStatus(), NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+
+        // Read origin data.
+        helper.TestReadOneKey(tableName, {1, 1, 1}, 100);
+    }    
+
+    Y_UNIT_TEST_TWIN(ShouldReturnBrokenLockWhenWriteInSeparateTransactions, EvWrite) {
+        TTestHelper helper;
+
+        auto runtime = helper.Server->GetRuntime();
+
+        const ui64 lockTxId = 1011121314;
+        const TString tableName = "table-1";
+        const ui64 tabletId = helper.Tables["table-1"].TabletId;
+        const ui64 nodeId = runtime->GetNodeId();
+
+        NLongTxService::TLockHandle lockHandle(lockTxId, runtime->GetActorSystem(0));
+
+        // Write in first transaction.
+        auto writeRequest = helper.MakeWriteRequest(tableName, ++helper.TxId, {1, 1, 1, 101});
+        writeRequest->Record.SetLockTxId(lockTxId);
+        writeRequest->Record.SetLockNodeId(nodeId);
+
+        NKikimrDataEvents::TEvWriteResult writeResult = helper.SendWrite(tabletId, std::move(writeRequest));
+
+        UNIT_ASSERT_VALUES_EQUAL(writeResult.GetStatus(), NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+        UNIT_ASSERT(writeResult.TxLocksSize() == 1);
+        const auto& writeLock = writeResult.GetTxLocks(0);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetLockId(), lockTxId);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetDataShard(), tabletId);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetGeneration(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetCounter(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock.GetHasWrites(), true);
+
+        // Breaks lock obtained above using write in separate transaction.
+        helper.WriteRowTwin(tableName, {1, 1, 1, 202}, EvWrite);
+
+        // Commit locks in first transaction. They should be broken.
+        auto writeRequest2 = std::make_unique<NKikimr::NEvents::TDataEvents::TEvWrite>(++helper.TxId, NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE);
+        NKikimrDataEvents::TKqpLocks& kqpLocks2 = *writeRequest2->Record.MutableLocks();
+        kqpLocks2.MutableLocks()->CopyFrom(writeResult.GetTxLocks());
+        kqpLocks2.AddSendingShards(tabletId);
+        kqpLocks2.AddReceivingShards(tabletId);
+        kqpLocks2.SetOp(::NKikimrDataEvents::TKqpLocks::Commit);
+
+        NKikimrDataEvents::TEvWriteResult writeResult2 = helper.SendWrite(tabletId, std::move(writeRequest2), NKikimrDataEvents::TEvWriteResult::STATUS_LOCKS_BROKEN);
+
+        UNIT_ASSERT_VALUES_EQUAL(writeResult2.GetStatus(), NKikimrDataEvents::TEvWriteResult::STATUS_LOCKS_BROKEN);
+        UNIT_ASSERT(writeResult2.TxLocksSize() == 1);
+        const auto& writeLock2 = writeResult2.GetTxLocks(0);
+        UNIT_ASSERT_VALUES_EQUAL(writeLock2.GetLockId(), writeLock.GetLockId());
+        UNIT_ASSERT_VALUES_EQUAL(writeLock2.GetDataShard(), writeLock.GetDataShard());
+        UNIT_ASSERT_VALUES_EQUAL(writeLock2.GetGeneration(), writeLock.GetGeneration());
+        UNIT_ASSERT_VALUES_EQUAL(writeLock2.GetCounter(), writeLock.GetCounter());
+        UNIT_ASSERT_VALUES_EQUAL(writeLock2.GetHasWrites(), writeLock.GetHasWrites());
+
+        // read written data
+        helper.TestReadOneKey(tableName, {1, 1, 1}, 202);
+    }
+
     Y_UNIT_TEST_TWIN(ShouldReturnBrokenLockWhenReadKey, EvWrite) {
         TTestHelper helper;
 
diff --git a/ydb/core/tx/datashard/datashard_write_operation.cpp b/ydb/core/tx/datashard/datashard_write_operation.cpp
index 5bd48c0a031..335dc5c6ff7 100644
--- a/ydb/core/tx/datashard/datashard_write_operation.cpp
+++ b/ydb/core/tx/datashard/datashard_write_operation.cpp
@@ -31,6 +31,8 @@ TValidatedWriteTx::TValidatedWriteTx(TDataShard* self, TTransactionContext& txc,
     ComputeTxSize();
     NActors::NMemory::TLabel<MemoryLabelValidatedDataTx>::Add(TxSize);
 
+    UserDb.SetIsWriteTx(true);
+    
     if (LockTxId()) {
         UserDb.SetLockTxId(LockTxId());
         UserDb.SetLockNodeId(LockNodeId());
@@ -43,10 +45,12 @@ TValidatedWriteTx::TValidatedWriteTx(TDataShard* self, TTransactionContext& txc,
 
     LOG_TRACE_S(Ctx, NKikimrServices::TX_DATASHARD, "Parsing write transaction for " << globalTxId << " at " << TabletId << ", record: " << GetRecord().ShortDebugString());
 
-    if (!ParseRecord(self->TableInfos))
-        return;
+    if (HasOperations()) {
+        if (!ParseOperations(self->TableInfos))
+            return;
 
-    SetTxKeys(RecordOperation().GetColumnIds());
+        SetTxKeys(RecordOperation().GetColumnIds());
+    }
 
     KqpSetTxLocksKeys(GetKqpLocks(), self->SysLocksTable(), KeyValidator);
     KeyValidator.GetInfo().SetLoaded();
@@ -56,7 +60,7 @@ TValidatedWriteTx::~TValidatedWriteTx() {
     NActors::NMemory::TLabel<MemoryLabelValidatedDataTx>::Sub(TxSize);
 }
 
-bool TValidatedWriteTx::ParseRecord(const TDataShard::TTableInfos& tableInfos) {
+bool TValidatedWriteTx::ParseOperations(const TDataShard::TTableInfos& tableInfos) {
     if (GetRecord().GetOperations().size() != 1)
     {
         ErrCode = NKikimrTxDataShard::TError::BAD_ARGUMENT;
@@ -194,11 +198,20 @@ void TValidatedWriteTx::SetTxKeys(const ::google::protobuf::RepeatedField<::NPro
 
 ui32 TValidatedWriteTx::ExtractKeys(bool allowErrors)
 {
+    if (!HasOperations())
+        return 0;
+
     SetTxKeys(RecordOperation().GetColumnIds());
 
     bool isValid = ReValidateKeys();
-    Y_ABORT_UNLESS(allowErrors || isValid, "Validation errors: %s", ErrStr.data());
-
+    if (allowErrors) {
+        if (!isValid) {
+            return 0;
+        }
+    } else {
+        Y_ABORT_UNLESS(isValid, "Validation errors: %s", ErrStr.data());
+    }
+    
     return KeysCount();
 }
 
diff --git a/ydb/core/tx/datashard/datashard_write_operation.h b/ydb/core/tx/datashard/datashard_write_operation.h
index b14d829244c..1ae40f39c02 100644
--- a/ydb/core/tx/datashard/datashard_write_operation.h
+++ b/ydb/core/tx/datashard/datashard_write_operation.h
@@ -40,6 +40,10 @@ public:
         return GetRecord().operations(0);
     }
 
+    ui64 GetTxId() const {
+        return UserDb.GetGlobalTxId();
+    }
+
     ui64 LockTxId() const {
         return GetRecord().locktxid();
     }
@@ -127,16 +131,23 @@ public:
         return Source != TActorId();
     }
 
-    inline const ::NKikimrDataEvents::TKqpLocks& GetKqpLocks() const {
+    const ::NKikimrDataEvents::TKqpLocks& GetKqpLocks() const {
         return GetRecord().locks();
     }
+    bool HasKqpLocks() const {
+        return GetRecord().has_locks();
+    }
 
-    bool ParseRecord(const TDataShard::TTableInfos& tableInfos);
+    bool ParseOperations(const TDataShard::TTableInfos& tableInfos);
     void SetTxKeys(const ::google::protobuf::RepeatedField<::NProtoBuf::uint32>& columnIds);
 
     ui32 ExtractKeys(bool allowErrors);
     bool ReValidateKeys();
 
+    ui64 HasOperations() const {
+        return GetRecord().operations().size() != 0;
+    }
+
     ui32 KeysCount() const {
         return TxInfo().WritesCount;
     }
diff --git a/ydb/core/tx/datashard/key_validator.cpp b/ydb/core/tx/datashard/key_validator.cpp
index 0dd217a4a74..ee9f9b5d1d4 100644
--- a/ydb/core/tx/datashard/key_validator.cpp
+++ b/ydb/core/tx/datashard/key_validator.cpp
@@ -66,10 +66,11 @@ void TKeyValidator::AddWriteRange(const TTableId& tableId, const TTableRange& ra
 }
 
 TKeyValidator::TValidateOptions::TValidateOptions(const TDataShardUserDb& userDb)
-    : LockTxId(userDb.GetLockTxId())
-    , LockNodeId(userDb.GetLockNodeId())
+    : IsLockTxId(static_cast<bool>(userDb.GetLockTxId()))
+    , IsLockNodeId(static_cast<bool>(userDb.GetLockNodeId()))
     , IsRepeatableSnapshot(userDb.GetIsRepeatableSnapshot())
     , IsImmediateTx(userDb.GetIsImmediateTx())
+    , IsWriteTx(userDb.GetIsWriteTx())
 {
 }
 
@@ -77,17 +78,16 @@ bool TKeyValidator::IsValidKey(TKeyDesc& key, const TValidateOptions& opt) const
     if (TSysTables::IsSystemTable(key.TableId))
         return true;
 
-    if (opt.LockTxId) {
+    if (key.RowOperation != TKeyDesc::ERowOperation::Read && !opt.IsWriteTx) {
         // Prevent updates/erases with LockTxId set, unless it's allowed for immediate mvcc txs
-        if (key.RowOperation != TKeyDesc::ERowOperation::Read &&
-            (!Self.GetEnableLockedWrites() || !opt.IsImmediateTx || !opt.IsRepeatableSnapshot || !opt.LockNodeId))
-        {
-            key.Status = TKeyDesc::EStatus::OperationNotSupported;
-            return false;
+        if (opt.IsLockTxId) {
+            if (!(Self.GetEnableLockedWrites() && opt.IsImmediateTx && opt.IsRepeatableSnapshot && opt.IsLockNodeId)) {
+                key.Status = TKeyDesc::EStatus::OperationNotSupported;
+                return false;
+            }
         }
-    } else if (opt.IsRepeatableSnapshot) {
         // Prevent updates/erases in repeatable mvcc txs
-        if (key.RowOperation != TKeyDesc::ERowOperation::Read) {
+        else if (opt.IsRepeatableSnapshot) {
             key.Status = TKeyDesc::EStatus::OperationNotSupported;
             return false;
         }
diff --git a/ydb/core/tx/datashard/key_validator.h b/ydb/core/tx/datashard/key_validator.h
index 0106aa17c96..9b3edd0c667 100644
--- a/ydb/core/tx/datashard/key_validator.h
+++ b/ydb/core/tx/datashard/key_validator.h
@@ -24,11 +24,12 @@ public:
     void AddWriteRange(const TTableId& tableId, const TTableRange& range, const TVector<NScheme::TTypeInfo>& keyTypes, const TVector<TColumnWriteMeta>& columns, bool isPureEraseOp);
     
     struct TValidateOptions {
-        ui64 LockTxId;
-        ui64 LockNodeId;
+        bool IsLockTxId;
+        bool IsLockNodeId;
         bool IsRepeatableSnapshot;
         bool IsImmediateTx;
-
+        bool IsWriteTx;
+        
         TValidateOptions(const TDataShardUserDb& userDb);
     };
 
diff --git a/ydb/core/tx/datashard/write_unit.cpp b/ydb/core/tx/datashard/write_unit.cpp
index 15d4065b279..5d5ea43b059 100644
--- a/ydb/core/tx/datashard/write_unit.cpp
+++ b/ydb/core/tx/datashard/write_unit.cpp
@@ -3,6 +3,7 @@
 #include "setup_sys_locks.h"
 #include "datashard_locks_db.h"
 #include "datashard_user_db.h"
+#include "datashard_kqp.h"
 
 #include <ydb/core/engine/mkql_engine_flat_host.h>
 
@@ -37,25 +38,47 @@ public:
         return !op->HasRuntimeConflicts();
     }
 
-    void DoExecute(TDataShard* self, TWriteOperation* writeOp, TTransactionContext& txc, const TActorContext& ctx) {
+    void AddLocksToResult(TWriteOperation* writeOp, const TActorContext& ctx) {
+        NEvents::TDataEvents::TEvWriteResult& writeResult = *writeOp->GetWriteResult();
+
+        auto locks = DataShard.SysLocksTable().ApplyLocks();
+        LOG_TRACE_S(ctx, NKikimrServices::TX_DATASHARD, "add locks to result: " << locks.size());
+        for (const auto& lock : locks) {
+            if (lock.IsError()) {
+                LOG_NOTICE_S(ctx, NKikimrServices::TX_DATASHARD, "Lock is not set for " << *writeOp << " at " << DataShard.TabletID() << " lock " << lock);
+            }
+
+            writeResult.AddTxLock(lock.LockId, lock.DataShard, lock.Generation, lock.Counter, lock.SchemeShard, lock.PathId, lock.HasWrites);
+
+            LOG_TRACE_S(ctx, NKikimrServices::TX_DATASHARD, "add lock to result: " << writeResult.Record.GetTxLocks().rbegin()->ShortDebugString());
+        }
+        DataShard.SubscribeNewLocks(ctx);
+    }
+
+    void DoUpdateToUserDb(TWriteOperation* writeOp, TTransactionContext& txc, const TActorContext& ctx) {
         TValidatedWriteTx::TPtr& writeTx = writeOp->GetWriteTx();
 
+        if (!writeTx->HasOperations()) {
+            LOG_DEBUG_S(ctx, NKikimrServices::TX_DATASHARD, "Skip empty write operation for " << *writeOp << " at " << DataShard.TabletID());
+            return;
+        }
+
         const ui64 tableId = writeTx->GetTableId().PathId.LocalPathId;
-        const TTableId fullTableId(self->GetPathOwnerId(), tableId);
-        const ui64 localTableId = self->GetLocalTableId(fullTableId);
+        const TTableId fullTableId(DataShard.GetPathOwnerId(), tableId);
+        const ui64 localTableId = DataShard.GetLocalTableId(fullTableId);
         if (localTableId == 0) {
             writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_INTERNAL_ERROR, TStringBuilder() << "Unknown table id " << tableId);
             return;
         }
-        const ui64 shadowTableId = self->GetShadowTableId(fullTableId);
-        const TUserTable& TableInfo_ = *self->GetUserTables().at(tableId);
+        const ui64 shadowTableId = DataShard.GetShadowTableId(fullTableId);
+        const TUserTable& TableInfo_ = *DataShard.GetUserTables().at(tableId);
         Y_ABORT_UNLESS(TableInfo_.LocalTid == localTableId);
         Y_ABORT_UNLESS(TableInfo_.ShadowTid == shadowTableId);
 
         const NTable::TScheme& scheme = txc.DB.GetScheme();
         const NTable::TScheme::TTableInfo* tableInfo = scheme.GetTableInfo(localTableId);
 
-        auto [readVersion, writeVersion] = self->GetReadWriteVersions(writeOp);
+        auto [readVersion, writeVersion] = DataShard.GetReadWriteVersions(writeOp);
         writeTx->SetReadVersion(readVersion);
         writeTx->SetWriteVersion(writeVersion);
 
@@ -94,12 +117,10 @@ public:
             writeTx->GetUserDb().UpdateRow(fullTableId, key, ops);
         }
 
-        self->IncCounter(COUNTER_WRITE_ROWS, matrix.GetRowCount());
-        self->IncCounter(COUNTER_WRITE_BYTES, matrix.GetBuffer().size());
+        DataShard.IncCounter(COUNTER_WRITE_ROWS, matrix.GetRowCount());
+        DataShard.IncCounter(COUNTER_WRITE_BYTES, matrix.GetBuffer().size());
 
-        writeOp->SetWriteResult(NEvents::TDataEvents::TEvWriteResult::BuildCommited(self->TabletID(), writeOp->GetTxId()));
-
-        LOG_DEBUG_S(ctx, NKikimrServices::TX_DATASHARD, "Executed write operation for " << *writeOp << " at " << self->TabletID());
+        LOG_DEBUG_S(ctx, NKikimrServices::TX_DATASHARD, "Executed write operation for " << *writeOp << " at " << DataShard.TabletID());
     }
 
     EExecutionStatus Execute(TOperation::TPtr op, TTransactionContext& txc, const TActorContext& ctx) override {
@@ -128,6 +149,8 @@ public:
         TDataShardLocksDb locksDb(DataShard, txc);
         TSetupSysLocks guardLocks(op, DataShard, &locksDb);
 
+        ui64 tabletId = DataShard.TabletID();
+
         const TValidatedWriteTx::TPtr& writeTx = writeOp->GetWriteTx();
 
         if (op->IsImmediate() && !writeOp->ReValidateKeys()) {
@@ -145,7 +168,144 @@ public:
         }
 
         try {
-            DoExecute(&DataShard, writeOp, txc, ctx);
+            const ui64 txId = writeTx->GetTxId();
+            const auto* kqpLocks = writeTx->HasKqpLocks() ? &writeTx->GetKqpLocks() : nullptr;
+            const auto& inReadSets = op->InReadSets();
+            auto& awaitingDecisions = op->AwaitingDecisions();
+            auto& outReadSets = op->OutReadSets();
+            bool useGenericReadSets = true;
+            TSysLocks& sysLocks = DataShard.SysLocksTable();
+
+            ui64 consumedMemory = writeTx->GetTxSize();
+            if (MaybeRequestMoreTxMemory(consumedMemory, txc)) {
+                LOG_TRACE_S(ctx, NKikimrServices::TX_DATASHARD, "Operation " << *op << " at " << DataShard.TabletID() << " requested " << txc.GetRequestedMemory() << " more memory");
+
+                DataShard.IncCounter(COUNTER_TX_WAIT_RESOURCE);
+                return EExecutionStatus::Restart;
+            }
+
+            if (guardLocks.LockTxId) {
+                switch (DataShard.SysLocksTable().EnsureCurrentLock()) {
+                    case EEnsureCurrentLock::Success:
+                        // Lock is valid, we may continue with reads and side-effects
+                        break;
+
+                    case EEnsureCurrentLock::Broken:
+                        // Lock is valid, but broken, we could abort early in some
+                        // cases, but it doesn't affect correctness.
+                        break;
+
+                    case EEnsureCurrentLock::TooMany:
+                        // Lock cannot be created, it's not necessarily a problem
+                        // for read-only transactions, for non-readonly we need to
+                        // abort;
+                        if (op->IsReadOnly()) {
+                            break;
+                        }
+
+                        [[fallthrough]];
+
+                    case EEnsureCurrentLock::Abort:
+                        // Lock cannot be created and we must abort
+                        LOG_TRACE_S(ctx, NKikimrServices::TX_DATASHARD, "Operation " << *op << " at " << tabletId << " aborting because it cannot acquire locks");
+                        writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_LOCKS_BROKEN, "Operation is aborting because it cannot acquire locks");
+                        return EExecutionStatus::Executed;
+                }
+            }
+
+            bool keepOutReadSets = !op->HasVolatilePrepareFlag();
+
+            Y_DEFER {
+                // We need to clear OutReadSets and AwaitingDecisions for
+                // volatile transactions, except when we commit them.
+                if (!keepOutReadSets) {
+                    outReadSets.clear();
+                    awaitingDecisions.clear();
+                }
+            };
+
+            auto [validated, brokenLocks] = op->HasVolatilePrepareFlag()
+                                                ? KqpValidateVolatileTx(tabletId, sysLocks, kqpLocks, useGenericReadSets, txId, op->DelayedInReadSets(), awaitingDecisions, outReadSets)
+                                                : KqpValidateLocks(tabletId, sysLocks, kqpLocks, useGenericReadSets, inReadSets);
+
+            if (!validated) {
+                writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_LOCKS_BROKEN, "Operation is aborting because locks are not valid");
+
+                for (auto& brokenLock : brokenLocks) {
+                    writeOp->GetWriteResult()->Record.MutableTxLocks()->Add()->Swap(&brokenLock);
+                }
+
+                KqpEraseLocks(tabletId, kqpLocks, sysLocks);
+                sysLocks.ApplyLocks();
+                DataShard.SubscribeNewLocks(ctx);
+                if (locksDb.HasChanges()) {
+                    op->SetWaitCompletionFlag(true);
+                    return EExecutionStatus::ExecutedNoMoreRestarts;
+                }
+                return EExecutionStatus::Executed;
+            }
+
+            auto [readVersion, writeVersion] = DataShard.GetReadWriteVersions(writeOp);
+            writeTx->SetReadVersion(readVersion);
+            writeTx->SetWriteVersion(writeVersion);
+
+            if (op->HasVolatilePrepareFlag()) {
+                writeTx->SetVolatileTxId(txId);
+            }
+
+
+            KqpCommitLocks(tabletId, kqpLocks, sysLocks, writeVersion, writeTx->GetUserDb());
+
+            DoUpdateToUserDb(writeOp, txc, ctx);
+
+            writeOp->SetWriteResult(NEvents::TDataEvents::TEvWriteResult::BuildCommited(DataShard.TabletID(), writeOp->GetTxId()));
+
+            if (Pipeline.AddLockDependencies(op, guardLocks)) {
+                writeTx->ResetCollectedChanges();
+                writeOp->ReleaseTxData(txc, ctx);
+                if (txc.DB.HasChanges()) {
+                    txc.DB.RollbackChanges();
+                }
+                return EExecutionStatus::Continue;
+            }
+
+            // Note: any transaction (e.g. immediate or non-volatile) may decide to commit as volatile due to dependencies
+            // Such transactions would have no participants and become immediately committed
+            auto commitTxIds = writeTx->GetVolatileCommitTxIds();
+            if (commitTxIds) {
+                TVector<ui64> participants(awaitingDecisions.begin(), awaitingDecisions.end());
+                DataShard.GetVolatileTxManager().PersistAddVolatileTx(
+                    txId,
+                    writeVersion,
+                    commitTxIds,
+                    writeTx->GetVolatileDependencies(),
+                    participants,
+                    writeTx->GetVolatileChangeGroup(),
+                    writeTx->GetVolatileCommitOrdered(),
+                    txc
+                );
+            }
+
+            if (op->HasVolatilePrepareFlag()) {
+                // Notify other shards about our expectations as soon as possible, even before we commit
+                for (ui64 target : op->AwaitingDecisions()) {
+                    if (DataShard.AddExpectation(target, op->GetStep(), op->GetTxId())) {
+                        DataShard.SendReadSetExpectation(ctx, op->GetStep(), op->GetTxId(), DataShard.TabletID(), target);
+                    }
+                }
+                if (!op->OutReadSets().empty()) {
+                    DataShard.PrepareAndSaveOutReadSets(op->GetStep(), op->GetTxId(), op->OutReadSets(), op->PreparedOutReadSets(), txc, ctx);
+                }
+                keepOutReadSets = true;
+            }
+
+            // Note: may erase persistent locks, must be after we persist volatile tx
+            AddLocksToResult(writeOp, ctx);
+
+            if (auto changes = std::move(writeTx->GetCollectedChanges())) {
+                op->ChangeRecords() = std::move(changes);
+            }
+
         } catch (const TNeedGlobalTxId&) {
             Y_VERIFY_S(op->GetGlobalTxId() == 0,
                 "Unexpected TNeedGlobalTxId exception for write operation with TxId# " << op->GetGlobalTxId());