[#2685] document classes now cache a collection db alias that can be different per thread

Author: Jesse McLaughlin

mongoengine/base/metaclasses.py

@@ -166,7 +166,7 @@ def __new__(mcs, name, bases, attrs):
         ) = mcs._import_classes()
 
         if issubclass(new_class, Document):
-            new_class._collection = None
+            new_class._collections = {}
 
         # Add class to the _document_registry
         _document_registry[new_class._class_name] = new_class

mongoengine/connection.py

@@ -18,7 +18,8 @@
     "get_db",
     "register_connection",
     "set_local_db_alias",
-    "del_local_db_alias"
+    "del_local_db_alias",
+    "get_local_db_alias"
 ]
 
 
@@ -398,9 +399,14 @@ def del_local_db_alias(alias):
     _local.db_alias[alias].pop()
 
 
-def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
+def get_local_db_alias(alias):
     if alias in _local.db_alias and _local.db_alias[alias]:
         alias = _local.db_alias[alias][-1]
+    return alias
+
+
+def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
+    alias = get_local_db_alias(alias)
 
     if reconnect:
         disconnect(alias)

mongoengine/context_managers.py

@@ -81,18 +81,22 @@ def __init__(self, cls, db_alias):
     def __enter__(self):
         """Change the db_alias and clear the cached collection."""
         self.cls._meta["db_alias"] = self.db_alias
-        self.cls._collection = None
+        self.cls._set_collection(None)
         return self.cls
 
     def __exit__(self, t, value, traceback):
         """Reset the db_alias and collection."""
         self.cls._meta["db_alias"] = self.ori_db_alias
-        self.cls._collection = self.collection
+        self.cls._set_collection(self.collection)
 
 
 class switch_collection:
     """switch_collection alias context manager.
 
+    Warning ::
+
+        ###  This is NOT completely thread-safe  ###
+
     Example ::
 
         class Group(Document):
@@ -123,12 +127,12 @@ def _get_collection_name(cls):
             return self.collection_name
 
         self.cls._get_collection_name = _get_collection_name
-        self.cls._collection = None
+        self.cls._set_collection(None)
         return self.cls
 
     def __exit__(self, t, value, traceback):
         """Reset the collection."""
-        self.cls._collection = self.ori_collection
+        self.cls._set_collection(self.ori_collection)
         self.cls._get_collection_name = self.ori_get_collection_name
 
 

mongoengine/document.py

@@ -15,7 +15,7 @@
     get_document,
 )
 from mongoengine.common import _import_class
-from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
+from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db, get_local_db_alias
 from mongoengine.context_managers import (
     set_write_concern,
     switch_collection,
@@ -196,15 +196,23 @@ def __hash__(self):
 
         return hash(self.pk)
 
+    @classmethod
+    def _get_local_db_alias(cls):
+        return get_local_db_alias(cls._meta.get("db_alias", DEFAULT_CONNECTION_NAME))
+
     @classmethod
     def _get_db(cls):
         """Some Model using other db_alias"""
-        return get_db(cls._meta.get("db_alias", DEFAULT_CONNECTION_NAME))
+        return get_db(cls._get_local_db_alias())
 
     @classmethod
     def _disconnect(cls):
-        """Detach the Document class from the (cached) database collection"""
-        cls._collection = None
+        """Detach the Document class from all (cached) database collections"""
+        cls._collections = {}
+
+    @classmethod
+    def _set_collection(cls, collection):
+        cls._collections[cls._get_local_db_alias()] = collection
 
     @classmethod
     def _get_collection(cls):
@@ -216,14 +224,15 @@ def _get_collection(cls):
         2. Creates indexes defined in this document's :attr:`meta` dictionary.
            This happens only if `auto_create_index` is True.
         """
-        if not hasattr(cls, "_collection") or cls._collection is None:
+        local_db_alias = cls._get_local_db_alias()
+        if local_db_alias not in cls._collections:
             # 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()
+                cls._collections[local_db_alias] = cls._get_capped_collection()
             else:
                 db = cls._get_db()
                 collection_name = cls._get_collection_name()
-                cls._collection = db[collection_name]
+                cls._collections[local_db_alias] = db[collection_name]
 
             # Ensure indexes on the collection unless auto_create_index was
             # set to False.
@@ -232,7 +241,7 @@ def _get_collection(cls):
             if cls._meta.get("auto_create_index", True) and db.client.is_primary:
                 cls.ensure_indexes()
 
-        return cls._collection
+        return cls._collections[local_db_alias]
 
     @classmethod
     def _get_capped_collection(cls):
@@ -260,7 +269,7 @@ def _get_capped_collection(cls):
             if options.get("max") != max_documents or options.get("size") != max_size:
                 raise InvalidCollectionError(
                     'Cannot create collection "{}" as a capped '
-                    "collection as it already exists".format(cls._collection)
+                    "collection as it already exists".format(collection_name)
                 )
 
             return collection
@@ -837,7 +846,7 @@ def drop_collection(cls):
             raise OperationError(
                 "Document %s has no collection defined (is it abstract ?)" % cls
             )
-        cls._collection = None
+        cls._set_collection(None)
         db = cls._get_db()
         db.drop_collection(coll_name)
 

mongoengine/queryset/base.py

@@ -54,9 +54,9 @@ class BaseQuerySet:
     __dereference = False
     _auto_dereference = True
 
-    def __init__(self, document, collection):
+    def __init__(self, document, db_alias=None):
         self._document = document
-        self._collection_obj = collection
+        self._db_alias = db_alias
         self._mongo_query = None
         self._query_obj = Q()
         self._cls_query = {}
@@ -74,6 +74,8 @@ def __init__(self, document, collection):
         self._as_pymongo = False
         self._search_text = None
 
+        self.__init_using_collection()
+
         # If inheritance is allowed, only return instances and instances of
         # subclasses of the class being used
         if document._meta.get("allow_inheritance") is True:
@@ -100,6 +102,12 @@ def __init__(self, document, collection):
         # it anytime we change _limit. Inspired by how it is done in pymongo.Cursor
         self._empty = False
 
+    def __init_using_collection(self):
+        self._using_collection = None
+        if self._db_alias is not None:
+            with switch_db(self._document, self._db_alias) as cls:
+                self._using_collection = cls._get_collection()
+
     def __call__(self, q_obj=None, **query):
         """Filter the selected documents by calling the
         :class:`~mongoengine.queryset.QuerySet` with a query.
@@ -137,9 +145,6 @@ def __getstate__(self):
 
         obj_dict = self.__dict__.copy()
 
-        # don't picke collection, instead pickle collection params
-        obj_dict.pop("_collection_obj")
-
         # don't pickle cursor
         obj_dict["_cursor_obj"] = None
 
@@ -152,11 +157,11 @@ def __setstate__(self, obj_dict):
         See https://github.com/MongoEngine/mongoengine/issues/442
         """
 
-        obj_dict["_collection_obj"] = obj_dict["_document"]._get_collection()
-
         # update attributes
         self.__dict__.update(obj_dict)
 
+        self.__init_using_collection()
+
         # forse load cursor
         # self._cursor
 
@@ -494,7 +499,7 @@ def delete(self, write_concern=None, _from_doc_delete=False, cascade_refs=None):
             if rule == CASCADE:
                 cascade_refs = set() if cascade_refs is None else cascade_refs
                 # Handle recursive reference
-                if doc._collection == document_cls._collection:
+                if doc._collection == document_cls._get_collection():
                     for ref in queryset:
                         cascade_refs.add(ref.id)
                 refs = document_cls.objects(
@@ -777,14 +782,11 @@ def using(self, alias):
         :param alias: The database alias
         """
 
-        with switch_db(self._document, alias) as cls:
-            collection = cls._get_collection()
-
-        return self._clone_into(self.__class__(self._document, collection))
+        return self._clone_into(self.__class__(self._document, alias))
 
     def clone(self):
         """Create a copy of the current queryset."""
-        return self._clone_into(self.__class__(self._document, self._collection_obj))
+        return self._clone_into(self.__class__(self._document, self._db_alias))
 
     def _clone_into(self, new_qs):
         """Copy all of the relevant properties of this queryset to
@@ -1531,7 +1533,7 @@ def sum(self, field):
         if isinstance(field_instances[-1], ListField):
             pipeline.insert(1, {"$unwind": "$" + field})
 
-        result = tuple(self._document._get_collection().aggregate(pipeline))
+        result = tuple(self._collection.aggregate(pipeline))
 
         if result:
             return result[0]["total"]
@@ -1558,7 +1560,7 @@ def average(self, field):
         if isinstance(field_instances[-1], ListField):
             pipeline.insert(1, {"$unwind": "$" + field})
 
-        result = tuple(self._document._get_collection().aggregate(pipeline))
+        result = tuple(self._collection.aggregate(pipeline))
         if result:
             return result[0]["total"]
         return 0
@@ -1620,7 +1622,8 @@ def _collection(self):
         """Property that returns the collection object. This allows us to
         perform operations only if the collection is accessed.
         """
-        return self._collection_obj
+        return self._document._get_collection() \
+            if self._using_collection is None else self._using_collection
 
     @property
     def _cursor_args(self):

mongoengine/queryset/manager.py

@@ -35,7 +35,7 @@ def __get__(self, instance, owner):
 
         # owner is the document that contains the QuerySetManager
         queryset_class = owner._meta.get("queryset_class", self.default)
-        queryset = queryset_class(owner, owner._get_collection())
+        queryset = queryset_class(owner)
         if self.get_queryset:
             arg_count = self.get_queryset.__code__.co_argcount
             if arg_count == 1: