MongoDBStorageMixin.cpp

Go to the documentation of this file.

#include "MongoDBStorageMixin.h"
 
#include <SimoxUtility/algorithm/string.h>
 
#include <ArmarXCore/core/exceptions/LocalException.h>
#include <ArmarXCore/core/exceptions/local/ExpressionException.h>
#include <ArmarXCore/core/logging/Logging.h>
 
#include "util/filesystem.h"
 
namespace armarx::armem::server::ltm::detail::mixin
{
    std::unique_ptr<mongocxx::instance> MongoDBStorageMixin::MongoCXXInstance = nullptr;
    std::atomic_bool MongoDBStorageMixin::MongoCXXInitialized = false;
    std::recursive_mutex MongoDBStorageMixin::MongoCXXConnectionPoolMutex;
    std::map<std::string, std::unique_ptr<mongocxx::pool>>
        MongoDBStorageMixin::MongoCXXConnectionPool = {};
 
    void
    MongoDBSettings::initializeFromArmarXConfig()
    {
        std::filesystem::path armarx_config_home = std::string(getenv("ARMARX_USER_CONFIG_DIR"));
        if (!util::fs::directoryExists(armarx_config_home))
        {
            ARMARX_WARNING << ("Could not find an ArmarX Config folder. Tried: '" +
                               armarx_config_home.string() +
                               "' (from ARMARX_USER_CONFIG_DIR). If not set via scenario params "
                               "this means that LTM is disabled.");
            return;
        }
 
        std::ifstream cFile(armarx_config_home / "default.cfg");
        if (cFile.is_open())
        {
            std::string line;
            while (getline(cFile, line))
            {
                line.erase(std::remove_if(line.begin(), line.end(), isspace), line.end());
                if (line.empty() || line[0] == '#')
                {
                    continue;
                }
                auto delimiterPos = line.find("=");
                const auto name = simox::alg::to_lower(line.substr(0, delimiterPos));
                const auto value = line.substr(delimiterPos + 1);
                if (host.empty() && name == "armarx.mongohost")
                {
                    host = value;
                }
                if (port == 0 && name == "armarx.mongoport")
                {
                    port = (unsigned int)std::stoi(value);
                }
                if (user.empty() && name == "armarx.mongouser")
                {
                    user = value;
                }
                if (password.empty() && name == "armarx.mongopassword")
                {
                    password = value;
                }
            }
        }
    }
 
    bool
    MongoDBSettings::isSet() const
    {
        // we always need a host and a port
        return !host.empty() and port != 0;
    }
 
    std::string
    MongoDBSettings::baseUri() const
    {
        std::stringstream ss;
        ss << "mongodb://";
 
        if (!user.empty())
        {
            ss << user;
            if (!password.empty())
            {
                ss << ":" << password;
            }
            ss << "@";
        }
        ss << host;
        return ss.str();
    }
 
    std::string
    MongoDBSettings::key() const
    {
        // TODO: What happens if a connection exists and you would like to open another one with a different user (e.g. that sees different things)?
        return "mongodb://" + host + ":" + std::to_string(port);
    }
 
    std::string
    MongoDBSettings::uri() const
    {
        return baseUri() + ":" + std::to_string(port) +
               "/?minPoolSize=" + std::to_string(minPoolSize) +
               "&maxPoolSize=" + std::to_string(maxPoolSize);
    }
 
    std::string
    MongoDBSettings::toString() const
    {
        return uri();
    }
 
    MongoDBStorageMixin::MongoDBStorageMixin(const MongoDBSettings& settings,
                                             const std::string& en,
                                             const armem::MemoryID& id) :
        settings(settings), exportName(en), _id(id)
    {
    }
 
    void
    MongoDBStorageMixin::setHost(const std::string& n)
    {
        settings.host = n;
    }
 
    void
    MongoDBStorageMixin::setPort(const unsigned int n)
    {
        settings.port = n;
    }
 
    void
    MongoDBStorageMixin::setUser(const std::string& n)
    {
        settings.user = n;
    }
 
    void
    MongoDBStorageMixin::setPassword(const std::string& n)
    {
        settings.password = n;
    }
 
    void
    MongoDBStorageMixin::connect() const
    {
        std::lock_guard l(MongoCXXConnectionPoolMutex);
        if (!MongoCXXInitialized.exchange(true))
        {
            ARMARX_IMPORTANT << "INITIALIZE MONGODB";
            MongoCXXInstance =
                std::make_unique<mongocxx::instance>(); // This should be done only once.
        }
 
        const auto key = settings.key();
        const auto uri = settings.uri();
        auto it = MongoCXXConnectionPool.find(key);
        if (it == MongoCXXConnectionPool.end())
        {
            ARMARX_INFO << "Establishing new connection to: " << uri << " from id: " << _id.str();
            mongocxx::uri u(uri);
            auto pool = std::make_unique<mongocxx::pool>(u);
            MongoCXXConnectionPool.emplace(settings.key(), std::move(pool));
        }
    }
 
    mongocxx::pool*
    MongoDBStorageMixin::getPool(bool tryToConnect) const
    {
        std::lock_guard l(MongoCXXConnectionPoolMutex);
        const auto key = settings.key();
        auto it = MongoCXXConnectionPool.find(key);
        if (it == MongoCXXConnectionPool.end())
        {
            if (tryToConnect)
            {
                // try to establish a connection and try again
                connect();
                return getPool(false);
            }
            return nullptr;
        }
        else
        {
            return it->second.get();
        }
    }
 
    mongocxx::pool*
    MongoDBStorageMixin::ensurePool() const
    {
        std::lock_guard l(MongoCXXConnectionPoolMutex);
        auto pool = getPool();
        if (pool)
        {
            return pool;
        }
        else
        {
            throw armarx::LocalException("Pool could not be ensured...");
        }
    }
 
    void
    MongoDBStorageMixin::start()
    {
        // ensure a connection is made if not done already
        connect();
    }
 
    void
    MongoDBStorageMixin::stop()
    {
    }
 
    bool
    MongoDBStorageMixin::connected() const
    {
        std::lock_guard l(MongoCXXConnectionPoolMutex);
 
        try
        {
            auto client = ensurePool()->acquire();
            auto admin = client->database("admin");
            auto result = admin.run_command(bsoncxx::builder::basic::make_document(
                bsoncxx::builder::basic::kvp("isMaster", 1)));
            return true;
        }
        catch (const std::exception& xcp)
        {
            return false;
        }
        return false; // should never happen
    }
 
    std::string
    MongoDBStorageMixin::getDocumentName() const
    {
        return util::mongodb::toDocumentID(_id);
    }
 
    std::string
    MongoDBStorageMixin::getCollectionName() const
    {
        return util::mongodb::toCollectionName(_id);
    }
 
    std::string
    MongoDBStorageMixin::getPreviousCollectionName() const
    {
        ARMARX_CHECK(!_id.memoryName.empty() and !_id.coreSegmentName.empty());
        return util::mongodb::toCollectionName(_id.removeLeafItem());
    }
 
    std::string
    MongoDBStorageMixin::getDatabaseName() const
    {
        return exportName;
    }
 
    std::optional<mongocxx::database>
    MongoDBStorageMixin::databaseExists() const
    {
        auto client = this->ensurePool()->acquire();
        return util::mongodb::databaseExists(*client, getDatabaseName());
    }
 
    std::optional<mongocxx::collection>
    MongoDBStorageMixin::collectionExists() const
    {
        auto db = databaseExists();
        if (!db)
        {
            return std::nullopt;
        }
 
        return util::mongodb::collectionExists(*db, getCollectionName());
    }
 
    std::optional<mongocxx::collection>
    MongoDBStorageMixin::previousCollectionExists() const
    {
        auto db = databaseExists();
        if (!db)
        {
            return std::nullopt;
        }
 
        return util::mongodb::collectionExists(*db, getPreviousCollectionName());
    }
 
    std::optional<nlohmann::json>
    MongoDBStorageMixin::documentExists() const
    {
        return documentExists(getDocumentName());
    }
 
    std::optional<nlohmann::json>
    MongoDBStorageMixin::documentExists(const std::string& id) const
    {
        auto coll = collectionExists();
        if (!coll)
        {
            return std::nullopt;
        }
 
        nlohmann::json query;
        query[ID] = getDocumentName();
        return util::mongodb::documentExists(*coll, query);
    }
 
    mongocxx::database
    MongoDBStorageMixin::ensureDatabaseExists(bool createIfNotExistent)
    {
        auto client = this->ensurePool()->acquire();
        return util::mongodb::ensureDatabaseExists(*client, getDatabaseName(), createIfNotExistent);
    }
 
    mongocxx::collection
    MongoDBStorageMixin::ensureCollectionExists(bool createIfNotExistent)
    {
        auto db = ensureDatabaseExists(createIfNotExistent);
        return util::mongodb::ensureCollectionExists(db, getCollectionName(), createIfNotExistent);
    }
 
    mongocxx::collection
    MongoDBStorageMixin::ensurePreviousCollectionExists(bool createIfNotExistent)
    {
        auto db = ensureDatabaseExists(createIfNotExistent);
        return util::mongodb::ensureCollectionExists(
            db, getPreviousCollectionName(), createIfNotExistent);
    }
 
    nlohmann::json
    MongoDBStorageMixin::ensureDocumentExists(bool createIfNotExistent)
    {
        return ensureDocumentExists(getDocumentName(), createIfNotExistent);
    }
 
    nlohmann::json
    MongoDBStorageMixin::ensureDocumentExists(const std::string& id, bool createIfNotExistent)
    {
        auto coll = ensureCollectionExists(createIfNotExistent);
 
        nlohmann::json query;
        query[ID] = getDocumentName();
        return util::mongodb::ensureDocumentExists(coll, query, createIfNotExistent);
    }
 
    void
    MongoDBStorageMixin::writeDataToDocument(const std::string& id, const nlohmann::json& data)
    {
        auto coll = ensureCollectionExists(true);
 
        nlohmann::json query;
        query[ID] = getDocumentName();
 
        nlohmann::json update;
        update[DATA] = data;
        util::mongodb::writeDataToDocument(coll, query, update);
    }
 
    void
    MongoDBStorageMixin::writeForeignKeyToPreviousDocument()
    {
        auto coll = ensurePreviousCollectionExists(true);
 
        nlohmann::json query;
        query[ID] = _id.str();
 
        nlohmann::json update;
        update[FOREIGN_KEY] = getCollectionName();
        util::mongodb::writeDataToDocument(coll, query, update);
    }
 
    void
    MongoDBStorageMixin::writeForeignKeyToPreviousDocument(const nlohmann::json& type)
    {
        auto coll = ensurePreviousCollectionExists(true);
 
        nlohmann::json query;
        query[ID] = _id.str();
 
        nlohmann::json update;
        update[FOREIGN_KEY] = getCollectionName();
        update[TYPE] = type;
 
        util::mongodb::writeDataToDocument(coll, query, update);
    }
 
    nlohmann::json
    MongoDBStorageMixin::readDataFromDocument(const std::string& id) const
    {
        auto coll = collectionExists();
        if (!coll)
        {
            // What to do here?
            return {};
        }
 
        if (!documentExists())
        {
            // What to do here?
            return {};
        }
 
        auto query =
            nlohmann::json::parse(std::string("{\"") + ID + "\": " + getDocumentName() + "}");
        return util::mongodb::readDataFromDocument(*coll, query);
    }
 
    void
    MongoDBStorageMixin::writeDataToDocument(const nlohmann::json& data)
    {
        writeDataToDocument(getDocumentName(), data);
    }
 
    nlohmann::json
    MongoDBStorageMixin::readDataFromDocument() const
    {
        return readDataFromDocument(getDocumentName());
    }
 
    std::vector<nlohmann::json>
    MongoDBStorageMixin::getAllDocuments() const
    {
        auto coll = collectionExists();
        if (!coll)
        {
            return {};
        }
        return util::mongodb::getAllDocuments(*coll);
    }
 
    MongoDBSettings
    MongoDBStorageMixin::getSettings() const
    {
        return settings;
    }
 
    void
    MongoDBStorageMixin::setMixinMemoryID(const MemoryID& n)
    {
        ARMARX_CHECK_NOT_EMPTY(_id.memoryName) << " The full id was: " << _id.str();
 
        _id = n;
    }
 
    void
    MongoDBStorageMixin::setMixinExportName(const std::string& n)
    {
        exportName = n;
    }
 
    void
    MongoDBStorageMixin::configureMixin(const nlohmann::json& json)
    {
    }
} // namespace armarx::armem::server::ltm::detail::mixin