pub/mongoengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix connecting to a list of hosts

Browse Source
This commit is contained in:
Jeff Tharp 2016-10-19 11:57:02 -07:00
parent 4ef5d1f0cd
commit f7ac8cea90
2 changed files with 63 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
mongoengine/connection.py
View File

@ -1,5 +1,5 @@
from pymongo import MongoClient, ReadPreference, uri_parser
from mongoengine.python_support import IS_PYMONGO_3
from mongoengine.python_support import (IS_PYMONGO_3, str_types)
__all__ = ['ConnectionError', 'connect', 'register_connection',
'DEFAULT_CONNECTION_NAME']
@ -56,25 +56,35 @@ def register_connection(alias, name=None, host=None, port=None,
'authentication_source': authentication_source
}
# Handle uri style connections
conn_host = conn_settings['host']
if conn_host.startswith('mongomock://'):
conn_settings['is_mock'] = True
# `mongomock://` is not a valid url prefix and must be replaced by `mongodb://`
conn_settings['host'] = conn_host.replace('mongomock://', 'mongodb://', 1)
elif '://' in conn_host:
uri_dict = uri_parser.parse_uri(conn_host)
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']
# 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)

34
tests/test_connection.py
View File

@ -88,6 +88,40 @@ class ConnectionTest(unittest.TestCase):
conn = get_connection('testdb7')
self.assertTrue(isinstance(conn, mongomock.MongoClient))
def test_connect_with_host_list(self):
"""Ensure that the connect() method works when host is a list
Uses mongomock to test w/o needing multiple mongod/mongos processes
"""
try:
import mongomock
except ImportError:
raise SkipTest('you need mongomock installed to run this testcase')
connect(host=['mongomock://localhost'])
conn = get_connection()
self.assertTrue(isinstance(conn, mongomock.MongoClient))
connect(host=['mongodb://localhost'], is_mock=True, alias='testdb2')
conn = get_connection('testdb2')
self.assertTrue(isinstance(conn, mongomock.MongoClient))
connect(host=['localhost'], is_mock=True, alias='testdb3')
conn = get_connection('testdb3')
self.assertTrue(isinstance(conn, mongomock.MongoClient))
connect(host=['mongomock://localhost:27017', 'mongomock://localhost:27018'], alias='testdb4')
conn = get_connection('testdb4')
self.assertTrue(isinstance(conn, mongomock.MongoClient))
connect(host=['mongodb://localhost:27017', 'mongodb://localhost:27018'], is_mock=True, alias='testdb5')
conn = get_connection('testdb5')
self.assertTrue(isinstance(conn, mongomock.MongoClient))
connect(host=['localhost:27017', 'localhost:27018'], is_mock=True, alias='testdb6')
conn = get_connection('testdb6')
self.assertTrue(isinstance(conn, mongomock.MongoClient))
def test_disconnect(self):
"""Ensure that the disconnect() method works properly
"""