From 83e3c5c7d8545d5ced30ce679fef800aeecbd60c Mon Sep 17 00:00:00 2001 From: Ross Lawley Date: Fri, 9 Dec 2011 08:26:39 -0800 Subject: [PATCH] Updated connection for pymongo 2.1 support closes #378 --- mongoengine/connection.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/mongoengine/connection.py b/mongoengine/connection.py index 1c0504a3..1284e69a 100644 --- a/mongoengine/connection.py +++ b/mongoengine/connection.py @@ -1,4 +1,4 @@ -from pymongo import Connection +from pymongo import Connection, version_tuple __all__ = ['ConnectionError', 'connect', 'register_connection', @@ -18,8 +18,8 @@ _dbs = {} def register_connection(alias, name, host='localhost', port=27017, - is_slave=False, slaves=None, username=None, - password=None): + is_slave=False, read_preference=False, slaves=None, + username=None, password=None): """Add a connection. :param alias: the name that will be used to refer to this connection @@ -27,11 +27,13 @@ def register_connection(alias, name, host='localhost', port=27017, :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 is_slave: whether the connection can act as a slave + :param is_slave: whether the connection can act as a slave ** Depreciated pymongo 2.0.1+ + :param read_preference: The read preference for the collection ** Added pymongo 2.1 :param slaves: a list of aliases of slave connections; each of these must be a registered connection that has :attr:`is_slave` set to ``True`` :param username: username to authenticate with :param password: password to authenticate with + """ global _connection_settings _connection_settings[alias] = { @@ -42,6 +44,7 @@ def register_connection(alias, name, host='localhost', port=27017, 'slaves': slaves or [], 'username': username, 'password': password, + 'read_preference': read_preference } @@ -70,11 +73,19 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False): raise ConnectionError(msg) conn_settings = _connection_settings[alias].copy() - # Get all the slave connections - slaves = [] - for slave_alias in conn_settings['slaves']: - slaves.append(get_connection(slave_alias)) - conn_settings['slaves'] = slaves + if version_tuple[0] >= 2 and version_tuple [1] > 0: + conn_settings.pop('name') + conn_settings.pop('slaves') + conn_settings.pop('is_slave') + conn_settings.pop('username') + conn_settings.pop('password') + else: + # Get all the slave connections + slaves = [] + for slave_alias in conn_settings['slaves']: + slaves.append(get_connection(slave_alias)) + conn_settings['slaves'] = slaves + conn_settings.pop('read_preference') try: _connections[alias] = Connection(**conn_settings)