Skip to content

Commit 8a34683

Browse files
committed
Resolved issue with Multiple DB settings.
1 parent bdc96c7 commit 8a34683

File tree

2 files changed

+52
-4
lines changed

2 files changed

+52
-4
lines changed

flask_mongoengine/connection.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,11 @@ def _validate_settings(is_test, temp_db, preserved, conn_host):
7373
'only when `TESTING` is set to true.'
7474
raise InvalidSettingsError(msg)
7575

76+
def __get_app_config(key):
77+
return (_app_instance.get(key, False)
78+
if isinstance(_app_instance, dict)
79+
else _app_instance.config.get(key, False))
80+
7681
def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
7782
global _connections
7883
set_global_attributes()
@@ -98,9 +103,9 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
98103
conn_settings.pop('password', None)
99104
conn_settings.pop('authentication_source', None)
100105

101-
is_test = _app_instance.config.get('TESTING', False)
102-
temp_db = _app_instance.config.get('TEMP_DB', False)
103-
preserved = _app_instance.config.get('PRESERVE_TEMP_DB', False)
106+
is_test = __get_app_config('TESTING')
107+
temp_db = __get_app_config('TEMP_DB')
108+
preserved = __get_app_config('PRESERVE_TEMP_DB')
104109

105110
# Validation
106111
_validate_settings(is_test, temp_db, preserved, conn_host)
@@ -351,7 +356,7 @@ def create_connection(config, app):
351356
@param app: instance of flask.Flask
352357
"""
353358
global _connection_settings, _app_instance
354-
_app_instance = app
359+
_app_instance = app if app else config
355360

356361
if config is None or not isinstance(config, dict):
357362
raise Exception("Invalid application configuration");
@@ -363,6 +368,7 @@ def create_connection(config, app):
363368
connections = {}
364369
for conn_setting in conn_settings:
365370
alias = conn_setting['alias']
371+
_connection_settings[alias] = conn_setting
366372
connections[alias] = get_connection(alias)
367373
return connections
368374
else:

tests/test_connection.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ def test_live_connection(self):
3939
'PASSWORD': None,
4040
'DB': 'test'
4141
}
42+
43+
self._do_persist(db)
44+
45+
def _do_persist(self, db):
4246
class Todo(db.Document):
4347
title = db.StringField(max_length=60)
4448
text = db.StringField()
@@ -57,6 +61,44 @@ class Todo(db.Document):
5761
f_to = Todo.objects().first()
5862
self.assertEqual(s_todo.title, f_to.title)
5963

64+
def test_multiple_connections(self):
65+
db = MongoEngine()
66+
self.app.config['TESTING'] = True
67+
self.app.config['MONGODB_SETTINGS'] = [
68+
{
69+
"ALIAS": "default",
70+
"DB": 'my_db1',
71+
"HOST": 'localhost',
72+
"PORT": 27017
73+
},
74+
{
75+
"ALIAS": "my_db2",
76+
"DB": 'my_db2',
77+
"HOST": 'localhost',
78+
"PORT": 27017
79+
},
80+
]
81+
class Todo(db.Document):
82+
title = db.StringField(max_length=60)
83+
text = db.StringField()
84+
done = db.BooleanField(default=False)
85+
meta = {"db_alias": "my_db2"}
86+
87+
db.init_app(self.app)
88+
Todo.drop_collection()
89+
90+
# Switch DB
91+
from mongoengine.context_managers import switch_db
92+
with switch_db(Todo, 'default') as Todo:
93+
todo = Todo()
94+
todo.text = "Sample"
95+
todo.title = "Testing"
96+
todo.done = True
97+
s_todo = todo.save()
98+
99+
f_to = Todo.objects().first()
100+
self.assertEqual(s_todo.title, f_to.title)
101+
60102
def test_mongodb_temp_instance(self):
61103
# String value used instead of boolean
62104
self.app.config['TESTING'] = True

0 commit comments

Comments
 (0)