#pragma once #include "yql_udf_resolver.h" #include #include #include #include #include #include namespace NYql { struct TFunctionInfo { TString Name; int ArgCount = 0; int OptionalArgCount = 0; bool IsTypeAwareness = false; TString CallableType; TString RunConfigType; bool IsStrict = false; bool SupportsBlocks = false; }; // todo: specify whether path is frozen struct TDownloadLink { bool IsUrl = false; TString Path; TString Md5; TDownloadLink() { } TDownloadLink(bool isUrl, const TString& path, const TString& md5) : IsUrl(isUrl) , Path(path) , Md5(md5) { } TDownloadLink(const TDownloadLink&) = default; TDownloadLink& operator=(const TDownloadLink&) = default; static TDownloadLink Url(const TString& path, const TString& md5 = "") { return { true, path, md5 }; } static TDownloadLink File(const TString& path, const TString& md5 = "") { return { false, path, md5 }; } bool operator==(const TDownloadLink& other) const { return std::tie(IsUrl, Path, Md5) == std::tie(other.IsUrl, other.Path, Md5); } bool operator!=(const TDownloadLink& other) const { return !(*this == other); } size_t Hash() const { return CombineHashes( CombineHashes((size_t)IsUrl, ComputeHash(Path)), ComputeHash(Md5) ); } }; struct TResourceInfo : public TThrRefBase { typedef TIntrusiveConstPtr TPtr; bool IsTrusted = false; TDownloadLink Link; TSet Modules; TMap Functions; TMap> ICaseFuncNames; void SetFunctions(const TVector& functions) { for (auto& f : functions) { Functions.emplace(f.Name, f); ICaseFuncNames[to_lower(f.Name)].insert(f.Name); } } }; inline bool operator<(const TResourceInfo::TPtr& p1, const TResourceInfo::TPtr& p2) { return p1.Get() < p2.Get(); } class TUdfIndex : public TThrRefBase { public: typedef TIntrusivePtr TPtr; public: // todo: trusted resources should not be replaceble regardless of specified mode enum class EOverrideMode { PreserveExisting, ReplaceWithNew, RaiseError }; enum class EStatus { Found, NotFound, Ambigious }; public: TUdfIndex(); void SetCaseSentiveSearch(bool caseSensitive); bool CanonizeModule(TString& moduleName) const; EStatus ContainsModule(const TString& moduleName) const; EStatus FindFunction(const TString& moduleName, const TString& functionName, TFunctionInfo& function) const; TResourceInfo::TPtr FindResourceByModule(const TString& moduleName) const; bool ContainsModuleStrict(const TString& moduleName) const; /* New resource can contain already registered module. In this case 'mode' will be used to resolve conflicts. For instance, if mode == ReplaceWithNew all functions from old resource will be removed and new functions will be registered. It is important to do it atomically because two .so cannot have intersecting module lists */ void RegisterResource(const TResourceInfo::TPtr& resource, EOverrideMode mode); void RegisterResources(const TVector& resources, EOverrideMode mode); TIntrusivePtr Clone() const; private: explicit TUdfIndex(const TMap& resources, bool caseSensitive); bool ContainsAnyModule(const TSet& modules) const; TSet FindResourcesByModules(const TSet& modules) const; void UnregisterResource(TResourceInfo::TPtr resource); private: // module => Resource TMap Resources_; bool CaseSensitive_ = true; TMap> ICaseModules_; }; void LoadRichMetadataToUdfIndex(const IUdfResolver& resolver, const TVector& paths, bool isTrusted, TUdfIndex::EOverrideMode mode, TUdfIndex& registry); void LoadRichMetadataToUdfIndex(const IUdfResolver& resolver, const TMap& pathsWithMd5, bool isTrusted, TUdfIndex::EOverrideMode mode, TUdfIndex& registry); void LoadRichMetadataToUdfIndex(const IUdfResolver& resolver, const TVector& blocks, TUdfIndex::EOverrideMode mode, TUdfIndex& registry); void LoadRichMetadataToUdfIndex(const IUdfResolver& resolver, const TUserDataBlock& block, TUdfIndex::EOverrideMode mode, TUdfIndex& registry); }