Make overriding settings easier and allow lazy connections.

flask_mongoengine/__init__.py

@@ -4,6 +4,8 @@
 from flask import abort, current_app
 
 import mongoengine
+import mongoengine.connection
+from mongoengine.connection import DEFAULT_CONNECTION_NAME
 
 from mongoengine.queryset import MultipleObjectsReturned, DoesNotExist, QuerySet
 from mongoengine.base import ValidationError
@@ -15,14 +17,19 @@
 from .json import overide_json_encoder
 
 
+# Dictionary to hold connection settings until `get_db` is executed. This
+# allows creation of lazy connection objects.
+pending_lazy_connections = {}
+
+
 def _include_mongoengine(obj):
     for module in mongoengine, mongoengine.fields:
         for key in module.__all__:
             if not hasattr(obj, key):
                 setattr(obj, key, getattr(module, key))
 
 
-def _create_connection(conn_settings):
+def _create_connection(conn_settings, lazy=False):
 
     # Handle multiple connections recursively
     if isinstance(conn_settings, list):
@@ -41,7 +48,18 @@ def _create_connection(conn_settings):
         uri_dict = uri_parser.parse_uri(conn_settings['host'])
         conn['db'] = uri_dict['database']
 
-    return mongoengine.connect(conn.pop('db', 'test'), **conn)
+    if lazy:
+        pending_lazy_connections[conn.get('alias')] = conn
+    else:
+        return mongoengine.connect(conn.pop('db', 'test'), **conn)
+
+
+def get_db(db_alias):
+    """Gets db after ensuring lazy connection is connected."""
+    if db_alias in pending_lazy_connections:
+        conn = pending_lazy_connections.get(db_alias)
+        mongoengine.connect(conn.pop('db', 'test'), **conn)
+    return mongoengine.connection.get_db(db_alias)
 
 
 class MongoEngine(object):
@@ -71,35 +89,38 @@ def init_app(self, app, config=None):
             # potentially new configuration would not be loaded.
             raise Exception('Extension already initialized')
 
-        if config:
-            # If passed an explicit config then we must make sure to ignore
-            # anything set in the application config.
-            connection = _create_connection(config)
+        if not config:
+            # If not passed a config then we read the connection settings
+            # from the app config.
+            config = app.config
+
+        if 'MONGODB_SETTINGS' in config:
+            # Connection settings provided as a dictionary.
+            settings = config['MONGODB_SETTINGS']
         else:
-            # Set default config
-            config = {}
-            config.setdefault('db', app.config.get('MONGODB_DB', None))
-            config.setdefault('host', app.config.get('MONGODB_HOST', None))
-            config.setdefault('port', app.config.get('MONGODB_PORT', None))
-            config.setdefault('username',
-                                app.config.get('MONGODB_USERNAME', None))
-            config.setdefault('password',
-                                app.config.get('MONGODB_PASSWORD', None))
-
-            # Before using default config we check for MONGODB_SETTINGS
-            if 'MONGODB_SETTINGS' in app.config:
-                connection = _create_connection(app.config['MONGODB_SETTINGS'])
-            else:
-                connection = _create_connection(config)
+            # Connection settings provided in standard format.
+            settings = {'alias': config.get('MONGODB_ALIAS', None),
+                        'db': config.get('MONGODB_DB', None),
+                        'host': config.get('MONGODB_HOST', None),
+                        'password': config.get('MONGODB_PASSWORD', None),
+                        'port': config.get('MONGODB_PORT', None),
+                        'username': config.get('MONGODB_USERNAME', None)}
+
+        # If the alias isn't set then use the default connection name,
+        if not settings.get('alias'):
+            settings['alias'] = DEFAULT_CONNECTION_NAME
 
         # Store objects in application instance so that multiple apps do
         # not end up accessing the same objects.
         app.extensions['mongoengine'][self] = {'app': app,
-                                               'conn': connection}
+                                               'settings': settings}
 
-    @property
-    def connection(self):
-        return current_app.extensions['mongoengine'][self]['conn']
+        # Creating a lazy connection requires that your documents inherit from
+        # the Document class defined in this extension. Otherwise mongoengine
+        # will not have an initialized connection on first attempt to access
+        # data.
+        lazy = bool(config.get('MONGODB_LAZY_CONNECTION', False))
+        _create_connection(settings, lazy=lazy)
 
 
 class BaseQuerySet(QuerySet):
@@ -145,6 +166,11 @@ def paginate_field(self, field_name, page, per_page, total=None):
         return ListFieldPagination(self.__class__.objects, self.pk, field_name,
                                    page, per_page, total=total)
 
+    @classmethod
+    def _get_db(cls):
+        """Override so we can establish connections inside flask-mongoengine"""
+        return get_db(cls._meta.get("db_alias", DEFAULT_CONNECTION_NAME))
+
 
 class DynamicDocument(mongoengine.DynamicDocument):
     """Abstract Dynamic document with extra helpers in the queryset class"""

tests/test_basic_app.py

@@ -13,15 +13,14 @@ class BasicAppTestCase(FlaskMongoEngineTestCase):
     def setUp(self):
         super(BasicAppTestCase, self).setUp()
         db = MongoEngine()
+        db.init_app(self.app)
 
         class Todo(db.Document):
             title = db.StringField(max_length=60)
             text = db.StringField()
             done = db.BooleanField(default=False)
             pub_date = db.DateTimeField(default=datetime.datetime.now)
 
-        db.init_app(self.app)
-
         Todo.drop_collection()
         self.Todo = Todo
 

tests/test_json_app.py

@@ -24,15 +24,14 @@ def setUp(self):
         self.app.config['MONGODB_DB'] = 'testing'
         self.app.config['TESTING'] = True
         db = MongoEngine()
+        db.init_app(self.app)
 
         class Todo(db.Document):
             title = db.StringField(max_length=60)
             text = db.StringField()
             done = db.BooleanField(default=False)
             pub_date = db.DateTimeField(default=datetime.datetime.now)
 
-        db.init_app(self.app)
-
         Todo.drop_collection()
         self.Todo = Todo
 

tests/test_session.py

@@ -16,7 +16,8 @@ def setUp(self):
         self.db_name = 'testing'
         self.app.config['MONGODB_DB'] = self.db_name
         self.app.config['TESTING'] = True
-        db = MongoEngine(self.app)
+        db = MongoEngine()
+        db.init_app(self.app)
         self.app.session_interface = MongoEngineSessionInterface(db)
 
         @self.app.route('/')