Merge pull request #2038 from bagerard/disconnect

Fix connect/disconnect functions

docs/changelog.rst

@@ -4,10 +4,17 @@ Changelog
 
 Development
 ===========
+- expose `mongoengine.connection.disconnect` and `mongoengine.connection.disconnect_all`
+- Fix disconnect function #566 #1599 #605 #607 #1213 #565
+- Improve connect/disconnect documentations
+- POTENTIAL BREAKING CHANGES: (associated with connect/disconnect fixes)
+    - calling `connect` 2 times with the same alias and different parameter will raise an error (should call disconnect first)
+    - disconnect now clears `mongoengine.connection._connection_settings`
+    - disconnect now clears the cached attribute `Document._collection`
 - POTENTIAL BREAKING CHANGE: Aggregate gives wrong results when used with a queryset having limit and skip #2029
 - mongoengine now requires pymongo>=3.5 #2017
 - Generate Unique Indices for SortedListField and EmbeddedDocumentListFields #2020
-- connect() fails immediately when db name contains invalid characters (e. g. when user mistakenly puts 'mongodb://127.0.0.1:27017' as db name, happened in #1718) or is if db name is of an invalid type
+- connect() fails immediately when db name contains invalid characters #2031 #1718
 - (Fill this out as you fix issues and develop your features).
 
 Changes in 0.17.0

docs/guide/connecting.rst

@@ -4,9 +4,11 @@
 Connecting to MongoDB
 =====================
 
-To connect to a running instance of :program:`mongod`, use the
-:func:`~mongoengine.connect` function. The first argument is the name of the
-database to connect to::
+Connections in MongoEngine are registered globally and are identified with aliases.
+If no `alias` is provided during the connection, it will use "default" as alias.
+
+To connect to a running instance of :program:`mongod`, use the :func:`~mongoengine.connect`
+function. The first argument is the name of the database to connect to::
 
     from mongoengine import connect
     connect('project1')
@@ -42,6 +44,9 @@ the :attr:`host` to
     will establish connection to ``production`` database using
     ``admin`` username and ``qwerty`` password.
 
+.. note:: Calling :func:`~mongoengine.connect` without argument will establish
+    a connection to the "test" database by default
+
 Replica Sets
 ============
 
@@ -71,6 +76,8 @@ is used.
 In the background this uses :func:`~mongoengine.register_connection` to
 store the data and you can register all aliases up front if required.
 
+Documents defined in different database
+---------------------------------------
 Individual documents can also support multiple databases by providing a
 `db_alias` in their meta data. This allows :class:`~pymongo.dbref.DBRef`
 objects to point across databases and collections. Below is an example schema,
@@ -93,6 +100,33 @@ using 3 different databases to store data::
             meta = {'db_alias': 'users-books-db'}
 
 
+Disconnecting an existing connection
+------------------------------------
+The function :func:`~mongoengine.disconnect` can be used to
+disconnect a particular connection. This can be used to change a
+connection globally::
+
+        from mongoengine import connect, disconnect
+        connect('a_db', alias='db1')
+
+        class User(Document):
+            name = StringField()
+            meta = {'db_alias': 'db1'}
+
+        disconnect(alias='db1')
+
+        connect('another_db', alias='db1')
+
+.. note:: Calling :func:`~mongoengine.disconnect` without argument
+    will disconnect the "default" connection
+
+.. note:: Since connections gets registered globally, it is important
+    to use the `disconnect` function from MongoEngine and not the
+    `disconnect()` method of an existing connection (pymongo.MongoClient)
+
+.. note:: :class:`~mongoengine.Document` are caching the pymongo collection.
+    using `disconnect` ensures that it gets cleaned as well
+
 Context Managers
 ================
 Sometimes you may want to switch the database or collection to query against.
@@ -119,7 +153,7 @@ access to the same User document across databases::
 
 Switch Collection
 -----------------
-The :class:`~mongoengine.context_managers.switch_collection` context manager
+The :func:`~mongoengine.context_managers.switch_collection` context manager
 allows you to change the collection for a given class allowing quick and easy
 access to the same Group document across collection::
 

mongoengine/base/common.py

@@ -13,7 +13,7 @@
 
 
 def get_document(name):
-    """Get a document class by name."""
+    """Get a registered Document class by name."""
     doc = _document_registry.get(name, None)
     if not doc:
         # Possible old style name
@@ -30,3 +30,12 @@ def get_document(name):
             been imported?
         """.strip() % name)
     return doc
+
+
+def _get_documents_by_db(connection_alias, default_connection_alias):
+    """Get all registered Documents class attached to a given database"""
+    def get_doc_alias(doc_cls):
+        return doc_cls._meta.get('db_alias', default_connection_alias)
+
+    return [doc_cls for doc_cls in _document_registry.values()
+            if get_doc_alias(doc_cls) == connection_alias]

mongoengine/connection.py

@@ -4,11 +4,15 @@
 
 from mongoengine.pymongo_support import IS_PYMONGO_3
 
-__all__ = ['MongoEngineConnectionError', 'connect', 'register_connection',
-           'DEFAULT_CONNECTION_NAME', 'get_db']
+__all__ = ['MongoEngineConnectionError', 'connect', 'disconnect', 'disconnect_all',
+           'register_connection', 'DEFAULT_CONNECTION_NAME', 'DEFAULT_DATABASE_NAME',
+           'get_db', 'get_connection']
 
 
 DEFAULT_CONNECTION_NAME = 'default'
+DEFAULT_DATABASE_NAME = 'test'
+DEFAULT_HOST = 'localhost'
+DEFAULT_PORT = 27017
 
 if IS_PYMONGO_3:
     READ_PREFERENCE = ReadPreference.PRIMARY
@@ -39,40 +43,39 @@ def check_db_name(name):
         _check_name(name)
 
 
-def register_connection(alias, db=None, name=None, host=None, port=None,
-                        read_preference=READ_PREFERENCE,
-                        username=None, password=None,
-                        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 db: the name of the database to use, for compatibility with connect
-    :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
+def _get_connection_settings(
+        db=None, name=None, host=None, port=None,
+        read_preference=READ_PREFERENCE,
+        username=None, password=None,
+        authentication_source=None,
+        authentication_mechanism=None,
+        **kwargs):
+    """Get the connection settings as a dict
+
+    : param db: the name of the database to use, for compatibility with connect
+    : 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 to authenticate with
-    :param authentication_source: database to authenticate against
-    :param authentication_mechanism: database authentication mechanisms.
+    : param username: username to authenticate with
+    : param password: password to authenticate with
+    : 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: ad-hoc parameters to be passed into the pymongo driver,
+    : param is_mock: explicitly use mongomock for this connection
+        (can also be done by using `mongomock: // ` as db host prefix)
+    : param kwargs: ad-hoc parameters to be passed into the pymongo driver,
         for example maxpoolsize, tz_aware, etc. See the documentation
         for pymongo's `MongoClient` for a full list.
 
     .. versionchanged:: 0.10.6 - added mongomock support
     """
     conn_settings = {
-        'name': name or db or 'test',
-        'host': host or 'localhost',
-        'port': port or 27017,
+        'name': name or db or DEFAULT_DATABASE_NAME,
+        'host': host or DEFAULT_HOST,
+        'port': port or DEFAULT_PORT,
         'read_preference': read_preference,
         'username': username,
         'password': password,
@@ -137,17 +140,75 @@ def register_connection(alias, db=None, name=None, host=None, port=None,
     kwargs.pop('is_slave', None)
 
     conn_settings.update(kwargs)
+    return conn_settings
+
+
+def register_connection(alias, db=None, name=None, host=None, port=None,
+                        read_preference=READ_PREFERENCE,
+                        username=None, password=None,
+                        authentication_source=None,
+                        authentication_mechanism=None,
+                        **kwargs):
+    """Register the connection settings.
+
+    : 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 db: the name of the database to use, for compatibility with connect
+    : 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 to authenticate with
+    : 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: ad-hoc parameters to be passed into the pymongo driver,
+        for example maxpoolsize, tz_aware, etc. See the documentation
+        for pymongo's `MongoClient` for a full list.
+
+    .. versionchanged:: 0.10.6 - added mongomock support
+    """
+    conn_settings = _get_connection_settings(
+        db=db, name=name, host=host, port=port,
+        read_preference=read_preference,
+        username=username, password=password,
+        authentication_source=authentication_source,
+        authentication_mechanism=authentication_mechanism,
+        **kwargs)
     _connection_settings[alias] = conn_settings
 
 
 def disconnect(alias=DEFAULT_CONNECTION_NAME):
     """Close the connection with a given alias."""
+    from mongoengine.base.common import _get_documents_by_db
+    from mongoengine import Document
+
     if alias in _connections:
         get_connection(alias=alias).close()
         del _connections[alias]
+
     if alias in _dbs:
+        # Detach all cached collections in Documents
+        for doc_cls in _get_documents_by_db(alias, DEFAULT_CONNECTION_NAME):
+            if issubclass(doc_cls, Document):     # Skip EmbeddedDocument
+                doc_cls._disconnect()
+
         del _dbs[alias]
 
+    if alias in _connection_settings:
+        del _connection_settings[alias]
+
+
+def disconnect_all():
+    """Close all registered database."""
+    for alias in list(_connections.keys()):
+        disconnect(alias)
+
 
 def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
     """Return a connection with a given alias."""
@@ -270,14 +331,24 @@ def connect(db=None, alias=DEFAULT_CONNECTION_NAME, **kwargs):
     provide username and password arguments as well.
 
     Multiple databases are supported by using aliases. Provide a separate
-    `alias` to connect to a different instance of :program:`mongod`.
+    `alias` to connect to a different instance of: program: `mongod`.
+
+    In order to replace a connection identified by a given alias, you'll
+    need to call ``disconnect`` first
 
     See the docstring for `register_connection` for more details about all
     supported kwargs.
 
     .. versionchanged:: 0.6 - added multiple database support.
     """
-    if alias not in _connections:
+    if alias in _connections:
+        prev_conn_setting = _connection_settings[alias]
+        new_conn_settings = _get_connection_settings(db, **kwargs)
+
+        if new_conn_settings != prev_conn_setting:
+            raise MongoEngineConnectionError(
+                'A different connection with alias `%s` was already registered. Use disconnect() first' % alias)
+    else:
         register_connection(alias, db, **kwargs)
 
     return get_connection(alias)

mongoengine/document.py

@@ -187,11 +187,17 @@ def _get_db(cls):
         """Some Model using other db_alias"""
         return get_db(cls._meta.get('db_alias', DEFAULT_CONNECTION_NAME))
 
+    @classmethod
+    def _disconnect(cls):
+        """Detach the Document class from the (cached) database collection"""
+        cls._collection = None
+
     @classmethod
     def _get_collection(cls):
-        """Return a PyMongo collection for the document."""
+        """Return the corresponding PyMongo collection of this document.
+        Upon the first call, it will ensure that indexes gets created. The returned collection then gets cached
+        """
         if not hasattr(cls, '_collection') or cls._collection is None:
-
             # Get the collection, either capped or regular.
             if cls._meta.get('max_size') or cls._meta.get('max_documents'):
                 cls._collection = cls._get_capped_collection()
@@ -789,13 +795,13 @@ def drop_collection(cls):
         .. versionchanged:: 0.10.7
             :class:`OperationError` exception raised if no collection available
         """
-        col_name = cls._get_collection_name()
-        if not col_name:
+        coll_name = cls._get_collection_name()
+        if not coll_name:
             raise OperationError('Document %s has no collection defined '
                                  '(is it abstract ?)' % cls)
         cls._collection = None
         db = cls._get_db()
-        db.drop_collection(col_name)
+        db.drop_collection(coll_name)
 
     @classmethod
     def create_index(cls, keys, background=False, **kwargs):