tursom/memgraph
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow multiple folders for query modules ()

Browse Source 
* Allow multiple query modules directories as arg
This commit is contained in:
antonio2368 2021-03-10 12:18:09 +01:00 committed by GitHub
parent 16715d5005
commit 35d789c56b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 78 additions and 36 deletions

4
CHANGELOG.md
View File

@ -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

28
src/memgraph.cpp
View File

@ -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<std::filesystem::path> 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));

64
src/query/procedure/module.cpp
View File

@ -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<std::filesystem::path> 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<utils::RWLock> 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<utils::RWLock> 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<utils::RWLock> 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<utils::RWLock> 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<utils::RWLock> guard(lock_);
auto found_it = modules_.find(name);

18
src/query/procedure/module.hpp
View File

@ -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<std::filesystem::path> 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<std::filesystem::path> modules_dirs_;
};
/// Single, global module registry.