Enhance token resolution

commit_hash:7f10df1e62935e3145c9bd8312b2e44ef8bf0bda

28 files changed, 186 insertions, 49 deletions

diff --git a/yql/essentials/core/facade/yql_facade.cpp b/yql/essentials/core/facade/yql_facade.cpp
index fefcad7f579..32c1e6b6094 100644
--- a/yql/essentials/core/facade/yql_facade.cpp
+++ b/yql/essentials/core/facade/yql_facade.cpp
@@ -2216,7 +2216,13 @@ TTypeAnnotationContextPtr TProgram::BuildTypeAnnotationContext(const TString& us
     }
 
     tokenResolvers.push_back(BuildDefaultTokenResolver(typeAnnotationContext->Credentials));
-    typeAnnotationContext->UserDataStorage->SetTokenResolver(BuildCompositeTokenResolver(std::move(tokenResolvers)));
+    auto tokenResolver = BuildCompositeTokenResolver(std::move(tokenResolvers));
+
+    typeAnnotationContext->UserDataStorage->SetTokenResolver(tokenResolver);
+
+    if (auto* urlListerManager = typeAnnotationContext->UrlListerManager.Get()) {
+        urlListerManager->SetTokenResolver(std::move(tokenResolver));
+    }
 
     return typeAnnotationContext;
 }
diff --git a/yql/essentials/core/qplayer/url_lister/qplayer_url_lister_manager.cpp b/yql/essentials/core/qplayer/url_lister/qplayer_url_lister_manager.cpp
index 3e1ffcf187b..ff83bdc46be 100644
--- a/yql/essentials/core/qplayer/url_lister/qplayer_url_lister_manager.cpp
+++ b/yql/essentials/core/qplayer/url_lister/qplayer_url_lister_manager.cpp
@@ -114,6 +114,13 @@ public:
         Underlying_->SetCredentials(credentials);
     }
 
+    void SetTokenResolver(std::function<TString(const TString&, const TString&)> tokenResolver) override {
+        if (QContext_.CanRead()) {
+            return;
+        }
+        Underlying_->SetTokenResolver(std::move(tokenResolver));
+    }
+
     void SetUrlPreprocessing(IUrlPreprocessing::TPtr urlPreprocessing) override {
         if (QContext_.CanRead()) {
             return;
diff --git a/yql/essentials/core/type_ann/type_ann_core.cpp b/yql/essentials/core/type_ann/type_ann_core.cpp
index 6338cf67adb..66e650008d8 100644
--- a/yql/essentials/core/type_ann/type_ann_core.cpp
+++ b/yql/essentials/core/type_ann/type_ann_core.cpp
@@ -8101,10 +8101,9 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
                 });
 
                 for (const auto& x : ctx.Types.DataSources) {
-                    auto tokens = x->GetClusterTokens();
-                    if (tokens) {
-                        for (const auto& t : *tokens) {
-                            description.SecureParams.insert(std::make_pair(TString("cluster:default_") + t.first, t.second));
+                    for (const auto& clusterName : x->GetValidClusters()) {
+                        if (auto token = x->ResolveClusterToken(clusterName)) {
+                            description.SecureParams.emplace(TString("cluster:default_") + TString(clusterName), *token);
                         }
                     }
                 }
@@ -12011,18 +12010,14 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
                     TStringBuf clusterName = p1;
                     if (clusterName.SkipPrefix("default_")) {
                         for (auto& x : ctx.Types.DataSources) {
-                            auto tokens = x->GetClusterTokens();
-                            auto token = tokens ? tokens->FindPtr(clusterName) : nullptr;
-                            if (token) {
+                            if (auto token = x->ResolveClusterToken(TString(clusterName))) {
                                 clusterCred.ConstructInPlace(TString(x->GetName()), "", *token);
                                 cred = clusterCred.Get();
                                 break;
                             }
                         }
                         for (auto& x : ctx.Types.DataSinks) {
-                            auto tokens = x->GetClusterTokens();
-                            auto token = tokens ? tokens->FindPtr(clusterName) : nullptr;
-                            if (token) {
+                            if (auto token = x->ResolveClusterToken(TString(clusterName))) {
                                 clusterCred.ConstructInPlace(TString(x->GetName()), "", *token);
                                 cred = clusterCred.Get();
                                 break;
diff --git a/yql/essentials/core/url_lister/interface/url_lister_manager.h b/yql/essentials/core/url_lister/interface/url_lister_manager.h
index ec8f88f13cc..3487ac3eb17 100644
--- a/yql/essentials/core/url_lister/interface/url_lister_manager.h
+++ b/yql/essentials/core/url_lister/interface/url_lister_manager.h
@@ -8,6 +8,9 @@
 #include <yql/essentials/core/url_preprocessing/interface/url_preprocessing.h>
 
 #include <util/generic/ptr.h>
+#include <util/generic/string.h>
+
+#include <functional>
 
 namespace NYql {
 
@@ -20,6 +23,7 @@ public:
     virtual TIntrusivePtr<IUrlListerManager> Clone() const = 0;
 
     virtual void SetCredentials(TCredentials::TPtr credentials) = 0;
+    virtual void SetTokenResolver(std::function<TString(const TString&, const TString&)> tokenResolver) = 0;
     virtual void SetUrlPreprocessing(IUrlPreprocessing::TPtr urlPreprocessing) = 0;
     virtual void SetParameters(const NYT::TNode& parameters) = 0;
 };
diff --git a/yql/essentials/core/url_lister/url_lister_manager.cpp b/yql/essentials/core/url_lister/url_lister_manager.cpp
index 6aa9af35638..d16a4592e81 100644
--- a/yql/essentials/core/url_lister/url_lister_manager.cpp
+++ b/yql/essentials/core/url_lister/url_lister_manager.cpp
@@ -22,7 +22,7 @@ public:
     TVector<TUrlListEntry> ListUrl(const TString& url, const TString& tokenName) const override {
         auto [preprocessedUrl, alias] = GetPreparedUrlAndAlias(url);
 
-        TString token = GetToken(tokenName, alias);
+        TString token = GetToken(tokenName, preprocessedUrl, alias);
 
         for (const auto& urlLister : UrlListers_) {
             if (urlLister->Accept(preprocessedUrl)) {
@@ -36,7 +36,7 @@ public:
     TVector<TUrlListEntry> ListUrlRecursive(const TString& url, const TString& tokenName, const TString& separator, ui32 foldersLimit) const override {
         auto [preprocessedUrl, alias] = GetPreparedUrlAndAlias(url);
 
-        TString token = GetToken(tokenName, alias);
+        TString token = GetToken(tokenName, preprocessedUrl, alias);
 
         for (const auto& urlLister : UrlListers_) {
             if (urlLister->Accept(preprocessedUrl)) {
@@ -64,6 +64,10 @@ public:
         Credentials_ = std::move(credentials);
     }
 
+    void SetTokenResolver(std::function<TString(const TString&, const TString&)> tokenResolver) override {
+        TokenResolver_ = std::move(tokenResolver);
+    }
+
     void SetUrlPreprocessing(IUrlPreprocessing::TPtr urlPreprocessing) override {
         UrlPreprocessing_ = std::move(urlPreprocessing);
     }
@@ -85,7 +89,7 @@ private:
         return {preprocessedUrl, alias};
     }
 
-    TString GetToken(const TString& tokenName, const TString& alias) const {
+    TString GetToken(const TString& tokenName, const TString& preprocessedUrl, const TString& alias) const {
         TMaybe<TString> token;
 
         if (tokenName) {
@@ -101,10 +105,12 @@ private:
             token = credential->Content;
         }
 
-        if (!token && alias && Credentials_) {
-            if (auto credential = Credentials_->FindCredential("default_" + alias)) {
-                token = credential->Content;
+        if (!token && preprocessedUrl && alias) {
+            if (!TokenResolver_) {
+                throw yexception() << "Missing token resolver";
             }
+
+            token = TokenResolver_(preprocessedUrl, alias);
         }
 
         return *token.OrElse("");
@@ -156,6 +162,7 @@ private:
     TVector<IUrlListerPtr> UrlListers_;
 
     TCredentials::TPtr Credentials_;
+    std::function<TString(const TString&, const TString&)> TokenResolver_;
     IUrlPreprocessing::TPtr UrlPreprocessing_;
     TMaybe<NYT::TNode> Parameters_;
 };
diff --git a/yql/essentials/core/yql_data_provider.h b/yql/essentials/core/yql_data_provider.h
index a60337bac7e..5f424481f57 100644
--- a/yql/essentials/core/yql_data_provider.h
+++ b/yql/essentials/core/yql_data_provider.h
@@ -116,6 +116,8 @@ public:
     virtual IGraphTransformer& GetConfigurationTransformer() = 0;
     virtual TExprNode::TPtr GetClusterInfo(const TString& cluster, TExprContext& ctx) = 0;
     virtual const THashMap<TString, TString>* GetClusterTokens() = 0;
+    virtual TMaybe<TString> ResolveClusterToken(const TString& cluster) = 0;
+    virtual const THashSet<TString>& GetValidClusters() = 0;
     virtual void AddCluster(const TString& name, const THashMap<TString, TString>& properties) = 0;
 
     //-- discovery & rewrite
diff --git a/yql/essentials/providers/common/config/yql_dispatch.h b/yql/essentials/providers/common/config/yql_dispatch.h
index a99d72f8341..a22137a4d8b 100644
--- a/yql/essentials/providers/common/config/yql_dispatch.h
+++ b/yql/essentials/providers/common/config/yql_dispatch.h
@@ -369,6 +369,12 @@ public:
         ValidClusters.insert(cluster);
     }
 
+    bool IsValidCluster(const TString& cluster) const {
+        return ValidClusters.contains(cluster);
+    }
+
+    const THashSet<TString>& GetValidClusters() const;
+
     template <typename TType, EConfSettingType SettingType>
     TSettingHandlerImpl<TType, SettingType>& AddSetting(const TString& name, TConfSetting<TType, SettingType>& setting) {
         TIntrusivePtr<TSettingHandlerImpl<TType, SettingType>> handler = new TSettingHandlerImpl<TType, SettingType>(name, setting);
@@ -422,9 +428,6 @@ public:
     void Enumerate(std::function<void(std::string_view)> callback);
 
 protected:
-    // FIXME switch usages to an acesssor
-    const THashSet<TString>& GetValidClusters() const;
-
     THashSet<TString> ValidClusters; // NOLINT(readability-identifier-naming)
     THashMap<TString, TSettingHandler::TPtr> Handlers_;
     TSet<TString> Names_;
diff --git a/yql/essentials/providers/common/provider/yql_data_provider_impl.cpp b/yql/essentials/providers/common/provider/yql_data_provider_impl.cpp
index 21cd47329b7..c40f8131b4a 100644
--- a/yql/essentials/providers/common/provider/yql_data_provider_impl.cpp
+++ b/yql/essentials/providers/common/provider/yql_data_provider_impl.cpp
@@ -120,6 +120,32 @@ const THashMap<TString, TString>* TDataProviderBase::GetClusterTokens() {
     return nullptr;
 }
 
+TMaybe<TString> TDataProviderBase::ResolveClusterToken(const TString& cluster) {
+    if (auto* tokens = GetClusterTokens()) {
+        if (auto* token = tokens->FindPtr(cluster)) {
+            return *token;
+        }
+    }
+
+    return {};
+}
+
+// TODO: drop this compatibility implementation once all descendants
+// provide their own overloads
+const THashSet<TString>& TDataProviderBase::GetValidClusters() {
+    if (ValidClusters_) {
+        return ValidClusters_;
+    }
+
+    if (auto* tokens = GetClusterTokens()) {
+        for (const auto& [clusterName, _] : *tokens) {
+            ValidClusters_.emplace(clusterName);
+        }
+    }
+
+    return ValidClusters_;
+}
+
 IGraphTransformer& TDataProviderBase::GetIODiscoveryTransformer() {
     return NullTransformer_;
 }
diff --git a/yql/essentials/providers/common/provider/yql_data_provider_impl.h b/yql/essentials/providers/common/provider/yql_data_provider_impl.h
index 99ae95bdfa0..fdae849126f 100644
--- a/yql/essentials/providers/common/provider/yql_data_provider_impl.h
+++ b/yql/essentials/providers/common/provider/yql_data_provider_impl.h
@@ -49,6 +49,8 @@ public:
     TExprNode::TPtr GetClusterInfo(const TString& cluster, TExprContext& ctx) override;
     void AddCluster(const TString& name, const THashMap<TString, TString>& properties) override;
     const THashMap<TString, TString>* GetClusterTokens() override;
+    TMaybe<TString> ResolveClusterToken(const TString& cluster) override;
+    const THashSet<TString>& GetValidClusters() override;
     IGraphTransformer& GetIODiscoveryTransformer() override;
     IGraphTransformer& GetEpochsTransformer() override;
     IGraphTransformer& GetIntentDeterminationTransformer() override;
@@ -102,6 +104,9 @@ protected:
     THolder<IGraphTransformer> DefConstraintTransformer_;
     TNullTransformer NullTransformer_;
     TTrackableNodeProcessorBase NullTrackableNodeProcessor_;
+
+    // TODO: remove after overriding GetValidClusters method in all descendants
+    THashSet<TString> ValidClusters_;
 };
 
 TExprNode::TPtr DefaultCleanupWorld(const TExprNode::TPtr& node, TExprContext& ctx);
diff --git a/yql/essentials/providers/common/provider/yql_provider.cpp b/yql/essentials/providers/common/provider/yql_provider.cpp
index 8d4bd13e6e9..1c7acd747fd 100644
--- a/yql/essentials/providers/common/provider/yql_provider.cpp
+++ b/yql/essentials/providers/common/provider/yql_provider.cpp
@@ -1101,23 +1101,15 @@ void GetToken(const TString& string, TString& out, const TTypeAnnotationContext&
                 TStringBuf clusterName = p1;
                 if (clusterName.SkipPrefix("default_")) {
                     for (auto& x : type.DataSources) {
-                        auto tokens = x->GetClusterTokens();
-                        if (tokens) {
-                            auto token = tokens->FindPtr(clusterName);
-                            if (token) {
-                                out = *token;
-                                return;
-                            }
+                        if (auto token = x->ResolveClusterToken(TString(clusterName))) {
+                            out = *token;
+                            return;
                         }
                     }
                     for (auto& x : type.DataSinks) {
-                        auto tokens = x->GetClusterTokens();
-                        if (tokens) {
-                            auto token = tokens->FindPtr(clusterName);
-                            if (token) {
-                                out = *token;
-                                return;
-                            }
+                        if (auto token = x->ResolveClusterToken(TString(clusterName))) {
+                            out = *token;
+                            return;
                         }
                     }
                 }
diff --git a/yt/yql/providers/yt/gateway/file/yql_yt_file.cpp b/yt/yql/providers/yt/gateway/file/yql_yt_file.cpp
index 7b86e66caa6..be2d4683f94 100644
--- a/yt/yql/providers/yt/gateway/file/yql_yt_file.cpp
+++ b/yt/yql/providers/yt/gateway/file/yql_yt_file.cpp
@@ -1648,6 +1648,10 @@ private:
         return MakeFuture<IYtGateway::TDownloadTableResult>();
     }
 
+    IYtTokenResolver::TPtr GetYtTokenResolver() const override {
+        return nullptr;
+    }
+
 private:
     TYtFileServices::TPtr Services_;
     TIntrusivePtr<NCommon::TMkqlCommonCallableCompiler> MkqlCompiler_;
diff --git a/yt/yql/providers/yt/gateway/lib/exec_ctx.cpp b/yt/yql/providers/yt/gateway/lib/exec_ctx.cpp
index 41f094d156d..ebfaf156460 100644
--- a/yt/yql/providers/yt/gateway/lib/exec_ctx.cpp
+++ b/yt/yql/providers/yt/gateway/lib/exec_ctx.cpp
@@ -389,6 +389,12 @@ TString TExecContextBaseSimple::GetAuth(const TYtSettings::TConstPtr& config) co
          auth = Clusters_->GetAuth(Cluster_);
     }
 
+    if (!auth || auth->empty()) {
+        if (auto ytTokenResolver = Gateway->GetYtTokenResolver()) {
+            auth = ytTokenResolver->ResolveClusterToken(Cluster_);
+        }
+    }
+
     return auth.GetOrElse(TString());
 }
 
diff --git a/yt/yql/providers/yt/gateway/native/yql_yt_native.cpp b/yt/yql/providers/yt/gateway/native/yql_yt_native.cpp
index 1044756c248..e4d2a56af2e 100644
--- a/yt/yql/providers/yt/gateway/native/yql_yt_native.cpp
+++ b/yt/yql/providers/yt/gateway/native/yql_yt_native.cpp
@@ -5919,6 +5919,13 @@ private:
                 if (!auth || auth->empty()) {
                     auth = Clusters_->GetAuth(options.Cluster());
                 }
+
+                if (!auth && Services_.YtTokenResolver) {
+                    if (auto token = Services_.YtTokenResolver->ResolveClusterToken(options.Cluster())) {
+                        auth = *token;
+                    }
+                }
+
                 clusterConnectionResult.Token = auth;
             }
             clusterConnectionResult.SetSuccess();
@@ -6168,6 +6175,10 @@ private:
         });
     }
 
+    IYtTokenResolver::TPtr GetYtTokenResolver() const override {
+        return Services_.YtTokenResolver;
+    }
+
 private:
     const TYtNativeServices Services_;
     const TConfigClusters::TPtr Clusters_;
diff --git a/yt/yql/providers/yt/gateway/native/yql_yt_native.h b/yt/yql/providers/yt/gateway/native/yql_yt_native.h
index b8f10227855..93f94df0fa2 100644
--- a/yt/yql/providers/yt/gateway/native/yql_yt_native.h
+++ b/yt/yql/providers/yt/gateway/native/yql_yt_native.h
@@ -3,6 +3,7 @@
 #include <yt/yql/providers/yt/gateway/lib/exec_ctx.h>
 #include <yt/yql/providers/yt/provider/yql_yt_gateway.h>
 #include <yt/yql/providers/yt/lib/secret_masker/secret_masker.h>
+#include <yt/yql/providers/yt/lib/yt_token_resolver/yt_token_resolver.h>
 
 #include <yql/essentials/core/file_storage/file_storage.h>
 #include <yql/essentials/minikql/mkql_function_registry.h>
@@ -19,6 +20,7 @@ struct TYtNativeServices: public TYtBaseServices {
 
     IMetricsRegistryPtr Metrics;
     ISecretMasker::TPtr SecretMasker;
+    IYtTokenResolver::TPtr YtTokenResolver;
 };
 
 IYtGateway::TPtr CreateYtNativeGateway(const TYtNativeServices& services);
diff --git a/yt/yql/providers/yt/gateway/qplayer/yql_yt_qplayer_gateway.cpp b/yt/yql/providers/yt/gateway/qplayer/yql_yt_qplayer_gateway.cpp
index 2fb87df5779..1fa0259c865 100644
--- a/yt/yql/providers/yt/gateway/qplayer/yql_yt_qplayer_gateway.cpp
+++ b/yt/yql/providers/yt/gateway/qplayer/yql_yt_qplayer_gateway.cpp
@@ -1199,6 +1199,10 @@ private:
         return Inner_->DownloadTable(std::move(options));
     }
 
+    IYtTokenResolver::TPtr GetYtTokenResolver() const override {
+        return Inner_->GetYtTokenResolver();
+    }
+
 private:
     const IYtGateway::TPtr Inner_;
     const TQContext QContext_;
diff --git a/yt/yql/providers/yt/lib/yt_token_resolver/ya.make b/yt/yql/providers/yt/lib/yt_token_resolver/ya.make
new file mode 100644
index 00000000000..d1c888441cb
--- /dev/null
+++ b/yt/yql/providers/yt/lib/yt_token_resolver/ya.make
@@ -0,0 +1,7 @@
+LIBRARY()
+
+SRCS(
+    yt_token_resolver.cpp
+)
+
+END()
diff --git a/yt/yql/providers/yt/lib/yt_token_resolver/yt_token_resolver.cpp b/yt/yql/providers/yt/lib/yt_token_resolver/yt_token_resolver.cpp
new file mode 100644
index 00000000000..6bf23bd017d
--- /dev/null
+++ b/yt/yql/providers/yt/lib/yt_token_resolver/yt_token_resolver.cpp
@@ -0,0 +1 @@
+#include "yt_token_resolver.h"
diff --git a/yt/yql/providers/yt/lib/yt_token_resolver/yt_token_resolver.h b/yt/yql/providers/yt/lib/yt_token_resolver/yt_token_resolver.h
new file mode 100644
index 00000000000..53961403b55
--- /dev/null
+++ b/yt/yql/providers/yt/lib/yt_token_resolver/yt_token_resolver.h
@@ -0,0 +1,17 @@
+#pragma once
+
+#include <util/generic/maybe.h>
+#include <util/generic/ptr.h>
+#include <util/generic/string.h>
+
+
+namespace NYql {
+
+class IYtTokenResolver: public TThrRefBase {
+public:
+    using TPtr = TIntrusivePtr<IYtTokenResolver>;
+
+    virtual TMaybe<TString> ResolveClusterToken(const TString& cluster) = 0;
+};
+
+} // namespace NYql
diff --git a/yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_write.cpp b/yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_write.cpp
index fa0f5fe145e..2e3b41414b1 100644
--- a/yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_write.cpp
+++ b/yt/yql/providers/yt/provider/phy_opt/yql_yt_phy_opt_write.cpp
@@ -275,7 +275,6 @@ TMaybeNode<TExprBase> TYtPhysicalOptProposalTransformer::YtDqProcessWrite(TExprB
         [] (const TExprNode::TPtr& node) { return node->IsCallable({TCoToFlow::CallableName(), TCoIterator::CallableName()}) && node->Head().IsCallable(TYtTableContent::CallableName()); });
         !contents.empty()) {
         TNodeOnNodeOwnedMap replaces(contents.size());
-        const bool addToken = !State_->Configuration->Auth.Get().GetOrElse(TString()).empty();
 
         for (const auto& cont : contents) {
             const TYtTableContent content(cont->HeadPtr());
@@ -285,9 +284,10 @@ TMaybeNode<TExprBase> TYtPhysicalOptProposalTransformer::YtDqProcessWrite(TExprB
             if (output) {
                 input = ConvertContentInputToRead(output.Cast(), {}, ctx);
             }
+
+            const auto cluster = input.Cast<TYtReadTable>().DataSource().Cluster().StringValue();
             TMaybeNode<TCoSecureParam> secParams;
-            if (addToken) {
-                const auto cluster = input.Cast<TYtReadTable>().DataSource().Cluster();
+            if (State_->ResolveClusterToken(cluster)) {
                 secParams = Build<TCoSecureParam>(ctx, node.Pos()).Name().Build(TString("cluster:default_").append(cluster)).Done();
             }
 
diff --git a/yt/yql/providers/yt/provider/ya.make b/yt/yql/providers/yt/provider/ya.make
index 78fa92d59de..aa528f84454 100644
--- a/yt/yql/providers/yt/provider/ya.make
+++ b/yt/yql/providers/yt/provider/ya.make
@@ -134,6 +134,7 @@ PEERDIR(
     yt/yql/providers/yt/lib/skiff
     yt/yql/providers/yt/lib/temp_files
     yt/yql/providers/yt/lib/yson_helpers
+    yt/yql/providers/yt/lib/yt_token_resolver
     yt/yql/providers/yt/opt
     yt/yql/providers/yt/gateway/qplayer
     yt/yql/providers/yt/proto
diff --git a/yt/yql/providers/yt/provider/yql_yt_datasource.cpp b/yt/yql/providers/yt/provider/yql_yt_datasource.cpp
index c48d679b038..7bff3315f04 100644
--- a/yt/yql/providers/yt/provider/yql_yt_datasource.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_datasource.cpp
@@ -171,6 +171,18 @@ public:
         return &State_->Configuration->Tokens;
     }
 
+    TMaybe<TString> ResolveClusterToken(const TString& cluster) override {
+        if (!State_->Configuration->IsValidCluster(cluster)) {
+            return {};
+        }
+
+        return State_->ResolveClusterToken(cluster);
+    }
+
+    const THashSet<TString>& GetValidClusters() override {
+        return State_->Configuration->GetValidClusters();
+    }
+
     bool ValidateParameters(TExprNode& node, TExprContext& ctx, TMaybe<TString>& cluster) override {
         if (node.IsCallable(TCoDataSource::CallableName())) {
             if (!EnsureArgsCount(node, 2, ctx)) {
diff --git a/yt/yql/providers/yt/provider/yql_yt_dq_integration.cpp b/yt/yql/providers/yt/provider/yql_yt_dq_integration.cpp
index ee20cdc47f9..b954c84d54f 100644
--- a/yt/yql/providers/yt/provider/yql_yt_dq_integration.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_dq_integration.cpp
@@ -693,8 +693,8 @@ public:
 
         if (auto maybeYtReadTable = TMaybeNode<TYtReadTable>(read)) {
             TMaybeNode<TCoSecureParam> secParams;
-            const auto cluster = maybeYtReadTable.Cast().DataSource().Cluster();
-            if (ytState->Configuration->Auth.Get().GetOrElse(TString()) || ytState->Configuration->Tokens.Value(cluster, "")) {
+            const auto cluster = maybeYtReadTable.Cast().DataSource().Cluster().StringValue();
+            if (ytState->ResolveClusterToken(cluster)) {
                 secParams = Build<TCoSecureParam>(ctx, read->Pos()).Name().Build(TString("cluster:default_").append(cluster)).Done();
             }
             return Build<TDqReadWrap>(ctx, read->Pos())
@@ -940,10 +940,10 @@ public:
             param("external_tx", GetGuidAsString(externalTx));
         }
         TString tokenName;
-        if (auto auth = ytState->Configuration->Auth.Get().GetOrElse(TString())) {
+        if (auto token = ytState->ResolveClusterToken(cluster)) {
             tokenName = TString("cluster:default_").append(cluster);
             if (!secureParams.contains(tokenName)) {
-                secureParams[tokenName] = auth;
+                secureParams[tokenName] = *token;
             }
         }
         param("token", tokenName);
diff --git a/yt/yql/providers/yt/provider/yql_yt_forwarding_gateway.cpp b/yt/yql/providers/yt/provider/yql_yt_forwarding_gateway.cpp
index 4ae7135fb89..27dfab8e306 100644
--- a/yt/yql/providers/yt/provider/yql_yt_forwarding_gateway.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_forwarding_gateway.cpp
@@ -157,4 +157,8 @@ NThreading::TFuture<IYtGateway::TDownloadTableResult> TYtForwardingGatewayBase::
     return Slave_->DownloadTable(std::move(options));
 }
 
+IYtTokenResolver::TPtr TYtForwardingGatewayBase::GetYtTokenResolver() const {
+    return Slave_->GetYtTokenResolver();
+}
+
 } // namspace NYql
diff --git a/yt/yql/providers/yt/provider/yql_yt_forwarding_gateway.h b/yt/yql/providers/yt/provider/yql_yt_forwarding_gateway.h
index 3848702bf73..9633adc0bde 100644
--- a/yt/yql/providers/yt/provider/yql_yt_forwarding_gateway.h
+++ b/yt/yql/providers/yt/provider/yql_yt_forwarding_gateway.h
@@ -81,6 +81,8 @@ public:
 
     NThreading::TFuture<TDownloadTableResult> DownloadTable(TDownloadTableOptions&& options) override;
 
+    IYtTokenResolver::TPtr GetYtTokenResolver() const override;
+
 protected:
     IYtGateway::TPtr Slave_;
 };
diff --git a/yt/yql/providers/yt/provider/yql_yt_gateway.h b/yt/yql/providers/yt/provider/yql_yt_gateway.h
index 6e07a96a653..147d74c5d5b 100644
--- a/yt/yql/providers/yt/provider/yql_yt_gateway.h
+++ b/yt/yql/providers/yt/provider/yql_yt_gateway.h
@@ -7,6 +7,7 @@
 #include <yt/yql/providers/yt/lib/full_capture/yql_yt_full_capture.h>
 #include <yt/yql/providers/yt/lib/row_spec/yql_row_spec.h>
 #include <yt/yql/providers/yt/lib/temp_files/temp_files.h>
+#include <yt/yql/providers/yt/lib/yt_token_resolver/yt_token_resolver.h>
 
 #include <yql/providers/stat/uploader/yql_stat_uploader.h>
 
@@ -815,6 +816,8 @@ public:
     virtual NThreading::TFuture<TDumpResult> Dump(TDumpOptions&& options) = 0;
 
     virtual NThreading::TFuture<TDownloadTableResult> DownloadTable(TDownloadTableOptions&& options) = 0;
+
+    virtual IYtTokenResolver::TPtr GetYtTokenResolver() const = 0;
 };
 
 }
diff --git a/yt/yql/providers/yt/provider/yql_yt_peephole.cpp b/yt/yql/providers/yt/provider/yql_yt_peephole.cpp
index 7094d4b75c9..420b657545c 100644
--- a/yt/yql/providers/yt/provider/yql_yt_peephole.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_peephole.cpp
@@ -104,7 +104,7 @@ private:
         YQL_ENSURE(!HasSetting(wideWrite.Settings().Ref(), "table"));
 
         TMaybeNode<TCoSecureParam> secParams;
-        if (State_->Configuration->Auth.Get().GetOrElse(TString())) {
+        if (State_->ResolveClusterToken(cluster)) {
             secParams = Build<TCoSecureParam>(ctx, node.Pos()).Name().Build(TString("cluster:default_").append(cluster)).Done();
         }
 
diff --git a/yt/yql/providers/yt/provider/yql_yt_provider.cpp b/yt/yql/providers/yt/provider/yql_yt_provider.cpp
index 86a8f1b85a5..bbc8ed41b78 100644
--- a/yt/yql/providers/yt/provider/yql_yt_provider.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_provider.cpp
@@ -481,16 +481,10 @@ TDataProviderInitializer GetYtNativeDataProviderInitializer(IYtGateway::TPtr gat
                 return {};
             }
 
-            // todo: get token by cluster name from Auth when it will be implemented
-            if (auto token = ytState->Configuration->Auth.Get()) {
+            if (auto token = ytState->ResolveClusterToken(cluster)) {
                 return *token;
             }
 
-            if (cluster) {
-                if (auto p = ytState->Configuration->Tokens.FindPtr(cluster)) {
-                    return *p;
-                }
-            }
             return {};
         };
 
@@ -581,4 +575,25 @@ bool TYtState::HybridTakesTooLong() const {
             > Configuration->HybridDqTimeSpentLimit.Get().GetOrElse(TDuration::Minutes(20));
 }
 
+TMaybe<TString> TYtState::ResolveClusterToken(const TString& cluster) {
+    // todo: get token by cluster name from Auth when it will be implemented
+    if (auto token = Configuration->Auth.Get()) {
+        return *token;
+    }
+
+    if (cluster) {
+        if (auto* token = Configuration->Tokens.FindPtr(cluster)) {
+            if (*token) {
+                return *token;
+            }
+        }
+
+        if (auto ytTokenResolver = Gateway->GetYtTokenResolver()) {
+            return ytTokenResolver->ResolveClusterToken(cluster);
+        }
+    }
+
+    return {};
+}
+
 }
diff --git a/yt/yql/providers/yt/provider/yql_yt_provider.h b/yt/yql/providers/yt/provider/yql_yt_provider.h
index ad97fddf991..5a1d3667528 100644
--- a/yt/yql/providers/yt/provider/yql_yt_provider.h
+++ b/yt/yql/providers/yt/provider/yql_yt_provider.h
@@ -98,6 +98,7 @@ struct TYtState {
     bool IsHybridEnabled() const;
     bool IsHybridEnabledForCluster(const std::string_view& cluster) const;
     bool HybridTakesTooLong() const;
+    TMaybe<TString> ResolveClusterToken(const TString& cluster);
 
     TYtState(TTypeAnnotationContext* types, const TQContext& qContext = {}) {
         Types = types;