Pass DiscoveryEndpoint in to DiscoveryMutator

Pass DiscoveryEndpoint in to DiscoveryMutator

Pull Request resolved: https://github.com/ydb-platform/ydb/pull/504

4 files changed, 21 insertions, 10 deletions

diff --git a/ydb/public/sdk/cpp/client/extensions/discovery_mutator/discovery_mutator_ut.cpp b/ydb/public/sdk/cpp/client/extensions/discovery_mutator/discovery_mutator_ut.cpp
index 7a5fa43cc3..5da6e26a21 100644
--- a/ydb/public/sdk/cpp/client/extensions/discovery_mutator/discovery_mutator_ut.cpp
+++ b/ydb/public/sdk/cpp/client/extensions/discovery_mutator/discovery_mutator_ut.cpp
@@ -12,18 +12,20 @@ using namespace NDiscoveryMutator;
 Y_UNIT_TEST_SUITE(DiscoveryMutator) {
     Y_UNIT_TEST(Simple) {
 
-        std::unordered_set<TString> dbs;
+        std::unordered_set<std::string_view> dbs;
+        TString discoveryEndpont = "localhost:100";
 
-        auto mutator = [&dbs](Ydb::Discovery::ListEndpointsResult* proto, TStatus status, const TString& database) {
-            UNIT_ASSERT_VALUES_EQUAL("localhost:100", status.GetEndpoint());
-            dbs.insert(database);
+        auto mutator = [&](Ydb::Discovery::ListEndpointsResult* proto, TStatus status, const IDiscoveryMutatorApi::TAuxInfo& aux) {
+            UNIT_ASSERT_VALUES_EQUAL(discoveryEndpont, status.GetEndpoint());
+            UNIT_ASSERT_VALUES_EQUAL(discoveryEndpont, aux.DiscoveryEndpoint);
+            dbs.insert(aux.Database);
             Y_UNUSED(proto);
             return status;
         };
         auto driver = TDriver(
             TDriverConfig()
                 .SetDatabase("db1")
-                .SetEndpoint("localhost:100"));
+                .SetEndpoint(discoveryEndpont));
 
         driver.AddExtension<TDiscoveryMutator>(TDiscoveryMutator::TParams(std::move(mutator)));
 
diff --git a/ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.cpp b/ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.cpp
index d855b2af08..7eec444eb5 100644
--- a/ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.cpp
+++ b/ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.cpp
@@ -346,11 +346,11 @@ TAsyncListEndpointsResult TGRpcConnectionsImpl::GetEndpoints(TDbDriverStatePtr d
         if (strong && result.DiscoveryStatus.IsTransportError()) {
             strong->StatCollector.IncDiscoveryFailDueTransportError();
         }
-        return NThreading::MakeFuture<TListEndpointsResult>(MutateDiscovery(std::move(result), strong->Database));
+        return NThreading::MakeFuture<TListEndpointsResult>(MutateDiscovery(std::move(result), *strong));
     });
 }
 
-TListEndpointsResult TGRpcConnectionsImpl::MutateDiscovery(TListEndpointsResult result, const TStringType& database) {
+TListEndpointsResult TGRpcConnectionsImpl::MutateDiscovery(TListEndpointsResult result, const TDbDriverState& dbDriverState) {
     std::lock_guard lock(ExtensionsLock_);
     if (!DiscoveryMutatorCb)
         return result;
@@ -358,7 +358,12 @@ TListEndpointsResult TGRpcConnectionsImpl::MutateDiscovery(TListEndpointsResult
     auto endpoint = result.DiscoveryStatus.Endpoint;
     auto ydbStatus = NYdb::TStatus(std::move(result.DiscoveryStatus));
 
-    ydbStatus = DiscoveryMutatorCb(&result.Result, std::move(ydbStatus), database);
+    auto aux = IDiscoveryMutatorApi::TAuxInfo {
+        .Database = dbDriverState.Database,
+        .DiscoveryEndpoint = dbDriverState.DiscoveryEndpoint
+    };
+
+    ydbStatus = DiscoveryMutatorCb(&result.Result, std::move(ydbStatus), aux);
 
     auto issues = ydbStatus.GetIssues();
 
diff --git a/ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.h b/ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.h
index b57b36e1d7..72887d2baf 100644
--- a/ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.h
+++ b/ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.h
@@ -578,7 +578,7 @@ public:
     }
 
     TAsyncListEndpointsResult GetEndpoints(TDbDriverStatePtr dbState) override;
-    TListEndpointsResult MutateDiscovery(TListEndpointsResult result, const TStringType& database);
+    TListEndpointsResult MutateDiscovery(TListEndpointsResult result, const TDbDriverState& dbDriverState);
 
 #ifndef YDB_GRPC_BYPASS_CHANNEL_POOL
     void DeleteChannels(const std::vector<std::string>& endpoints) override {
diff --git a/ydb/public/sdk/cpp/client/ydb_extension/extension.h b/ydb/public/sdk/cpp/client/ydb_extension/extension.h
index 1dd8589501..2a00b7875f 100644
--- a/ydb/public/sdk/cpp/client/ydb_extension/extension.h
+++ b/ydb/public/sdk/cpp/client/ydb_extension/extension.h
@@ -48,7 +48,11 @@ class DestroyedClientException: public yexception {};
 
 class IDiscoveryMutatorApi: public IExtensionApi {
 public:
-    using TMutatorCb = std::function<TStatus(Ydb::Discovery::ListEndpointsResult* proto, TStatus status, const TStringType& database)>;
+    struct TAuxInfo {
+        std::string_view Database;
+        std::string_view DiscoveryEndpoint;
+    };
+    using TMutatorCb = std::function<TStatus(Ydb::Discovery::ListEndpointsResult* proto, TStatus status, const TAuxInfo& aux)>;
 
     static IDiscoveryMutatorApi* Create(TDriver driver);
 public: