Fix loading session without cursor (ydb-platform#9246)

zverevgeny · Jan 1, 2025 · eac202d · eac202d
1 parent 47e4b5d
commit eac202d
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 28 deletions.
diff --git a/ydb/core/tx/columnshard/data_sharing/manager/sessions.cpp b/ydb/core/tx/columnshard/data_sharing/manager/sessions.cpp
@@ -67,11 +67,17 @@ bool TSessionsManager::Load(NTable::TDatabase& database, const TColumnEngineForL
             NKikimrColumnShardDataSharingProto::TSourceSession protoSession;
             AFL_VERIFY(protoSession.ParseFromString(rowset.GetValue<Schema::SourceSessions::Details>()));
 
-            NKikimrColumnShardDataSharingProto::TSourceSession::TCursorDynamic protoSessionCursorDynamic;
-            AFL_VERIFY(protoSessionCursorDynamic.ParseFromString(rowset.GetValue<Schema::SourceSessions::CursorDynamic>()));
+            std::optional<NKikimrColumnShardDataSharingProto::TSourceSession::TCursorDynamic> protoSessionCursorDynamic;
+            if (rowset.HaveValue<Schema::SourceSessions::CursorDynamic>()) {
+                protoSessionCursorDynamic = NKikimrColumnShardDataSharingProto::TSourceSession::TCursorDynamic{};
+                AFL_VERIFY(protoSessionCursorDynamic->ParseFromString(rowset.GetValue<Schema::SourceSessions::CursorDynamic>()));
+            }
 
-            NKikimrColumnShardDataSharingProto::TSourceSession::TCursorStatic protoSessionCursorStatic;
-            AFL_VERIFY(protoSessionCursorStatic.ParseFromString(rowset.GetValue<Schema::SourceSessions::CursorStatic>()));
+            std::optional<NKikimrColumnShardDataSharingProto::TSourceSession::TCursorStatic> protoSessionCursorStatic;
+            if (rowset.HaveValue<Schema::SourceSessions::CursorStatic>()) {
+                protoSessionCursorStatic = NKikimrColumnShardDataSharingProto::TSourceSession::TCursorStatic{};
+                AFL_VERIFY(protoSessionCursorStatic->ParseFromString(rowset.GetValue<Schema::SourceSessions::CursorStatic>()));
+            }
 
             AFL_VERIFY(index);
             session->DeserializeFromProto(protoSession, protoSessionCursorDynamic, protoSessionCursorStatic).Validate();

diff --git a/ydb/core/tx/columnshard/data_sharing/source/session/cursor.cpp b/ydb/core/tx/columnshard/data_sharing/source/session/cursor.cpp
@@ -1,6 +1,9 @@
 #include "source.h"
-#include <ydb/core/tx/columnshard/engines/column_engine_logs.h>
+
+#include <ydb/core/tx/columnshard/columnshard_schema.h>
 #include <ydb/core/tx/columnshard/data_sharing/destination/events/transfer.h>
+#include <ydb/core/tx/columnshard/engines/column_engine_logs.h>
+
 #include <ydb/library/formats/arrow/hash/xx_hash.h>
 
 namespace NKikimr::NOlap::NDataSharing {
@@ -130,18 +133,29 @@ NKikimr::TConclusionStatus TSourceCursor::DeserializeFromProto(const NKikimrColu
         PathPortionHashes.emplace(i.GetPathId(), i.GetHash());
     }
     AFL_VERIFY(PathPortionHashes.size());
-    StaticSaved = true;
+    IsStaticSaved = true;
     return TConclusionStatus::Success();
 }
 
 TSourceCursor::TSourceCursor(const TTabletId selfTabletId, const std::set<ui64>& pathIds, const TTransferContext transferContext)
     : SelfTabletId(selfTabletId)
     , TransferContext(transferContext)
-    , PathIds(pathIds)
-{
+    , PathIds(pathIds) {
 }
 
-bool TSourceCursor::Start(const std::shared_ptr<IStoragesManager>& storagesManager, const THashMap<ui64, std::vector<std::shared_ptr<TPortionInfo>>>& portions, const TVersionedIndex& index) {
+void TSourceCursor::SaveToDatabase(NIceDb::TNiceDb& db, const TString& sessionId) {
+    using SourceSessions = NKikimr::NColumnShard::Schema::SourceSessions;
+    db.Table<SourceSessions>().Key(sessionId).Update(
+        NIceDb::TUpdate<SourceSessions::CursorDynamic>(SerializeDynamicToProto().SerializeAsString()));
+    if (!IsStaticSaved) {
+        db.Table<SourceSessions>().Key(sessionId).Update(
+            NIceDb::TUpdate<SourceSessions::CursorStatic>(SerializeStaticToProto().SerializeAsString()));
+        IsStaticSaved = true;
+    }
+}
+
+bool TSourceCursor::Start(const std::shared_ptr<IStoragesManager>& storagesManager,
+    const THashMap<ui64, std::vector<std::shared_ptr<TPortionInfo>>>& portions, const TVersionedIndex& index) {
     AFL_VERIFY(!IsStartedFlag);
     std::map<ui64, std::map<ui32, std::shared_ptr<TPortionInfo>>> local;
     std::vector<std::shared_ptr<TPortionInfo>> portionsLock;
@@ -177,5 +191,4 @@ bool TSourceCursor::Start(const std::shared_ptr<IStoragesManager>& storagesManag
     IsStartedFlag = true;
     return true;
 }
-
-}
+}   // namespace NKikimr::NOlap::NDataSharing
diff --git a/ydb/core/tx/columnshard/data_sharing/source/session/cursor.h b/ydb/core/tx/columnshard/data_sharing/source/session/cursor.h
@@ -8,6 +8,10 @@ class TColumnEngineForLogs;
 class TVersionedIndex;
 }
 
+namespace NKikimr::NIceDb {
+class TNiceDb;
+}
+
 namespace NKikimr::NOlap::NDataSharing {
 
 class TSharedBlobsManager;
@@ -30,8 +34,11 @@ class TSourceCursor {
     std::set<ui64> PathIds;
     THashMap<ui64, TString> PathPortionHashes;
     bool IsStartedFlag = false;
-    YDB_ACCESSOR(bool, StaticSaved, false);
+    bool IsStaticSaved = false;
     void BuildSelection(const std::shared_ptr<IStoragesManager>& storagesManager, const TVersionedIndex& index);
+    NKikimrColumnShardDataSharingProto::TSourceSession::TCursorDynamic SerializeDynamicToProto() const;
+    NKikimrColumnShardDataSharingProto::TSourceSession::TCursorStatic SerializeStaticToProto() const;
+
 public:
     bool IsAckDataReceived() const {
         return AckReceivedForPackIdx == PackIdx;
@@ -96,11 +103,10 @@ class TSourceCursor {
 
     TSourceCursor(const TTabletId selfTabletId, const std::set<ui64>& pathIds, const TTransferContext transferContext);
 
-    bool Start(const std::shared_ptr<IStoragesManager>& storagesManager, const THashMap<ui64, std::vector<std::shared_ptr<TPortionInfo>>>& portions, const TVersionedIndex& index);
-
-    NKikimrColumnShardDataSharingProto::TSourceSession::TCursorDynamic SerializeDynamicToProto() const;
-    NKikimrColumnShardDataSharingProto::TSourceSession::TCursorStatic SerializeStaticToProto() const;
+    void SaveToDatabase(class NIceDb::TNiceDb& db, const TString& sessionId);
 
+    bool Start(const std::shared_ptr<IStoragesManager>& storagesManager,
+        const THashMap<ui64, std::vector<std::shared_ptr<TPortionInfo>>>& portions, const TVersionedIndex& index);
     [[nodiscard]] TConclusionStatus DeserializeFromProto(const NKikimrColumnShardDataSharingProto::TSourceSession::TCursorDynamic& proto,
         const NKikimrColumnShardDataSharingProto::TSourceSession::TCursorStatic& protoStatic);
 };

diff --git a/ydb/core/tx/columnshard/data_sharing/source/session/source.cpp b/ydb/core/tx/columnshard/data_sharing/source/session/source.cpp
@@ -1,8 +1,9 @@
 #include "source.h"
-#include <ydb/core/tx/columnshard/data_sharing/source/transactions/tx_finish_ack_to_source.h>
+
+#include <ydb/core/tx/columnshard/data_locks/locks/list.h>
 #include <ydb/core/tx/columnshard/data_sharing/source/transactions/tx_data_ack_to_source.h>
+#include <ydb/core/tx/columnshard/data_sharing/source/transactions/tx_finish_ack_to_source.h>
 #include <ydb/core/tx/columnshard/data_sharing/source/transactions/tx_write_source_cursor.h>
-#include <ydb/core/tx/columnshard/data_locks/locks/list.h>
 #include <ydb/core/tx/columnshard/engines/column_engine_logs.h>
 
 namespace NKikimr::NOlap::NDataSharing {
@@ -68,6 +69,10 @@ TConclusion<std::unique_ptr<NTabletFlatExecutor::ITransaction>> TSourceSession::
     }
 }
 
+void TSourceSession::SaveCursorToDatabase(NIceDb::TNiceDb& db) {
+    GetCursorVerified()->SaveToDatabase(db, GetSessionId());
+}
+
 void TSourceSession::ActualizeDestination(const NColumnShard::TColumnShard& shard, const std::shared_ptr<NDataLocks::TManager>& dataLocksManager) {
     AFL_VERIFY(IsInProgress() || IsPrepared());
     AFL_VERIFY(Cursor);

diff --git a/ydb/core/tx/columnshard/data_sharing/source/session/source.h b/ydb/core/tx/columnshard/data_sharing/source/session/source.h
@@ -3,6 +3,10 @@
 #include <ydb/core/tx/columnshard/data_sharing/common/session/common.h>
 #include <ydb/core/tx/columnshard/common/tablet_id.h>
 
+namespace NKikimr::NIceDb {
+class TNiceDb;
+}
+
 namespace NKikimr::NOlap::NDataSharing {
 
 class TSharedBlobsManager;
@@ -58,7 +62,9 @@ class TSourceSession: public TCommonSession {
         AFL_VERIFY(!!Cursor);
         return Cursor;
     }
-/*
+
+    void SaveCursorToDatabase(NIceDb::TNiceDb& db);
+    /*
     bool TryNextCursor(const ui32 packIdx, const std::shared_ptr<IStoragesManager>& storagesManager, const TVersionedIndex& index) {
         AFL_VERIFY(Cursor);
         if (packIdx != Cursor->GetPackIdx()) {

diff --git a/ydb/core/tx/columnshard/data_sharing/source/transactions/tx_data_ack_to_source.cpp b/ydb/core/tx/columnshard/data_sharing/source/transactions/tx_data_ack_to_source.cpp
@@ -22,13 +22,7 @@ bool TTxDataAckToSource::DoExecute(NTabletFlatExecutor::TTransactionContext& txc
     }
 
     NIceDb::TNiceDb db(txc.DB);
-    db.Table<Schema::SourceSessions>().Key(Session->GetSessionId())
-        .Update(NIceDb::TUpdate<Schema::SourceSessions::CursorDynamic>(Session->GetCursorVerified()->SerializeDynamicToProto().SerializeAsString()));
-    if (!Session->GetCursorVerified()->GetStaticSaved()) {
-        db.Table<Schema::SourceSessions>().Key(Session->GetSessionId())
-            .Update(NIceDb::TUpdate<Schema::SourceSessions::CursorStatic>(Session->GetCursorVerified()->SerializeStaticToProto().SerializeAsString()));
-        Session->GetCursorVerified()->SetStaticSaved(true);
-    }
+    Session->SaveCursorToDatabase(db);
     std::swap(SharedBlobIds, sharedTabletBlobIds);
     return true;
 }

diff --git a/ydb/core/tx/columnshard/data_sharing/source/transactions/tx_write_source_cursor.cpp b/ydb/core/tx/columnshard/data_sharing/source/transactions/tx_write_source_cursor.cpp
@@ -6,8 +6,7 @@ namespace NKikimr::NOlap::NDataSharing {
 bool TTxWriteSourceCursor::DoExecute(NTabletFlatExecutor::TTransactionContext& txc, const TActorContext& /*ctx*/) {
     using namespace NColumnShard;
     NIceDb::TNiceDb db(txc.DB);
-    db.Table<Schema::SourceSessions>().Key(Session->GetSessionId())
-        .Update(NIceDb::TUpdate<Schema::SourceSessions::CursorDynamic>(Session->GetCursorVerified()->SerializeDynamicToProto().SerializeAsString()));
+    Session->SaveCursorToDatabase(db);
     return true;
 }