Add OLAP support with topics in BufferWriteActor (ydb-platform#12686)

pixcc · Dec 25, 2024 · 97167eb · 97167eb
1 parent 308501b
commit 97167eb
Show file tree

Hide file tree

Showing 4 changed files with 211 additions and 46 deletions.
diff --git a/ydb/core/kqp/runtime/kqp_write_actor.cpp b/ydb/core/kqp/runtime/kqp_write_actor.cpp
@@ -1419,7 +1419,6 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
         State = EState::WRITING;
         Alloc->Release();
         Counters->BufferActorsCount->Inc();
-        TxManager->AddTopicsToShards();
     }
 
     void Bootstrap() {

diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -1378,6 +1378,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             && !txCtx->BufferActorId
             && (txCtx->HasTableWrite || request.TopicOperations.GetSize() != 0)) {
             txCtx->TxManager->SetTopicOperations(std::move(request.TopicOperations));
+            txCtx->TxManager->AddTopicsToShards();
 
             TKqpBufferWriterSettings settings {
                 .SessionActorId = SelfId(),
@@ -1388,6 +1389,9 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             };
             auto* actor = CreateKqpBufferWriterActor(std::move(settings));
             txCtx->BufferActorId = RegisterWithSameMailbox(actor);
+        } else if (Settings.TableService.GetEnableOltpSink() && txCtx->BufferActorId) {
+            txCtx->TxManager->SetTopicOperations(std::move(request.TopicOperations));
+            txCtx->TxManager->AddTopicsToShards();
         }
 
         auto executerActor = CreateKqpExecuter(std::move(request), Settings.Database,

diff --git a/ydb/core/testlib/test_client.h b/ydb/core/testlib/test_client.h
@@ -231,6 +231,18 @@ namespace Tests {
             AppConfig->MutableTableServiceConfig()->SetEnableOltpSink(withOltpSink);
             return *this;
         }
+        TServerSettings& SetEnableOlapSink(bool withOlapSink) {
+            AppConfig->MutableTableServiceConfig()->SetEnableOlapSink(withOlapSink);
+            return *this;
+        }
+        TServerSettings& SetEnableHtapTx(bool withHtapTx) {
+            AppConfig->MutableTableServiceConfig()->SetEnableHtapTx(withHtapTx);
+            return *this;
+        }
+        TServerSettings& SetAllowOlapDataQuery(bool withAllowOlapDataQuery) {
+            AppConfig->MutableTableServiceConfig()->SetAllowOlapDataQuery(withAllowOlapDataQuery);
+            return *this;
+        }
 
 
         // Add additional grpc services