connection.py

import collections
import threading
import warnings

from pymongo import MongoClient, ReadPreference, uri_parser
from pymongo.common import _UUID_REPRESENTATIONS

try:
    from pymongo.database_shared import _check_name
except ImportError:
    from pymongo.database import _check_name

# DriverInfo was added in PyMongo 3.7.
try:
    from pymongo.driver_info import DriverInfo
except ImportError:
    DriverInfo = None

import mongoengine
from mongoengine.pymongo_support import PYMONGO_VERSION

__all__ = [
    "DEFAULT_CONNECTION_NAME",
    "DEFAULT_DATABASE_NAME",
    "ConnectionFailure",
    "connect",
    "disconnect",
    "disconnect_all",
    "get_connection",
    "get_db",
    "register_connection",
]


DEFAULT_CONNECTION_NAME = "default"
DEFAULT_DATABASE_NAME = "test"
DEFAULT_HOST = "localhost"
DEFAULT_PORT = 27017

_connection_settings = {}
_connections = {}
_dbs = {}


READ_PREFERENCE = ReadPreference.PRIMARY


class ConnectionFailure(Exception):
    """Error raised when the database connection can't be established or
    when a connection with a requested alias can't be retrieved.
    """

    pass


def _check_db_name(name):
    """Check if a database name is valid.
    This functionality is copied from pymongo Database class constructor.
    """
    if not isinstance(name, str):
        raise TypeError("name must be an instance of %s" % str)
    elif name != "$external":
        _check_name(name)


def _get_connection_settings(
    db=None,
    name=None,
    host=None,
    port=None,
    read_preference=READ_PREFERENCE,
    username=None,
    password=None,
    authentication_source=None,
    authentication_mechanism=None,
    authmechanismproperties=None,
    **kwargs,
):
    """Get the connection settings as a dict

    :param db: the name of the database to use, for compatibility with connect
    :param name: the name of the specific database to use
    :param host: the host name of the: program: `mongod` instance to connect to
    :param port: the port that the: program: `mongod` instance is running on
    :param read_preference: The read preference for the collection
    :param username: username to authenticate with
    :param password: password to authenticate with
    :param authentication_source: database to authenticate against
    :param authentication_mechanism: database authentication mechanisms.
        By default, use SCRAM-SHA-1 with MongoDB 3.0 and later,
        MONGODB-CR (MongoDB Challenge Response protocol) for older servers.
    :param mongo_client_class: using alternative connection client other than
        pymongo.MongoClient, e.g. mongomock, montydb, that provides pymongo alike
        interface but not necessarily for connecting to a real mongo instance.
    :param kwargs: ad-hoc parameters to be passed into the pymongo driver,
        for example maxpoolsize, tz_aware, etc. See the documentation
        for pymongo's `MongoClient` for a full list.
    """
    conn_settings = {
        "name": name or db or DEFAULT_DATABASE_NAME,
        "host": host or DEFAULT_HOST,
        "port": port or DEFAULT_PORT,
        "read_preference": read_preference,
        "username": username,
        "password": password,
        "authentication_source": authentication_source,
        "authentication_mechanism": authentication_mechanism,
        "authmechanismproperties": authmechanismproperties,
    }

    _check_db_name(conn_settings["name"])
    conn_host = conn_settings["host"]

    # Host can be a list or a string, so if string, force to a list.
    if isinstance(conn_host, str):
        conn_host = [conn_host]

    resolved_hosts = []
    for entity in conn_host:
        # Reject old mongomock integration
        # To be removed in a few versions after 0.27.0
        if entity.startswith("mongomock://") or kwargs.get("is_mock"):
            raise Exception(
                "Use of mongomock:// URI or 'is_mock' were removed in favor of 'mongo_client_class=mongomock.MongoClient'. "
                "Check the CHANGELOG for more info"
            )

        # Handle URI style connections, only updating connection params which
        # were explicitly specified in the URI.
        if "://" in entity:
            uri_dict = uri_parser.parse_uri(entity)
            resolved_hosts.append(entity)

            database = uri_dict.get("database")
            if database:
                conn_settings["name"] = database

            for param in ("read_preference", "username", "password"):
                if uri_dict.get(param):
                    conn_settings[param] = uri_dict[param]

            uri_options = uri_dict[
                "options"
            ]  # uri_options is a _CaseInsensitiveDictionary
            if "replicaset" in uri_options:
                conn_settings["replicaSet"] = uri_options["replicaset"]
            if "authsource" in uri_options:
                conn_settings["authentication_source"] = uri_options["authsource"]
            if "authmechanism" in uri_options:
                conn_settings["authentication_mechanism"] = uri_options["authmechanism"]
            if "readpreference" in uri_options:
                read_preferences = (
                    ReadPreference.NEAREST,
                    ReadPreference.PRIMARY,
                    ReadPreference.PRIMARY_PREFERRED,
                    ReadPreference.SECONDARY,
                    ReadPreference.SECONDARY_PREFERRED,
                )

                # Starting with PyMongo v3.5, the "readpreference" option is
                # returned as a string (e.g. "secondaryPreferred") and not an
                # int (e.g. 3).
                # TODO simplify the code below once we drop support for
                # PyMongo v3.4.
                read_pf_mode = uri_options["readpreference"]
                if isinstance(read_pf_mode, str):
                    read_pf_mode = read_pf_mode.lower()
                for preference in read_preferences:
                    if (
                        preference.name.lower() == read_pf_mode
                        or preference.mode == read_pf_mode
                    ):
                        ReadPrefClass = preference.__class__
                        break

                if "readpreferencetags" in uri_options:
                    conn_settings["read_preference"] = ReadPrefClass(
                        tag_sets=uri_options["readpreferencetags"]
                    )
                else:
                    conn_settings["read_preference"] = ReadPrefClass()

            if "authmechanismproperties" in uri_options:
                conn_settings["authmechanismproperties"] = uri_options[
                    "authmechanismproperties"
                ]
            if "uuidrepresentation" in uri_options:
                REV_UUID_REPRESENTATIONS = {
                    v: k for k, v in _UUID_REPRESENTATIONS.items()
                }
                conn_settings["uuidrepresentation"] = REV_UUID_REPRESENTATIONS[
                    uri_options["uuidrepresentation"]
                ]
        else:
            resolved_hosts.append(entity)
    conn_settings["host"] = resolved_hosts

    # Deprecated parameters that should not be passed on
    kwargs.pop("slaves", None)
    kwargs.pop("is_slave", None)

    keys = {
        key.lower() for key in kwargs.keys()
    }  # pymongo options are case insensitive
    if "uuidrepresentation" not in keys and "uuidrepresentation" not in conn_settings:
        warnings.warn(
            "No uuidRepresentation is specified! Falling back to "
            "'pythonLegacy' which is the default for pymongo 3.x. "
            "For compatibility with other MongoDB drivers this should be "
            "specified as 'standard' or '{java,csharp}Legacy' to work with "
            "older drivers in those languages. This will be changed to "
            "'unspecified' in a future release.",
            DeprecationWarning,
            stacklevel=3,
        )
        kwargs["uuidRepresentation"] = "pythonLegacy"

    conn_settings.update(kwargs)
    return conn_settings


def register_connection(
    alias,
    db=None,
    name=None,
    host=None,
    port=None,
    read_preference=READ_PREFERENCE,
    username=None,
    password=None,
    authentication_source=None,
    authentication_mechanism=None,
    authmechanismproperties=None,
    **kwargs,
):
    """Register the connection settings.

    :param alias: the name that will be used to refer to this connection throughout MongoEngine
    :param db: the name of the database to use, for compatibility with connect
    :param name: the name of the specific database to use
    :param host: the host name of the: program: `mongod` instance to connect to
    :param port: the port that the: program: `mongod` instance is running on
    :param read_preference: The read preference for the collection
    :param username: username to authenticate with
    :param password: password to authenticate with
    :param authentication_source: database to authenticate against
    :param authentication_mechanism: database authentication mechanisms.
        By default, use SCRAM-SHA-1 with MongoDB 3.0 and later,
        MONGODB-CR (MongoDB Challenge Response protocol) for older servers.
    :param mongo_client_class: using alternative connection client other than
        pymongo.MongoClient, e.g. mongomock, montydb, that provides pymongo alike
        interface but not necessarily for connecting to a real mongo instance.
    :param kwargs: ad-hoc parameters to be passed into the pymongo driver,
        for example maxpoolsize, tz_aware, etc. See the documentation
        for pymongo's `MongoClient` for a full list.
    """
    conn_settings = _get_connection_settings(
        db=db,
        name=name,
        host=host,
        port=port,
        read_preference=read_preference,
        username=username,
        password=password,
        authentication_source=authentication_source,
        authentication_mechanism=authentication_mechanism,
        authmechanismproperties=authmechanismproperties,
        **kwargs,
    )
    _connection_settings[alias] = conn_settings


def disconnect(alias=DEFAULT_CONNECTION_NAME):
    """Close the connection with a given alias."""
    from mongoengine import Document
    from mongoengine.base.common import _get_documents_by_db

    connection = _connections.pop(alias, None)
    if connection:
        # MongoEngine may share the same MongoClient across multiple aliases
        # if connection settings are the same so we only close
        # the client if we're removing the final reference.
        # Important to use 'is' instead of '==' because clients connected to the same cluster
        # will compare equal even with different options
        if all(connection is not c for c in _connections.values()):
            connection.close()

    if alias in _dbs:
        # Detach all cached collections in Documents
        for doc_cls in _get_documents_by_db(alias, DEFAULT_CONNECTION_NAME):
            if issubclass(doc_cls, Document):  # Skip EmbeddedDocument
                doc_cls._disconnect()

        del _dbs[alias]

    if alias in _connection_settings:
        del _connection_settings[alias]


def disconnect_all():
    """Close all registered database."""
    for alias in list(_connections.keys()):
        disconnect(alias)


def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
    """Return a connection with a given alias."""

    # Connect to the database if not already connected
    if reconnect:
        disconnect(alias)

    # If the requested alias already exists in the _connections list, return
    # it immediately.
    if alias in _connections:
        return _connections[alias]

    # Validate that the requested alias exists in the _connection_settings.
    # Raise ConnectionFailure if it doesn't.
    if alias not in _connection_settings:
        if alias == DEFAULT_CONNECTION_NAME:
            msg = "You have not defined a default connection"
        else:
            msg = 'Connection with alias "%s" has not been defined' % alias
        raise ConnectionFailure(msg)

    def _clean_settings(settings_dict):
        if PYMONGO_VERSION < (4,):
            irrelevant_fields_set = {
                "name",
                "username",
                "password",
                "authentication_source",
                "authentication_mechanism",
                "authmechanismproperties",
            }
            rename_fields = {}
        else:
            irrelevant_fields_set = {"name"}
            rename_fields = {
                "authentication_source": "authSource",
                "authentication_mechanism": "authMechanism",
            }
        return {
            rename_fields.get(k, k): v
            for k, v in settings_dict.items()
            if k not in irrelevant_fields_set and v is not None
        }

    raw_conn_settings = _connection_settings[alias].copy()

    # Retrieve a copy of the connection settings associated with the requested
    # alias and remove the database name and authentication info (we don't
    # care about them at this point).
    conn_settings = _clean_settings(raw_conn_settings)
    if DriverInfo is not None:
        conn_settings.setdefault(
            "driver", DriverInfo("MongoEngine", mongoengine.__version__)
        )

    # Determine if we should use PyMongo's or mongomock's MongoClient.
    if "mongo_client_class" in conn_settings:
        mongo_client_class = conn_settings.pop("mongo_client_class")
    else:
        mongo_client_class = MongoClient

    # Re-use existing connection if one is suitable.
    existing_connection = _find_existing_connection(raw_conn_settings)
    if existing_connection:
        connection = existing_connection
    else:
        connection = _create_connection(
            alias=alias, mongo_client_class=mongo_client_class, **conn_settings
        )
    _connections[alias] = connection
    return _connections[alias]


def _create_connection(alias, mongo_client_class, **connection_settings):
    """
    Create the new connection for this alias. Raise
    ConnectionFailure if it can't be established.
    """
    try:
        return mongo_client_class(**connection_settings)
    except Exception as e:
        raise ConnectionFailure(f"Cannot connect to database {alias} :\n{e}")


def _find_existing_connection(connection_settings):
    """
    Check if an existing connection could be reused

    Iterate over all of the connection settings and if an existing connection
    with the same parameters is suitable, return it

    :param connection_settings: the settings of the new connection
    :return: An existing connection or None
    """
    connection_settings_bis = (
        (db_alias, settings.copy())
        for db_alias, settings in _connection_settings.items()
    )

    def _clean_settings(settings_dict):
        # Only remove the name but it's important to
        # keep the username/password/authentication_source/authentication_mechanism
        # to identify if the connection could be shared (cfr https://github.com/MongoEngine/mongoengine/issues/2047)
        return {k: v for k, v in settings_dict.items() if k != "name"}

    cleaned_conn_settings = _clean_settings(connection_settings)
    for db_alias, connection_settings in connection_settings_bis:
        db_conn_settings = _clean_settings(connection_settings)
        if cleaned_conn_settings == db_conn_settings and _connections.get(db_alias):
            return _connections[db_alias]


def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
    if reconnect:
        disconnect(alias)

    if alias not in _dbs:
        conn = get_connection(alias)
        conn_settings = _connection_settings[alias]
        db = conn[conn_settings["name"]]
        # Authenticate if necessary
        if (
            PYMONGO_VERSION < (4,)
            and conn_settings["username"]
            and (
                conn_settings["password"]
                or conn_settings["authentication_mechanism"] == "MONGODB-X509"
            )
            and conn_settings["authmechanismproperties"] is None
        ):
            auth_kwargs = {"source": conn_settings["authentication_source"]}
            if conn_settings["authentication_mechanism"] is not None:
                auth_kwargs["mechanism"] = conn_settings["authentication_mechanism"]
            db.authenticate(
                conn_settings["username"], conn_settings["password"], **auth_kwargs
            )
        _dbs[alias] = db
    return _dbs[alias]


def connect(db=None, alias=DEFAULT_CONNECTION_NAME, **kwargs):
    """Connect to the database specified by the 'db' argument.

    Connection settings may be provided here as well if the database is not
    running on the default port on localhost. If authentication is needed,
    provide username and password arguments as well.

    Multiple databases are supported by using aliases. Provide a separate
    `alias` to connect to a different instance of: program: `mongod`.

    In order to replace a connection identified by a given alias, you'll
    need to call ``disconnect`` first

    See the docstring for `register_connection` for more details about all
    supported kwargs.
    """
    if alias in _connections:
        prev_conn_setting = _connection_settings[alias]
        new_conn_settings = _get_connection_settings(db, **kwargs)

        if new_conn_settings != prev_conn_setting:
            err_msg = (
                "A different connection with alias `{}` was already "
                "registered. Use disconnect() first"
            ).format(alias)
            raise ConnectionFailure(err_msg)
    else:
        register_connection(alias, db, **kwargs)

    return get_connection(alias)


# Support old naming convention
_get_connection = get_connection
_get_db = get_db


class _LocalSessions(threading.local):
    def __init__(self):
        self.sessions = collections.deque()

    def append(self, session):
        self.sessions.append(session)

    def get_current(self):
        if len(self.sessions):
            return self.sessions[-1]

    def clear_current(self):
        if len(self.sessions):
            self.sessions.pop()

    def clear_all(self):
        self.sessions.clear()


_local_sessions = _LocalSessions()


def _set_session(session):
    _local_sessions.append(session)


def _get_session():
    return _local_sessions.get_current()


def _clear_session():
    return _local_sessions.clear_current()