From 35d789c56ba11136cbecd312607d64b1c5eba7a3 Mon Sep 17 00:00:00 2001 From: antonio2368 Date: Wed, 10 Mar 2021 12:18:09 +0100 Subject: [PATCH] Allow multiple folders for query modules (#105) * Allow multiple query modules directories as arg --- CHANGELOG.md | 4 ++- src/memgraph.cpp | 28 +++++++++++---- src/query/procedure/module.cpp | 64 ++++++++++++++++++++++------------ src/query/procedure/module.hpp | 18 ++++++---- 4 files changed, 78 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 383731586..97bc4f122 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,10 +2,12 @@ ## Future -TODO: Don't forget to add items on the fly. ### Major Feature and Improvements * Added replication to community version. +* Add support for multiple query modules directories at the same time. + You can now define multiple, comma-separated paths to directories from + which the modules will be loaded using the `--query-modules-directory` flag. ### Bug Fixes diff --git a/src/memgraph.cpp b/src/memgraph.cpp index 600bd25b7..9a902db4f 100644 --- a/src/memgraph.cpp +++ b/src/memgraph.cpp @@ -137,12 +137,28 @@ DEFINE_uint64(query_execution_timeout_sec, 180, "Maximum allowed query execution time. Queries exceeding this " "limit will be aborted. Value of 0 means no limit."); -DEFINE_VALIDATED_string(query_modules_directory, "", "Directory where modules with custom query procedures are stored.", +namespace { +std::vector query_modules_directories; +} // namespace +DEFINE_VALIDATED_string(query_modules_directory, "", + "Directory where modules with custom query procedures are stored. " + "NOTE: Multiple comma-separated directories can be defined.", { + query_modules_directories.clear(); if (value.empty()) return true; - if (utils::DirExists(value)) return true; - std::cout << "Expected --" << flagname << " to point to a directory." << std::endl; - return false; + const auto directories = utils::Split(value, ","); + for (const auto &dir : directories) { + if (!utils::DirExists(dir)) { + std::cout << "Expected --" << flagname << " to point to directories." << std::endl; + std::cout << dir << " is not a directory." << std::endl; + return false; + } + } + query_modules_directories.reserve(directories.size()); + std::transform(directories.begin(), directories.end(), + std::back_inserter(query_modules_directories), + [](const auto &dir) { return dir; }); + return true; }); // Logging flags @@ -944,8 +960,8 @@ int main(int argc, char **argv) { SessionData session_data{&db, &interpreter_context}; #endif - query::procedure::gModuleRegistry.SetModulesDirectory(FLAGS_query_modules_directory); - query::procedure::gModuleRegistry.UnloadAndLoadModulesFromDirectory(); + query::procedure::gModuleRegistry.SetModulesDirectory(query_modules_directories); + query::procedure::gModuleRegistry.UnloadAndLoadModulesFromDirectories(); #ifdef MG_ENTERPRISE AuthQueryHandler auth_handler(&auth, std::regex(FLAGS_auth_user_or_role_name_regex)); diff --git a/src/query/procedure/module.cpp b/src/query/procedure/module.cpp index b4eebc7c1..99bc54dc9 100644 --- a/src/query/procedure/module.cpp +++ b/src/query/procedure/module.cpp @@ -78,7 +78,7 @@ void RegisterMgLoad(ModuleRegistry *module_registry, utils::RWLock *lock, Builti }; auto load_all_cb = [module_registry, with_unlock_shared](const mgp_list *, const mgp_graph *, mgp_result *, mgp_memory *) { - with_unlock_shared([&]() { module_registry->UnloadAndLoadModulesFromDirectory(); }); + with_unlock_shared([&]() { module_registry->UnloadAndLoadModulesFromDirectories(); }); }; mgp_proc load_all("load_all", load_all_cb, utils::NewDeleteResource()); module->AddProcedure("load_all", std::move(load_all)); @@ -400,24 +400,16 @@ ModuleRegistry::ModuleRegistry() { modules_.emplace("mg", std::move(module)); } -void ModuleRegistry::SetModulesDirectory(const std::filesystem::path &modules_dir) { modules_dir_ = modules_dir; } +void ModuleRegistry::SetModulesDirectory(std::vector modules_dirs) { + modules_dirs_ = std::move(modules_dirs); +} -bool ModuleRegistry::LoadOrReloadModuleFromName(const std::string_view &name) { - if (modules_dir_.empty()) return false; - if (name.empty()) return false; - std::unique_lock guard(lock_); - auto found_it = modules_.find(name); - if (found_it != modules_.end()) { - if (!found_it->second->Close()) { - spdlog::warn("Failed to close module {}", found_it->first); - } - modules_.erase(found_it); - } - if (!utils::DirExists(modules_dir_)) { - spdlog::error("Module directory {} doesn't exist", modules_dir_); +bool ModuleRegistry::LoadModuleIfFound(const std::filesystem::path &modules_dir, const std::string_view name) { + if (!utils::DirExists(modules_dir)) { + spdlog::error("Module directory {} doesn't exist", modules_dir); return false; } - for (const auto &entry : std::filesystem::directory_iterator(modules_dir_)) { + for (const auto &entry : std::filesystem::directory_iterator(modules_dir)) { const auto &path = entry.path(); if (entry.is_regular_file() && path.stem() == name) { auto module = LoadModuleFromFile(path); @@ -428,15 +420,33 @@ bool ModuleRegistry::LoadOrReloadModuleFromName(const std::string_view &name) { return false; } -void ModuleRegistry::UnloadAndLoadModulesFromDirectory() { - if (modules_dir_.empty()) return; - if (!utils::DirExists(modules_dir_)) { - spdlog::error("Module directory {} doesn't exist", modules_dir_); +bool ModuleRegistry::LoadOrReloadModuleFromName(const std::string_view name) { + if (modules_dirs_.empty()) return false; + if (name.empty()) return false; + std::unique_lock guard(lock_); + auto found_it = modules_.find(name); + if (found_it != modules_.end()) { + if (!found_it->second->Close()) { + spdlog::warn("Failed to close module {}", found_it->first); + } + modules_.erase(found_it); + } + + for (const auto &module_dir : modules_dirs_) { + if (LoadModuleIfFound(module_dir, name)) { + return true; + } + } + return false; +} + +void ModuleRegistry::LoadModulesFromDirectory(const std::filesystem::path &modules_dir) { + if (modules_dir.empty()) return; + if (!utils::DirExists(modules_dir)) { + spdlog::error("Module directory {} doesn't exist", modules_dir); return; } - std::unique_lock guard(lock_); - DoUnloadAllModules(); - for (const auto &entry : std::filesystem::directory_iterator(modules_dir_)) { + for (const auto &entry : std::filesystem::directory_iterator(modules_dir)) { const auto &path = entry.path(); if (entry.is_regular_file()) { std::string name = path.stem(); @@ -448,6 +458,14 @@ void ModuleRegistry::UnloadAndLoadModulesFromDirectory() { } } +void ModuleRegistry::UnloadAndLoadModulesFromDirectories() { + std::unique_lock guard(lock_); + DoUnloadAllModules(); + for (const auto &module_dir : modules_dirs_) { + LoadModulesFromDirectory(module_dir); + } +} + ModulePtr ModuleRegistry::GetModuleNamed(const std::string_view &name) const { std::shared_lock guard(lock_); auto found_it = modules_.find(name); diff --git a/src/query/procedure/module.hpp b/src/query/procedure/module.hpp index 51cc46d4c..cdae588c8 100644 --- a/src/query/procedure/module.hpp +++ b/src/query/procedure/module.hpp @@ -57,11 +57,17 @@ class ModuleRegistry final { void DoUnloadAllModules(); + /// Loads the module if it's in the modules_dir directory + /// @return Whether the module was loaded + bool LoadModuleIfFound(const std::filesystem::path &modules_dir, std::string_view name); + + void LoadModulesFromDirectory(const std::filesystem::path &modules_dir); + public: ModuleRegistry(); - /// Set the modules directory that will be used when (re)loading modules. - void SetModulesDirectory(const std::filesystem::path &modules_dir); + /// Set the modules directories that will be used when (re)loading modules. + void SetModulesDirectory(std::vector modules_dir); /// Atomically load or reload a module with a particular name from the given /// directory. @@ -74,13 +80,13 @@ class ModuleRegistry final { /// /// Return true if the module was loaded or reloaded successfully, false /// otherwise. - bool LoadOrReloadModuleFromName(const std::string_view &name); + bool LoadOrReloadModuleFromName(const std::string_view name); /// Atomically unload all modules and then load all possible modules from the - /// given directory. + /// set directories. /// /// Takes a write lock. - void UnloadAndLoadModulesFromDirectory(); + void UnloadAndLoadModulesFromDirectories(); /// Find a module with given name or return nullptr. /// Takes a read lock. @@ -91,7 +97,7 @@ class ModuleRegistry final { void UnloadAllModules(); private: - std::filesystem::path modules_dir_; + std::vector modules_dirs_; }; /// Single, global module registry.