autonomous-app 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

autonomous/db/connection.py

@@ -0,0 +1,472 @@
+import warnings
+
+from pymongo import MongoClient, ReadPreference, uri_parser
+from pymongo.common import _UUID_REPRESENTATIONS
+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 autonomous.db
+from autonomous.db.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,
+        )
+        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 autonomous.db import Document
+    from autonomous.db.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", autonomous.db.__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/autonomous.db/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