Fix connect/disconnect functions

- expose disconnect
- disconnect cleans _connection_settings
- disconnect cleans cached collection in Document._collection
- re-connecting with the same alias raise an error (must call disconnect in between)
This commit is contained in:
Bastien Gérard
2019-04-15 21:24:07 +02:00
parent 9bb3dfd639
commit d1467c2f73
6 changed files with 312 additions and 26 deletions

View File

@@ -188,10 +188,16 @@ class Document(six.with_metaclass(TopLevelDocumentMetaclass, BaseDocument)):
return get_db(cls._meta.get('db_alias', DEFAULT_CONNECTION_NAME))
@classmethod
def _get_collection(cls):
"""Return a PyMongo collection for the document."""
if not hasattr(cls, '_collection') or cls._collection is None:
def _disconnect(cls):
"""Detach the Document class from the (cached) database collection"""
cls._collection = None
@classmethod
def _get_collection(cls):
"""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()