ayon-core/openpype/lib/mongo.py

import os
import sys
import time
import logging
import pymongo

if sys.version_info[0] == 2:
    from urlparse import urlparse, parse_qs
else:
    from urllib.parse import urlparse, parse_qs


class MongoEnvNotSet(Exception):
    pass


def decompose_url(url):
    components = {
        "scheme": None,
        "host": None,
        "port": None,
        "username": None,
        "password": None,
        "auth_db": None
    }

    result = urlparse(url)
    if result.scheme is None:
        _url = "mongodb://{}".format(url)
        result = urlparse(_url)

    components["scheme"] = result.scheme
    components["host"] = result.hostname
    try:
        components["port"] = result.port
    except ValueError:
        raise RuntimeError("invalid port specified")
    components["username"] = result.username
    components["password"] = result.password

    try:
        components["auth_db"] = parse_qs(result.query)['authSource'][0]
    except KeyError:
        # no auth db provided, mongo will use the one we are connecting to
        pass

    return components


def compose_url(scheme=None,
                host=None,
                username=None,
                password=None,
                port=None,
                auth_db=None):

    url = "{scheme}://"

    if username and password:
        url += "{username}:{password}@"

    url += "{host}"
    if port:
        url += ":{port}"

    if auth_db:
        url += "?authSource={auth_db}"

    return url.format(**{
        "scheme": scheme,
        "host": host,
        "username": username,
        "password": password,
        "port": port,
        "auth_db": auth_db
    })


def get_default_components():
    mongo_url = os.environ.get("OPENPYPE_MONGO")
    if mongo_url is None:
        raise MongoEnvNotSet(
            "URL for Mongo logging connection is not set."
        )
    return decompose_url(mongo_url)


def extract_port_from_url(url):
    parsed_url = urlparse(url)
    if parsed_url.scheme is None:
        _url = "mongodb://{}".format(url)
        parsed_url = urlparse(_url)
    return parsed_url.port


def validate_mongo_connection(mongo_uri):
    """Check if provided mongodb URL is valid.

    Args:
        mongo_uri (str): URL to validate.

    Raises:
        ValueError: When port in mongo uri is not valid.
        pymongo.errors.InvalidURI: If passed mongo is invalid.
        pymongo.errors.ServerSelectionTimeoutError: If connection timeout
            passed so probably couldn't connect to mongo server.

    """
    parsed = urlparse(mongo_uri)
    # Force validation of scheme
    if parsed.scheme not in ["mongodb", "mongodb+srv"]:
        raise pymongo.errors.InvalidURI((
            "Invalid URI scheme:"
            " URI must begin with 'mongodb://' or 'mongodb+srv://'"
        ))
    # we have mongo connection string. Let's try if we can connect.
    components = decompose_url(mongo_uri)
    mongo_args = {
        "host": compose_url(**components),
        "serverSelectionTimeoutMS": 1000
    }
    port = components.get("port")
    if port is not None:
        mongo_args["port"] = int(port)

    # Create connection
    client = pymongo.MongoClient(**mongo_args)
    client.server_info()
    client.close()


class OpenPypeMongoConnection:
    """Singleton MongoDB connection.

    Keeps MongoDB connections by url.
    """
    mongo_clients = {}
    log = logging.getLogger("OpenPypeMongoConnection")

    @staticmethod
    def get_default_mongo_url():
        return os.environ["OPENPYPE_MONGO"]

    @classmethod
    def get_mongo_client(cls, mongo_url=None):
        if mongo_url is None:
            mongo_url = cls.get_default_mongo_url()

        connection = cls.mongo_clients.get(mongo_url)
        if connection:
            # Naive validation of existing connection
            try:
                connection.server_info()
            except Exception:
                connection = None

        if not connection:
            cls.log.debug("Creating mongo connection to {}".format(mongo_url))
            connection = cls.create_connection(mongo_url)
            cls.mongo_clients[mongo_url] = connection

        return connection

    @classmethod
    def create_connection(cls, mongo_url, timeout=None):
        if timeout is None:
            timeout = int(os.environ.get("AVALON_TIMEOUT") or 1000)

        kwargs = {
            "host": mongo_url,
            "serverSelectionTimeoutMS": timeout
        }

        port = extract_port_from_url(mongo_url)
        if port is not None:
            kwargs["port"] = int(port)

        mongo_client = pymongo.MongoClient(**kwargs)

        for _retry in range(3):
            try:
                t1 = time.time()
                mongo_client.server_info()

            except Exception:
                cls.log.warning("Retrying...")
                time.sleep(1)
                timeout *= 1.5

            else:
                break

        else:
            raise IOError((
                "ERROR: Couldn't connect to {} in less than {:.3f}ms"
            ).format(mongo_url, timeout))

        cls.log.info("Connected to {}, delay {:.3f}s".format(
            mongo_url, time.time() - t1
        ))
        return mongo_client