<commit_before>from pymongo import MongoClient, ReadPreference, uri_parser
from mongoengine.python_support import (IS_PYMONGO_3, str_types)

__all__ = ['ConnectionError', 'connect', 'register_connection',
           'DEFAULT_CONNECTION_NAME']


DEFAULT_CONNECTION_NAME = 'default'
if IS_PYMONGO_3:
    READ_PREFERENCE = ReadPreference.PRIMARY
else:
    from pymongo import MongoReplicaSetClient
    READ_PREFERENCE = False


class ConnectionError(Exception):
    pass


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


def register_connection(alias, name=None, host=None, port=None,
                        read_preference=READ_PREFERENCE,
                        username=None, password=<PASSWORD>, authentication_source=None,
                        **kwargs):
    """Add a connection.

    :param alias: the name that will be used to refer to this connection
        throughout MongoEngine
    :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
       ** Added pymongo 2.1
    :param username: username to authenticate with
    :param password: <PASSWORD>
    :param authentication_source: database to authenticate against
    :param is_mock: explicitly use mongomock for this connection
        (can also be done by using `mongomock://` as db host prefix)
    :param kwargs: allow ad-hoc parameters to be passed into the pymongo driver

    .. versionchanged:: 0.10.6 - added mongomock support
    """
    global _connection_settings

    conn_settings = {
        'name': name or 'test',
        'host': host or 'localhost',
        'port': port or 27017,
        'read_preference': read_preference,
        'username': username,
        'password': password,
        'authentication_source': authentication_source
    }

    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_types):
        conn_host = [conn_host]

    resolved_hosts = []
    for entity in conn_host:
        # Handle uri style connections
        if entity.startswith('mongomock://'):
            conn_settings['is_mock'] = True
            # `mongomock://` is not a valid url prefix and must be replaced by `mongodb://`
            resolved_hosts.append(entity.replace('mongomock://', 'mongodb://', 1))
        elif '://' in entity:
            uri_dict = uri_parser.parse_uri(entity)
            resolved_hosts.append(entity)
            conn_settings.update({
                'name': uri_dict.get('database') or name,
                'username': uri_dict.get('username'),
                'password': uri_dict.get('password'),
                'read_preference': read_preference,
            })
            uri_options = uri_dict['options']
            if 'replicaset' in uri_options:
                conn_settings['replicaSet'] = True
            if 'authsource' in uri_options:
                conn_settings['authentication_source'] = uri_options['authsource']
        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)

    conn_settings.update(kwargs)
    _connection_settings[alias] = conn_settings


def disconnect(alias=DEFAULT_CONNECTION_NAME):
    global _connections
    global _dbs

    if alias in _connections:
        get_connection(alias=alias).close()
        del _connections[alias]
    if alias in _dbs:
        del _dbs[alias]


def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
    global _connections
    # Connect to the database if not already connected
    if reconnect:
        disconnect(alias)

    if alias not in _connections:
        if alias not in _connection_settings:
            msg = 'Connection with alias "%s" has not been defined' % alias
            if alias == DEFAULT_CONNECTION_NAME:
                msg = 'You have not defined a default connection'
            raise ConnectionError(msg)
        conn_settings = _connection_settings[alias].copy()

        conn_settings.pop('name', None)
        conn_settings.pop('username', None)
        conn_settings.pop('password', None)
        conn_settings.pop('authentication_source', None)

        is_mock = conn_settings.pop('is_mock', None)
        if is_mock:
            # Use MongoClient from mongomock
            try:
                import mongomock
            except ImportError:
                raise RuntimeError('You need mongomock installed '
                                   'to mock MongoEngine.')
            connection_class = mongomock.MongoClient
        else:
            # Use MongoClient from pymongo
            connection_class = MongoClient

        if 'replicaSet' in conn_settings:
            # Discard port since it can't be used on MongoReplicaSetClient
            conn_settings.pop('port', None)
            # Discard replicaSet if not base string
            if not isinstance(conn_settings['replicaSet'], basestring):
                conn_settings.pop('replicaSet', None)
            if not IS_PYMONGO_3:
                connection_class = MongoReplicaSetClient
                conn_settings['hosts_or_uri'] = conn_settings.pop('host', None)

        try:
            connection = None
            # check for shared connections
            connection_settings_iterator = (
                (db_alias, settings.copy()) for db_alias, settings in _connection_settings.iteritems())
            for db_alias, connection_settings in connection_settings_iterator:
                connection_settings.pop('name', None)
                connection_settings.pop('username', None)
                connection_settings.pop('password', None)
                connection_settings.pop('authentication_source', None)
                if conn_settings == connection_settings and _connections.get(db_alias, None):
                    connection = _connections[db_alias]
                    break

            _connections[alias] = connection if connection else connection_class(**conn_settings)
        except Exception, e:
            raise ConnectionError("Cannot connect to database %s :\n%s" % (alias, e))
    return _connections[alias]


def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
    global _dbs
    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 conn_settings['username'] and conn_settings['password']:
            db.authenticate(conn_settings['username'],
                            conn_settings['password'],
                            source=conn_settings['authentication_source'])
        _dbs[alias] = db
    return _dbs[alias]


def connect(db=None, alias=DEFAULT_CONNECTION_NAME, **kwargs):
    """Connect to the database specified by the 'db' argument.<commit_msg>Support for authentication mechanism #905 (#1333)<commit_after>from pymongo import MongoClient, ReadPreference, uri_parser
from mongoengine.python_support import (IS_PYMONGO_3, str_types)

__all__ = ['ConnectionError', 'connect', 'register_connection',
           'DEFAULT_CONNECTION_NAME']


DEFAULT_CONNECTION_NAME = 'default'

if IS_PYMONGO_3:
    READ_PREFERENCE = ReadPreference.PRIMARY
else:
    from pymongo import MongoReplicaSetClient
    READ_PREFERENCE = False


class ConnectionError(Exception):
    pass


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


def register_connection(alias, name=None, host=None, port=None,
                        read_preference=READ_PREFERENCE,
                        username=None, password=<PASSWORD>, authentication_source=None,
                        authentication_mechanism=None,
                        **kwargs):
    """Add a connection.

    :param alias: the name that will be used to refer to this connection
        throughout MongoEngine
    :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
       ** Added pymongo 2.1
    :param username: username to authenticate with
    :param password: <PASSWORD>
    :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 is_mock: explicitly use mongomock for this connection
        (can also be done by using `mongomock://` as db host prefix)
    :param kwargs: allow ad-hoc parameters to be passed into the pymongo driver

    .. versionchanged:: 0.10.6 - added mongomock support
    """
    global _connection_settings

    conn_settings = {
        'name': name or 'test',
        'host': host or 'localhost',
        'port': port or 27017,
        'read_preference': read_preference,
        'username': username,
        'password': password,
        'authentication_source': authentication_source,
        'authentication_mechanism': authentication_mechanism
    }

    # Handle uri style connections
    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_types):
        conn_host = [conn_host]

    resolved_hosts = []
    for entity in conn_host:
        # Handle uri style connections
        if entity.startswith('mongomock://'):
            conn_settings['is_mock'] = True
            # `mongomock://` is not a valid url prefix and must be replaced by `mongodb://`
            resolved_hosts.append(entity.replace('mongomock://', 'mongodb://', 1))
        elif '://' in entity:
            uri_dict = uri_parser.parse_uri(entity)
            resolved_hosts.append(entity)
            conn_settings.update({
                'name': uri_dict.get('database') or name,
                'username': uri_dict.get('username'),
                'password': uri_dict.get('password'),
                'read_preference': read_preference,
            })
            uri_options = uri_dict['options']
            if 'replicaset' in uri_options:
                conn_settings['replicaSet'] = True
            if 'authsource' in uri_options:
                conn_settings['authentication_source'] = uri_options['authsource']
            if 'authmechanism' in uri_options:
                conn_settings['authentication_mechanism'] = uri_options['authmechanism']
        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)

    conn_settings.update(kwargs)
    _connection_settings[alias] = conn_settings


def disconnect(alias=DEFAULT_CONNECTION_NAME):
    global _connections
    global _dbs

    if alias in _connections:
        get_connection(alias=alias).close()
        del _connections[alias]
    if alias in _dbs:
        del _dbs[alias]


def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
    global _connections
    # Connect to the database if not already connected
    if reconnect:
        disconnect(alias)

    if alias not in _connections:
        if alias not in _connection_settings:
            msg = 'Connection with alias "%s" has not been defined' % alias
            if alias == DEFAULT_CONNECTION_NAME:
                msg = 'You have not defined a default connection'
            raise ConnectionError(msg)
        conn_settings = _connection_settings[alias].copy()

        conn_settings.pop('name', None)
        conn_settings.pop('username', None)
        conn_settings.pop('password', None)
        conn_settings.pop('authentication_source', None)
        conn_settings.pop('authentication_mechanism', None)

        is_mock = conn_settings.pop('is_mock', None)
        if is_mock:
            # Use MongoClient from mongomock
            try:
                import mongomock
            except ImportError:
                raise RuntimeError('You need mongomock installed '
                                   'to mock MongoEngine.')
            connection_class = mongomock.MongoClient
        else:
            # Use MongoClient from pymongo
            connection_class = MongoClient

        if 'replicaSet' in conn_settings:
            # Discard port since it can't be used on MongoReplicaSetClient
            conn_settings.pop('port', None)
            # Discard replicaSet if not base string
            if not isinstance(conn_settings['replicaSet'], basestring):
                conn_settings.pop('replicaSet', None)
            if not IS_PYMONGO_3:
                connection_class = MongoReplicaSetClient
                conn_settings['hosts_or_uri'] = conn_settings.pop('host', None)

        try:
            connection = None
            # check for shared connections
            connection_settings_iterator = (
                (db_alias, settings.copy()) for db_alias, settings in _connection_settings.iteritems())
            for db_alias, connection_settings in connection_settings_iterator:
                connection_settings.pop('name', None)
                connection_settings.pop('username', None)
                connection_settings.pop('password', None)
                connection_settings.pop('authentication_source', None)
                connection_settings.pop('authentication_mechanism', None)
                if conn_settings == connection_settings and _connections.get(db_alias, None):
                    connection = _connections[db_alias]
                    break

            _connections[alias] = connection if connection else connection_class(**conn_settings)
        except Exception, e:
            raise ConnectionError("Cannot connect to database %s :\n%s" % (alias, e))
    return _connections[alias]


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

    if alias not in _dbs:
        conn = get_connection(alias)
        conn_settings = _connection_settings[alias]
        db = conn[conn_settings['name']]
        auth_kwargs = {'source': conn_settings['authentication_source']}
        if conn_settings['authentication_mechanism'] is not None:
            auth_kwargs['mechanism'] = conn_settings['authentication_mechanism']
        # Authenticate if necessary
        if conn_settings['username'] and (conn_settings['password'] or
                                          conn_settings['authentication_mechanism'] == 'MONGODB-X509'):
            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.
{'fraction_non_alphanumeric': 0.08384136348738119, 'fraction_numerical': 0.0028187479514913145, 'mean_word_length': 2.5503839888294157, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7117626', 'n_tokens_mistral': 4136, 'n_tokens_neox': 3952, 'n_words': 1114}