From ce411192d2dfda7f849bc194da30d98252328785 Mon Sep 17 00:00:00 2001
From: Dmitry Kardymon <kardymon-d@ydb.tech>
Date: Tue, 26 Dec 2023 16:44:32 +0300
Subject: [PATCH] Add discovery mutator for underlay (#566)

* Add discovery mutator for underlay

* Fix prefix

* TString -> StringBuf
---
 .../shared_resources/shared_resources.cpp     | 28 ++++++++++++++++++-
 ydb/core/fq/libs/shared_resources/ya.make     |  1 +
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/ydb/core/fq/libs/shared_resources/shared_resources.cpp b/ydb/core/fq/libs/shared_resources/shared_resources.cpp
index 73e8ea4be1d6..071e70288620 100644
--- a/ydb/core/fq/libs/shared_resources/shared_resources.cpp
+++ b/ydb/core/fq/libs/shared_resources/shared_resources.cpp
@@ -1,11 +1,15 @@
 #include "shared_resources.h"
 
-#include <ydb/library/services/services.pb.h>
 #include <ydb/core/fq/libs/events/events.h>
+#include <ydb/library/actors/core/actorsystem.h>
 #include <ydb/library/logger/actor.h>
+#include <ydb/library/services/services.pb.h>
 
 #include <library/cpp/actors/core/actorsystem.h>
 
+#include <ydb/public/api/protos/ydb_discovery.pb.h>
+#include <ydb/public/sdk/cpp/client/extensions/discovery_mutator/discovery_mutator.h>
+
 #include <util/generic/cast.h>
 #include <util/generic/strbuf.h>
 #include <util/stream/file.h>
@@ -33,6 +37,7 @@ struct TYqSharedResourcesImpl : public TActorSystemPtrMixin, public TYqSharedRes
         : TYqSharedResources(NYdb::TDriver(GetYdbDriverConfig(config.GetCommon().GetYdbDriverConfig())))
     {
         CreateDbPoolHolder(PrepareDbPoolConfig(config), credentialsProviderFactory, counters);
+        AddUnderlayDiscoveryMutator();
     }
 
     NDbPool::TConfig PrepareDbPoolConfig(const NFq::NConfig::TConfig& config) {
@@ -81,6 +86,27 @@ struct TYqSharedResourcesImpl : public TActorSystemPtrMixin, public TYqSharedRes
         const ::NMonitoring::TDynamicCounterPtr& counters) {
         DbPoolHolder = MakeIntrusive<NDbPool::TDbPoolHolder>(config, CoreYdbDriver, credentialsProviderFactory, counters);
     }
+
+    void AddUnderlayDiscoveryMutator() {
+
+        auto mutator = [](Ydb::Discovery::ListEndpointsResult* proto, NYdb::TStatus status, const NYdb::IDiscoveryMutatorApi::TAuxInfo& aux) {
+            TStringBuf underlayPrefix{"u-"};
+            if (!aux.DiscoveryEndpoint.starts_with(underlayPrefix) || !proto) {
+                return status;
+            }
+
+            for (size_t i = 0; i < proto->endpointsSize(); ++i) {
+                Ydb::Discovery::EndpointInfo* endpointInfo = proto->Mutableendpoints(i);
+                const TString& address = endpointInfo->address();
+                if (address.StartsWith(underlayPrefix)) {
+                    continue;
+                }
+                endpointInfo->set_address(underlayPrefix + address);
+            }
+            return status;
+        };
+        UserSpaceYdbDriver.AddExtension<NDiscoveryMutator::TDiscoveryMutator>(NDiscoveryMutator::TDiscoveryMutator::TParams(std::move(mutator)));
+    }
 };
 
 } // namespace
diff --git a/ydb/core/fq/libs/shared_resources/ya.make b/ydb/core/fq/libs/shared_resources/ya.make
index e278b408bd0b..d4348a4ea963 100644
--- a/ydb/core/fq/libs/shared_resources/ya.make
+++ b/ydb/core/fq/libs/shared_resources/ya.make
@@ -21,6 +21,7 @@ PEERDIR(
     ydb/library/logger
     ydb/library/security
     ydb/public/sdk/cpp/client/ydb_driver
+    ydb/public/sdk/cpp/client/ydb_extension
     ydb/public/sdk/cpp/client/ydb_table
 )