[V3 Config] Add MongoDB driver (#972)

* Basic Mongo Driver

* Update docstrings

* WIP motor support

* Use motor

* Update docs, add selective importer

* Make use of selective importer

* Fix docs

* Fix config storage location for JSON

* Add delimiters in the drivers doc section

* Make async things async

* Add basic config information for mongo driver

* get info from basic setup into config

* IT WORKS

* Add dependency for RJM's PR.

docs/framework_config.rst

@@ -191,5 +191,19 @@ Driver Reference
 ****************
 
 .. automodule:: redbot.core.drivers
+    :members:
 
+Base Driver
+^^^^^^^^^^^
 .. autoclass:: redbot.core.drivers.red_base.BaseDriver
+    :members:
+
+JSON Driver
+^^^^^^^^^^^
+.. autoclass:: redbot.core.drivers.red_json.JSON
+    :members:
+
+Mongo Driver
+^^^^^^^^^^^^
+.. autoclass:: redbot.core.drivers.red_mongo.Mongo
+    :members:

redbot/core/bank.py

@@ -435,7 +435,7 @@ async def set_global(global_: bool, user: Union[discord.User, discord.Member]) -
     if (await is_global()) is global_:
         return global_
 
-    if is_global():
+    if await is_global():
         await _conf.clear_all_users()
     else:
         await _conf.clear_all_members()

redbot/core/config.py

@@ -5,7 +5,7 @@ from typing import Callable, Union, Tuple
 import discord
 
 from .data_manager import cog_data_path, core_data_path
-from .drivers.red_json import JSON as JSONDriver
+from .drivers import get_driver
 
 log = logging.getLogger("red.config")
 
@@ -440,7 +440,18 @@ class Config:
         cog_name = cog_path_override.stem
         uuid = str(hash(identifier))
 
-        spawner = JSONDriver(cog_name, data_path_override=cog_path_override)
+        # We have to import this here otherwise we have a circular dependency
+        from .data_manager import basic_config
+
+        log.debug("Basic config: \n\n{}".format(basic_config))
+
+        driver_name = basic_config.get('STORAGE_TYPE', 'JSON')
+        driver_details = basic_config.get('STORAGE_DETAILS', {})
+
+        log.debug("Using driver: '{}'".format(driver_name))
+
+        spawner = get_driver(driver_name, cog_name, data_path_override=cog_path_override,
+                             **driver_details)
         return cls(cog_name=cog_name, unique_identifier=uuid,
                    force_registration=force_registration,
                    driver_spawn=spawner)
@@ -458,7 +469,16 @@ class Config:
             See `force_registration`.
 
         """
-        driver_spawn = JSONDriver("Core", data_path_override=core_data_path())
+        core_path = core_data_path()
+
+        # We have to import this here otherwise we have a circular dependency
+        from .data_manager import basic_config
+
+        driver_name = basic_config.get('STORAGE_TYPE', 'JSON')
+        driver_details = basic_config.get('STORAGE_DETAILS', {})
+
+        driver_spawn = get_driver(driver_name, "Core", data_path_override=core_path,
+                                  **driver_details)
         return cls(cog_name="Core", driver_spawn=driver_spawn,
                    unique_identifier='0',
                    force_registration=force_registration)
@@ -848,7 +868,7 @@ class Config:
 
         """
         if not scopes:
-            group = Group(identifiers=(self.unique_identifier),
+            group = Group(identifiers=(self.unique_identifier, ),
                           defaults={},
                           spawner=self.spawner)
         else:

redbot/core/data_manager.py

@@ -11,6 +11,8 @@ __all__ = ['load_basic_configuration', 'cog_data_path', 'core_data_path',
 jsonio = None
 basic_config = None
 
+instance_name = None
+
 basic_config_default = {
     "DATA_PATH": None,
     "COG_PATH_APPEND": "cogs",
@@ -21,12 +23,15 @@ config_dir = Path(appdirs.AppDirs("Red-DiscordBot").user_config_dir)
 config_file = config_dir / 'config.json'
 
 
-def load_basic_configuration(instance_name: str):
+def load_basic_configuration(instance_name_: str):
     global jsonio
     global basic_config
+    global instance_name
 
     jsonio = JsonIO(config_file)
 
+    instance_name = instance_name_
+
     try:
         config = jsonio._load_json()
         basic_config = config[instance_name]

redbot/core/drivers/__init__.py

@@ -0,0 +1,31 @@
+__all__ = ["get_driver"]
+
+
+def get_driver(type, *args, **kwargs):
+    """
+    Selectively import/load driver classes based on the selected type. This
+    is required so that dependencies can differ between installs (e.g. so that
+    you don't need to install a mongo dependency if you will just be running a
+    json data backend).
+
+    .. note::
+
+        See the respective classes for information on what ``args`` and ``kwargs``
+        should be.
+
+    :param str type:
+        One of: json, mongo
+    :param args:
+        Dependent on driver type.
+    :param kwargs:
+        Dependent on driver type.
+    :return:
+        Subclass of :py:class:`.red_base.BaseDriver`.
+    """
+    if type == "JSON":
+        from .red_json import JSON
+        return JSON(*args, **kwargs)
+    elif type == "MongoDB":
+        from .red_mongo import Mongo
+        return Mongo(*args, **kwargs)
+    raise RuntimeError("Invalid driver type: '{}'".format(type))

redbot/core/drivers/red_base.py

@@ -3,11 +3,40 @@ from typing import Tuple
 __all__ = ["BaseDriver"]
 
 class BaseDriver:
+    def __init__(self, cog_name):
+        self.cog_name = cog_name
+
     def get_driver(self):
         raise NotImplementedError
 
     async def get(self, identifiers: Tuple[str]):
+        """
+        Finds the value indicate by the given identifiers.
+
+        :param identifiers:
+            A list of identifiers that correspond to nested dict accesses.
+        :return:
+            Stored value.
+        """
+        raise NotImplementedError
+
+    def get_config_details(self):
+        """
+        Asks users for additional configuration information necessary
+        to use this config driver.
+
+        :return:
+            Dict of configuration details.
+        """
         raise NotImplementedError
 
     async def set(self, identifiers: Tuple[str], value):
+        """
+        Sets the value of the key indicated by the given identifiers.
+
+        :param identifiers:
+            A list of identifiers that correspond to nested dict accesses.
+        :param value:
+            Any JSON serializable python object.
+        """
         raise NotImplementedError

redbot/core/drivers/red_json.py

@@ -9,10 +9,20 @@ __all__ = ["JSON"]
 
 
 class JSON(BaseDriver):
+    """
+    Subclass of :py:class:`.red_base.BaseDriver`.
+
+    .. py:attribute:: file_name
+
+        The name of the file in which to store JSON data.
+
+    .. py:attribute:: data_path
+
+        The path in which to store the file indicated by :py:attr:`file_name`.
+    """
     def __init__(self, cog_name, *, data_path_override: Path=None,
                  file_name_override: str="settings.json"):
-        super().__init__()
+        super().__init__(cog_name)
-        self.cog_name = cog_name
         self.file_name = file_name_override
         if data_path_override:
             self.data_path = data_path_override

redbot/core/drivers/red_mongo.py

@@ -1,215 +1,129 @@
-import pymongo as m
+from typing import Tuple
+
+import motor.motor_asyncio
 from .red_base import BaseDriver
 
-
+__all__ = ["Mongo"]
-__all__ = ["Mongo", "RedMongoException", "MultipleMatches",
-           "MissingCollection"]
-
-
-class RedMongoException(Exception):
-    """Base Red Mongo Exception class"""
-    pass
-
-
-class MultipleMatches(RedMongoException):
-    """Raised when multiple documents match a single cog_name and
-        cog_identifier pair."""
-    pass
-
-
-class MissingCollection(RedMongoException):
-    """Raised when a collection is missing from the mongo db"""
-    pass
 
 
 class Mongo(BaseDriver):
-    def __init__(self, host, port=27017, admin_user=None, admin_pass=None,
+    """
-                 **kwargs):
+    Subclass of :py:class:`.red_base.BaseDriver`.
-        self.conn = m.MongoClient(host=host, port=port, **kwargs)
+    """
+    def __init__(self, cog_name, **kwargs):
+        super().__init__(cog_name)
+        self.host = kwargs['HOST']
+        self.port = kwargs['PORT']
+        admin_user = kwargs['USERNAME']
+        admin_pass = kwargs['PASSWORD']
+
+        from ..data_manager import instance_name
+
+        self.instance_name = instance_name
+
+        self.conn = None
 
         self.admin_user = admin_user
         self.admin_pass = admin_pass
 
-        self._db = self.conn.red
+    async def _authenticate(self):
-        if self.admin_user is not None and self.admin_pass is not None:
+        self.conn = motor.motor_asyncio.AsyncIOMotorClient(host=self.host, port=self.port)
-            self._db.authenticate(self.admin_user, self.admin_pass)
 
-        self._global = self._db.GLOBAL
+        if None not in (self.admin_pass, self.admin_user):
-        self._guild = self._db.GUILD
+            await self.db.authenticate(self.admin_user, self.admin_pass)
-        self._channel = self._db.CHANNEL
-        self._role = self._db.ROLE
-        self._member = self._db.MEMBER
-        self._user = self._db.USER
 
-    def get_global(self, cog_name, cog_identifier, _, key, *, default=None):
+    async def _ensure_connected(self):
-        doc = self._global.find(
+        if self.conn is None:
-            {"cog_name": cog_name, "cog_identifier": cog_identifier},
+            await self._authenticate()
-            projection=[key, ], batch_size=2)
-        if doc.count() == 2:
-            raise MultipleMatches("Too many matching documents at the GLOBAL"
-                                  " level: ({}, {})".format(cog_name,
-                                                            cog_identifier))
-        elif doc.count() == 1:
-            return doc[0].get(key, default)
-        return default
 
-    def get_guild(self, cog_name, cog_identifier, guild_id, key, *,
+    @property
-                   default=None):
+    def db(self) -> motor.core.Database:
-        doc = self._guild.find(
+        """
-            {"cog_name": cog_name, "cog_identifier": cog_identifier,
+        Gets the mongo database for this cog's name.
-             "guild_id": guild_id},
-            projection=[key, ], batch_size=2)
-        if doc.count() == 2:
-            raise MultipleMatches("Too many matching documents at the GUILD"
-                                  " level: ({}, {}, {})".format(
-                                      cog_name, cog_identifier, guild_id))
-        elif doc.count() == 1:
-            return doc[0].get(key, default)
-        return default
 
-    def get_channel(self, cog_name, cog_identifier, channel_id, key, *,
+        .. warning::
-                    default=None):
-        doc = self._channel.find(
-            {"cog_name": cog_name, "cog_identifier": cog_identifier,
-             "channel_id": channel_id},
-            projection=[key, ], batch_size=2)
-        if doc.count() == 2:
-            raise MultipleMatches("Too many matching documents at the CHANNEL"
-                                  " level: ({}, {}, {})".format(
-                                      cog_name, cog_identifier, channel_id))
-        elif doc.count() == 1:
-            return doc[0].get(key, default)
-        return default
 
-    def get_role(self, cog_name, cog_identifier, role_id, key, *,
+            Right now this will cause a new connection to be made every time the
-                 default=None):
+            database is accessed. We will want to create a connection pool down the
-        doc = self._role.find(
+            line to limit the number of connections.
-            {"cog_name": cog_name, "cog_identifier": cog_identifier,
-             "role_id": role_id},
-            projection=[key, ], batch_size=2)
-        if doc.count() == 2:
-            raise MultipleMatches("Too many matching documents at the ROLE"
-                                  " level: ({}, {}, {})".format(
-                                      cog_name, cog_identifier, role_id))
-        elif doc.count() == 1:
-            return doc[0].get(key, default)
-        return default
 
-    def get_member(self, cog_name, cog_identifier, user_id, guild_id, key, *,
+        :return:
-                   default=None):
+            PyMongo Database object.
-        doc = self._member.find(
+        """
-            {"cog_name": cog_name, "cog_identifier": cog_identifier,
+        db_name = "RED_{}".format(self.instance_name)
-             "user_id": user_id, "guild_id": guild_id},
+        return self.conn[db_name]
-            projection=[key, ], batch_size=2)
-        if doc.count() == 2:
-            raise MultipleMatches("Too many matching documents at the MEMBER"
-                                  " level: ({}, {}, mid {}, sid {})".format(
-                                      cog_name, cog_identifier, user_id,
-                                      guild_id))
-        elif doc.count() == 1:
-            return doc[0].get(key, default)
-        return default
 
-    def get_user(self, cog_name, cog_identifier, user_id, key, *,
+    def get_collection(self) -> motor.core.Collection:
-                 default=None):
+        """
-        doc = self._user.find(
+        Gets a specified collection within the PyMongo database for this cog.
-            {"cog_name": cog_name, "cog_identifier": cog_identifier,
-             "user_id": user_id},
-            projection=[key, ], batch_size=2)
-        if doc.count() == 2:
-            raise MultipleMatches("Too many matching documents at the USER"
-                                  " level: ({}, {}, mid {})".format(
-                                      cog_name, cog_identifier, user_id))
-        elif doc.count() == 1:
-            return doc[0].get(key, default)
-        else:
-            return default
 
-    def set_global(self, cog_name, cog_identifier, key, value, clear=False):
+        Unless you are doing custom stuff ``collection_name`` should be one of the class
-        filter = {"cog_name": cog_name, "cog_identifier": cog_identifier}
+        attributes of :py:class:`core.config.Config`.
-        data = {"$set": {key: value}}
-        if self._global.count(filter) > 1:
-            raise MultipleMatches("Too many matching documents at the GLOBAL"
-                                  " level: ({}, {})".format(cog_name,
-                                                            cog_identifier))
-        else:
-            if clear:
-                self._global.delete_one(filter)
-            else:
-                self._global.update_one(filter, data, upsert=True)
 
-    def set_guild(self, cog_name, cog_identifier, guild_id, key, value,
+        :param str collection_name:
-                   clear=False):
+        :return:
-        filter = {"cog_name": cog_name, "cog_identifier": cog_identifier,
+            PyMongo collection object.
-                  "guild_id": guild_id}
+        """
-        data = {"$set": {key: value}}
+        return self.db[self.cog_name]
-        if self._guild.count(filter) > 1:
-            raise MultipleMatches("Too many matching documents at the GUILD"
-                                  " level: ({}, {}, {})".format(
-                                      cog_name, cog_identifier, guild_id))
-        else:
-            if clear:
-                self._guild.delete_one(filter)
-            else:
-                self._guild.update_one(filter, data, upsert=True)
 
-    def set_channel(self, cog_name, cog_identifier, channel_id, key, value,
+    @staticmethod
-                    clear=False):
+    def _parse_identifiers(identifiers):
-        filter = {"cog_name": cog_name, "cog_identifier": cog_identifier,
+        uuid, identifiers = identifiers[0], identifiers[1:]
-                  "channel_id": channel_id}
+        return uuid, identifiers
-        data = {"$set": {key: value}}
-        if self._channel.count(filter) > 1:
-            raise MultipleMatches("Too many matching documents at the CHANNEL"
-                                  " level: ({}, {}, {})".format(
-                                      cog_name, cog_identifier, channel_id))
-        else:
-            if clear:
-                self._channel.delete_one(filter)
-            else:
-                self._channel.update_one(filter, data, upsert=True)
 
-    def set_role(self, cog_name, cog_identifier, role_id, key, value,
+    async def get(self, identifiers: Tuple[str]):
-                 clear=False):
+        await self._ensure_connected()
-        filter = {"cog_name": cog_name, "cog_identifier": cog_identifier,
+        uuid, identifiers = self._parse_identifiers(identifiers)
-                  "role_id": role_id}
-        data = {"$set": {key: value}}
-        if self._role.count(filter) > 1:
-            raise MultipleMatches("Too many matching documents at the ROLE"
-                                  " level: ({}, {}, {})".format(
-                                      cog_name, cog_identifier, role_id))
-        else:
-            if clear:
-                self._role.delete_one(filter)
-            else:
-                self._role.update_one(filter, data, upsert=True)
 
-    def set_member(self, cog_name, cog_identifier, user_id, guild_id, key,
+        mongo_collection = self.get_collection()
-                   value, clear=False):
-        filter = {"cog_name": cog_name, "cog_identifier": cog_identifier,
-                  "guild_id": guild_id, "user_id": user_id}
-        data = {"$set": {key: value}}
-        if self._member.count(filter) > 1:
-            raise MultipleMatches("Too many matching documents at the MEMBER"
-                                  " level: ({}, {}, mid {}, sid {})".format(
-                                      cog_name, cog_identifier, user_id,
-                                      guild_id))
-        else:
-            if clear:
-                self._member.delete_one(filter)
-            else:
-                self._member.update_one(filter, data, upsert=True)
 
-    def set_user(self, cog_name, cog_identifier, user_id, key, value,
+        dot_identifiers = '.'.join(identifiers)
-                 clear=False):
+
-        filter = {"cog_name": cog_name, "cog_identifier": cog_identifier,
+        partial = await mongo_collection.find_one(
-                  "user_id": user_id}
+            filter={'_id': uuid},
-        data = {"$set": {key: value}}
+            projection={dot_identifiers: True}
-        if self._user.count(filter) > 1:
+        )
-            raise MultipleMatches("Too many matching documents at the USER"
+
-                                  " level: ({}, {}, mid {})".format(
+        if partial is None:
-                                      cog_name, cog_identifier, user_id))
+            raise KeyError("No matching document was found and Config expects"
-        else:
+                           " a KeyError.")
-            if clear:
+
-                self._user.delete_one(filter)
+        for i in identifiers:
-            else:
+            partial = partial[i]
-                self._user.update_one(filter, data, upsert=True)
+        return partial
+
+    async def set(self, identifiers: Tuple[str], value):
+        await self._ensure_connected()
+        uuid, identifiers = self._parse_identifiers(identifiers)
+
+        dot_identifiers = '.'.join(identifiers)
+
+        mongo_collection = self.get_collection()
+
+        await mongo_collection.update_one(
+            {'_id': uuid},
+            update={"$set": {dot_identifiers: value}},
+            upsert=True
+        )
+
+    def get_driver(self):
+        return self
+
+
+def get_config_details():
+    host = input("Enter host address: ")
+    port = int(input("Enter host port: "))
+
+    admin_uname = input("Enter login username: ")
+    admin_password = input("Enter login password: ")
+
+    if admin_uname == "":
+        admin_uname = admin_password = None
+
+    ret = {
+        'HOST': host,
+        'PORT': port,
+        'USERNAME': admin_uname,
+        'PASSWORD': admin_password
+    }
+    return ret

redbot/setup.py

@@ -89,7 +89,13 @@ def basic_setup():
             if storage not in storage_dict:
                 storage = None
 
-    default_dirs['STORAGE_TYPE'] = storage_dict[storage]
+    default_dirs['STORAGE_TYPE'] = storage_dict.get(storage, 1)
+
+    if storage_dict.get(storage, 1) == "MongoDB":
+        from redbot.core.drivers.red_mongo import get_config_details
+        default_dirs['STORAGE_DETAILS'] = get_config_details()
+    else:
+        default_dirs['STORAGE_DETAILS'] = {}
 
     name = ""
     while len(name) == 0:

requirements.txt

@@ -1,3 +1,4 @@
 appdirs
 youtube_dl
-raven
+raven
+colorama

setup.py

@@ -106,7 +106,7 @@ setup(
     ],
     extras_require={
         'test': ['pytest>=3', 'pytest-asyncio'],
-        'mongo': ['pymongo', 'motor'],
+        'mongo': ['motor'],
         'docs': ['sphinx', 'sphinxcontrib-asyncio', 'sphinx_rtd_theme'],
         'voice': ['PyNaCl']
     }