Intermediate changes

3 files changed, 22 insertions, 0 deletions

diff --git a/yt/yt/client/federated/client.cpp b/yt/yt/client/federated/client.cpp
index 0a52560296..c5b200d409 100644
--- a/yt/yt/client/federated/client.cpp
+++ b/yt/yt/client/federated/client.cpp
@@ -5,6 +5,7 @@
 
 #include <yt/yt/client/api/client.h>
 #include <yt/yt/client/api/transaction.h>
+#include <yt/yt/client/api/rpc_proxy/transaction_impl.h>
 #include <yt/yt/client/api/dynamic_table_transaction_mixin.h>
 #include <yt/yt/client/api/queue_transaction_mixin.h>
 
@@ -58,6 +59,7 @@ class TTransaction
     : public virtual ITransaction
     , public TDynamicTableTransactionMixin
     , public TQueueTransactionMixin
+    , public IFederatedClientTransactionMixin
 {
 public:
     TTransaction(TClientPtr client, int clientIndex, ITransactionPtr underlying);
@@ -197,6 +199,12 @@ public:
         Underlying_->UnsubscribeAborted(handler);
     }
 
+    std::optional<TString> TryGetStickyProxyAddress() const override
+    {
+        const auto* derived = Underlying_->TryAs<NRpcProxy::TTransaction>();
+        return nullptr == derived ? std::nullopt : std::make_optional(derived->GetStickyProxyAddress());
+    }
+
     UNIMPLEMENTED_METHOD(TFuture<void>, SetNode, (const NYPath::TYPath&, const NYson::TYsonString&, const TSetNodeOptions&));
     UNIMPLEMENTED_METHOD(TFuture<void>, MultisetAttributesNode, (const NYPath::TYPath&, const NYTree::IMapNodePtr&, const TMultisetAttributesNodeOptions&));
     UNIMPLEMENTED_METHOD(TFuture<void>, RemoveNode, (const NYPath::TYPath&, const TRemoveNodeOptions&));
diff --git a/yt/yt/client/federated/client.h b/yt/yt/client/federated/client.h
index 1028ab98c3..4a519486bb 100644
--- a/yt/yt/client/federated/client.h
+++ b/yt/yt/client/federated/client.h
@@ -20,6 +20,14 @@ namespace NYT::NClient::NFederated {
 
 ////////////////////////////////////////////////////////////////////////////////
 
+struct IFederatedClientTransactionMixin
+{
+    //! Try fetch sticky proxy address if underlying transaction supports it.
+    virtual std::optional<TString> TryGetStickyProxyAddress() const = 0;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
 //! Creates a federated client with given underlying clients.
 NApi::IClientPtr CreateClient(
     std::vector<NApi::IClientPtr> clients,
diff --git a/yt/yt/client/federated/unittests/client_ut.cpp b/yt/yt/client/federated/unittests/client_ut.cpp
index 8588e6dc66..0bc009fb6f 100644
--- a/yt/yt/client/federated/unittests/client_ut.cpp
+++ b/yt/yt/client/federated/unittests/client_ut.cpp
@@ -298,6 +298,12 @@ TEST(TFederatedClientTest, Transactions)
 
     auto transaction = federatedClient->StartTransaction(NTransactionClient::ETransactionType::Tablet).Get().Value();
 
+    // Check mock transaction doesn't work with sticky proxy address.
+    {
+        auto* derived = transaction->As<IFederatedClientTransactionMixin>();
+        ASSERT_EQ(std::nullopt, derived->TryGetStickyProxyAddress());
+    }
+
     // From `vla`.
     {
         auto result = transaction->LookupRows(data.Path, data.NameTable, data.Keys);