[yql] Default tokens for file services

12 files changed, 44 insertions, 56 deletions

diff --git a/ydb/library/yql/core/facade/yql_facade.cpp b/ydb/library/yql/core/facade/yql_facade.cpp
index 43ec044998..56c7a43dbf 100644
--- a/ydb/library/yql/core/facade/yql_facade.cpp
+++ b/ydb/library/yql/core/facade/yql_facade.cpp
@@ -87,7 +87,18 @@ TProgram::TStatus SyncExecution(
     return status;
 }
 
-std::function<TMaybe<TString>(const TString&)> BuildCompositeTokenResolver(TVector<std::function<TMaybe<TString>(const TString&)>>&& children) {
+std::function<TString(const TString&, const TString&)> BuildDefaultTokenResolver(const TTypeAnnotationContext* typeCtx) {
+    return [typeCtx](const TString& /*url*/, const TString& alias) -> TString {
+        if (alias) {
+            if (auto cred = typeCtx->FindCredential(TString("default_").append(alias))) {
+                return cred->Content;
+            }
+        }
+        return {};
+    };
+}
+
+std::function<TString(const TString&, const TString&)> BuildCompositeTokenResolver(TVector<std::function<TString(const TString&, const TString&)>>&& children) {
     if (children.empty()) {
         return {};
     }
@@ -96,14 +107,14 @@ std::function<TMaybe<TString>(const TString&)> BuildCompositeTokenResolver(TVect
         return std::move(children[0]);
     }
 
-    return[children = std::move(children)](const TString& url)->TMaybe<TString> {
+    return [children = std::move(children)](const TString& url, const TString& alias) -> TString {
         for (auto& c : children) {
-            if (auto r = c(url)) {
+            if (auto r = c(url, alias)) {
                 return r;
             }
         }
 
-        return Nothing();
+        return {};
     };
 }
 
@@ -1263,6 +1274,11 @@ void TProgram::CloseLastSession() {
             dp.CloseSession(sessionId);
         }
     }
+
+    // Token resolver may keep some references to provider internal's. So reset it to release provider's data
+    if (FileStorage_) {
+        FileStorage_->SetTokenResolver({});
+    }
 }
 
 TString TProgram::ResultsAsString() const {
@@ -1307,7 +1323,7 @@ TTypeAnnotationContextPtr TProgram::BuildTypeAnnotationContext(const TString& us
     PlanBuilder_ = CreatePlanBuilder(*typeAnnotationContext);
     THashSet<TString> providerNames;
     TVector<TString> fullResultDataSinks;
-    TVector<std::function<TMaybe<TString>(const TString& url)>> tokenResolvers;
+    TVector<std::function<TString(const TString&, const TString&)>> tokenResolvers;
     for (const auto& dpi : DataProvidersInit_) {
         auto dp = dpi(
             username,
@@ -1385,7 +1401,11 @@ TTypeAnnotationContextPtr TProgram::BuildTypeAnnotationContext(const TString& us
         typeAnnotationContext->AddDataSource(ConfigProviderName, configProvider);
     }
 
-    typeAnnotationContext->UserDataStorage->SetTokenResolver(BuildCompositeTokenResolver(std::move(tokenResolvers)));
+    tokenResolvers.push_back(BuildDefaultTokenResolver(typeAnnotationContext.Get()));
+    if (FileStorage_) {
+        FileStorage_->SetTokenResolver(BuildCompositeTokenResolver(std::move(tokenResolvers)));
+    }
+
     return typeAnnotationContext;
 }
 
diff --git a/ydb/library/yql/core/file_storage/file_storage.cpp b/ydb/library/yql/core/file_storage/file_storage.cpp
index 29518b39da..84c77ffd17 100644
--- a/ydb/library/yql/core/file_storage/file_storage.cpp
+++ b/ydb/library/yql/core/file_storage/file_storage.cpp
@@ -25,6 +25,7 @@
 #include <util/generic/guid.h>
 #include <util/generic/yexception.h>
 
+#include <util/string/builder.h>
 #include <util/stream/str.h>
 #include <util/stream/file.h>
 #include <util/stream/null.h>
@@ -178,6 +179,9 @@ public:
         return Config;
     }
 
+    void SetTokenResolver(std::function<TString(const TString&, const TString&)> /*tokenResolver*/) final {
+    }
+
 private:
     TFileLinkPtr PutUrl(const THttpURL& url, const TString& token, const NFS::IDownloaderPtr& downloader) {
         return WithRetry<TDownloadError>(Config.GetRetryCount(), [&, this]() {
@@ -189,7 +193,7 @@ private:
     }
 
     TFileLinkPtr DoPutUrl(const THttpURL& url, const TString& token, const NFS::IDownloaderPtr& downloader) {
-        const auto urlMetaFile = BuildUrlMetaFileName(url);
+        const auto urlMetaFile = BuildUrlMetaFileName(url, token);
         auto lock = MultiResourceLock.Acquire(urlMetaFile); // let's use meta file as lock name
 
         TUrlMeta urlMeta;
@@ -231,7 +235,7 @@ private:
         }
 
         // todo: remove oldContentLink ?
-        const auto urlContentFile = BuildUrlContentFileName(url, etag, lastModified);
+        const auto urlContentFile = BuildUrlContentFileName(url, token, etag, lastModified);
         TFileLinkPtr result = Storage.Put(urlContentFile, TString(), TString(), puller);
 
         // save meta using rename for atomicity
@@ -247,13 +251,13 @@ private:
         return result;
     }
 
-    static TString BuildUrlMetaFileName(const THttpURL& url) {
-        return MD5::Calc(url.PrintS(THttpURL::FlagNoFrag | THttpURL::FlagHostAscii)) + ".url_meta";
+    static TString BuildUrlMetaFileName(const THttpURL& url, const TString& token) {
+        return MD5::Calc(TStringBuilder() << token << url.PrintS(THttpURL::FlagNoFrag | THttpURL::FlagHostAscii)) + ".url_meta";
     }
 
-    static TString BuildUrlContentFileName(const THttpURL& url, const TString& etag, const TString& lastModified) {
+    static TString BuildUrlContentFileName(const THttpURL& url, const TString& token, const TString& etag, const TString& lastModified) {
         TString needle = etag ? etag : lastModified;
-        return MD5::Calc(needle + url.PrintS(THttpURL::FlagNoFrag | THttpURL::FlagHostAscii)) + ".url";
+        return MD5::Calc(TStringBuilder() << token << needle << url.PrintS(THttpURL::FlagNoFrag | THttpURL::FlagHostAscii)) + ".url";
     }
 
 private:
diff --git a/ydb/library/yql/core/file_storage/file_storage.h b/ydb/library/yql/core/file_storage/file_storage.h
index b95ccabaaf..624badc8e1 100644
--- a/ydb/library/yql/core/file_storage/file_storage.h
+++ b/ydb/library/yql/core/file_storage/file_storage.h
@@ -12,6 +12,7 @@
 #include <util/generic/string.h>
 
 #include <vector>
+#include <functional>
 
 namespace NYql {
 
@@ -31,6 +32,7 @@ struct IFileStorage: public TThrRefBase {
     virtual TFsPath GetRoot() const = 0;
     virtual TFsPath GetTemp() const = 0;
     virtual const TFileStorageConfig& GetConfig() const = 0;
+    virtual void SetTokenResolver(std::function<TString(const TString&, const TString&)> tokenResolver) = 0;
 };
 
 using TFileStoragePtr = TIntrusivePtr<IFileStorage>;
diff --git a/ydb/library/yql/core/file_storage/file_storage_decorator.cpp b/ydb/library/yql/core/file_storage/file_storage_decorator.cpp
index ec9e5bfac8..f2b8e8db29 100644
--- a/ydb/library/yql/core/file_storage/file_storage_decorator.cpp
+++ b/ydb/library/yql/core/file_storage/file_storage_decorator.cpp
@@ -38,5 +38,8 @@ TFsPath TFileStorageDecorator::GetTemp() const {
 const TFileStorageConfig& TFileStorageDecorator::GetConfig() const {
     return Inner_->GetConfig();
 }
+void TFileStorageDecorator::SetTokenResolver(std::function<TString(const TString&, const TString&)> tokenResolver) {
+    Inner_->SetTokenResolver(std::move(tokenResolver));
+}
 
 } // NYql
diff --git a/ydb/library/yql/core/file_storage/file_storage_decorator.h b/ydb/library/yql/core/file_storage/file_storage_decorator.h
index 9d6f2b8f3e..dcb8d323eb 100644
--- a/ydb/library/yql/core/file_storage/file_storage_decorator.h
+++ b/ydb/library/yql/core/file_storage/file_storage_decorator.h
@@ -7,6 +7,7 @@ namespace NYql {
 class TFileStorageDecorator: public IFileStorage {
 public:
     TFileStorageDecorator(TFileStoragePtr fs);
+    ~TFileStorageDecorator() = default;
 
     TFileLinkPtr PutFile(const TString& file, const TString& outFileName) override;
     TFileLinkPtr PutFileStripped(const TString& file, const TString& originalMd5) override;
@@ -18,6 +19,7 @@ public:
     TFsPath GetRoot() const override;
     TFsPath GetTemp() const override;
     const TFileStorageConfig& GetConfig() const override;
+    void SetTokenResolver(std::function<TString(const TString&, const TString&)> tokenResolver) override;
 
 protected:
     TFileStoragePtr Inner_;
diff --git a/ydb/library/yql/core/yql_data_provider.h b/ydb/library/yql/core/yql_data_provider.h
index ab4fefae5f..b6f5b56b23 100644
--- a/ydb/library/yql/core/yql_data_provider.h
+++ b/ydb/library/yql/core/yql_data_provider.h
@@ -228,7 +228,7 @@ struct TDataProviderInfo {
 
     std::function<void(const TString& sessionId)> CleanupSession;
 
-    std::function<TMaybe<TString>(const TString& url)> TokenResolver;
+    std::function<TString(const TString& url, const TString& alias)> TokenResolver;
 };
 
 using THiddenQueryAborter = std::function<void()>; // aborts hidden query, which is running within a separate TProgram
diff --git a/ydb/library/yql/core/yql_type_annotation.cpp b/ydb/library/yql/core/yql_type_annotation.cpp
index 743210dd64..7d0815a701 100644
--- a/ydb/library/yql/core/yql_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_type_annotation.cpp
@@ -38,7 +38,6 @@ bool TTypeAnnotationContext::DoInitialize(TExprContext& ctx) {
     }
 
     Y_ENSURE(UserDataStorage);
-    UserDataStorage->FillUserDataTokens();
 
     // Disable "in progress" constraints
     //DisableConstraintCheck.emplace(TSortedConstraintNode::Name());
@@ -219,7 +218,6 @@ bool TModuleResolver::AddFromUrl(const TStringBuf& file, const TStringBuf& url,
     block.Type = EUserDataType::URL;
     block.Data = url;
     block.Data = SubstParameters(block.Data);
-    UserData->TryFillUserDataToken(block);
     UserData->AddUserDataBlock(file, block);
 
     return AddFromFile(file, ctx, syntaxVersion, packageVersion);
diff --git a/ydb/library/yql/core/yql_user_data_storage.cpp b/ydb/library/yql/core/yql_user_data_storage.cpp
index 1a9d725d27..3b9e01524a 100644
--- a/ydb/library/yql/core/yql_user_data_storage.cpp
+++ b/ydb/library/yql/core/yql_user_data_storage.cpp
@@ -45,10 +45,6 @@ TUserDataStorage::TUserDataStorage(TFileStoragePtr fileStorage, TUserDataTable d
 {
 }
 
-void TUserDataStorage::SetTokenResolver(TTokenResolver tokenResolver) {
-    TokenResolver_ = std::move(tokenResolver);
-}
-
 void TUserDataStorage::AddUserDataBlock(const TStringBuf& name, const TUserDataBlock& block) {
     const auto key = ComposeUserDataKey(name);
     AddUserDataBlock(key, block);
@@ -156,29 +152,6 @@ TMaybe<std::map<TUserDataKey, const TUserDataBlock*>> TUserDataStorage::FindUser
     return res;
 }
 
-void TUserDataStorage::FillUserDataTokens() {
-    for (auto& p : UserData_) {
-        TryFillUserDataToken(p.second);
-    }
-}
-
-void TUserDataStorage::TryFillUserDataToken(TUserDataBlock& block) const {
-    if (block.Type != EUserDataType::URL) {
-        return;
-    }
-
-    // no need to overwrite existing value
-    if (block.UrlToken) {
-        return;
-    }
-
-    if (!TokenResolver_) {
-        return;
-    }
-
-    block.UrlToken = TokenResolver_(block.Data).GetOrElse({});
-}
-
 std::map<TString, const TUserDataBlock*> TUserDataStorage::GetDirectoryContent(const TStringBuf& path, ui32 maxFileCount) const {
     const auto fullPath = MakeFolderName(path);
 
diff --git a/ydb/library/yql/core/yql_user_data_storage.h b/ydb/library/yql/core/yql_user_data_storage.h
index f1575cfffe..6da0eb2548 100644
--- a/ydb/library/yql/core/yql_user_data_storage.h
+++ b/ydb/library/yql/core/yql_user_data_storage.h
@@ -14,11 +14,9 @@ namespace NYql {
 class TUserDataStorage : public TThrRefBase {
 public:
     typedef TIntrusivePtr<TUserDataStorage> TPtr;
-    typedef std::function<TMaybe<TString>(const TString& url)> TTokenResolver;
 
 public:
     TUserDataStorage(TFileStoragePtr fileStorage, TUserDataTable data, IUdfResolver::TPtr udfResolver, TUdfIndex::TPtr udfIndex);
-    void SetTokenResolver(TTokenResolver tokenResolver);
 
     void AddUserDataBlock(const TStringBuf& name, const TUserDataBlock& block);
     void AddUserDataBlock(const TUserDataKey& key, const TUserDataBlock& block);
@@ -40,8 +38,6 @@ public:
     TMaybe<std::map<TUserDataKey, const TUserDataBlock*>> FindUserDataFolder(const TStringBuf& name, ui32 maxFileCount = ~0u) const;
     static TMaybe<std::map<TUserDataKey, const TUserDataBlock*>> FindUserDataFolder(const TUserDataTable& userData, const TStringBuf& name, ui32 maxFileCount = ~0u);
 
-    void FillUserDataTokens();
-    void TryFillUserDataToken(TUserDataBlock& block) const;
     std::map<TString, const TUserDataBlock*> GetDirectoryContent(const TStringBuf& path, ui32 maxFileCount = ~0u) const;
     static TString MakeFullName(const TStringBuf& name);
     static TString MakeFolderName(const TStringBuf& name);
@@ -73,7 +69,6 @@ private:
     TUserDataTable UserData_;
     IUdfResolver::TPtr UdfResolver;
     TUdfIndex::TPtr UdfIndex;
-    TTokenResolver TokenResolver_;
 
     THashSet<TUserDataKey, TUserDataKey::THash, TUserDataKey::TEqualTo> ScannedUdfs;
     std::function<void(const TUserDataBlock& block)> ScanUdfStrategy_;
diff --git a/ydb/library/yql/providers/common/provider/yql_provider.cpp b/ydb/library/yql/providers/common/provider/yql_provider.cpp
index aa99dcc44a..1effbb2b41 100644
--- a/ydb/library/yql/providers/common/provider/yql_provider.cpp
+++ b/ydb/library/yql/providers/common/provider/yql_provider.cpp
@@ -551,7 +551,6 @@ bool FillUsedFilesImpl(
                 const auto it = crutches.find(geobase);
                 if (crutches.cend() != it) {
                     auto pragma = it->second;
-                    types.UserDataStorage->TryFillUserDataToken(pragma);
                     types.UserDataStorage->AddUserDataBlock(geobase, pragma);
                     files.emplace(geobase, pragma).first->second.Usage.Set(EUserDataBlockUsage::Path);
                 }
diff --git a/ydb/library/yql/providers/config/yql_config_provider.cpp b/ydb/library/yql/providers/config/yql_config_provider.cpp
index 8acc03b651..60b0b7106c 100644
--- a/ydb/library/yql/providers/config/yql_config_provider.cpp
+++ b/ydb/library/yql/providers/config/yql_config_provider.cpp
@@ -845,8 +845,6 @@ namespace {
             block.Data = url;
             if (token) {
                 block.UrlToken = token;
-            } else {
-                Types.UserDataStorage->TryFillUserDataToken(block);
             }
             Types.UserDataStorage->AddUserDataBlock(key, block);
             return true;
diff --git a/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp b/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp
index 3c297af994..100347422b 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp
@@ -108,12 +108,6 @@ TDataProviderInitializer GetDqDataProviderInitializer(
             }
         };
 
-        info.TokenResolver = [](const TString& url) -> TMaybe<TString> {
-            Y_UNUSED(url);
-
-            return Nothing();
-        };
-
         return info;
     };
 }