def _fetch_objects(self, doc_type=None): """Fetch all references and convert to their document objects""" object_map = {} for collection, dbrefs in self.reference_map.items(): # we use getattr instead of hasattr because hasattr swallows any exception under python2 # so it could hide nasty things without raising exceptions (cfr bug #1688)) ref_document_cls_exists = getattr(collection, "objects", None) is not None if ref_document_cls_exists: col_name = collection._get_collection_name() refs = [ dbref for dbref in dbrefs if (col_name, dbref) not in object_map ] references = collection.objects.in_bulk(refs) for key, doc in references.items(): object_map[(col_name, key)] = doc else: # Generic reference: use the refs data to convert to document if isinstance(doc_type, (ListField, DictField, MapField)): continue refs = [ dbref for dbref in dbrefs if (collection, dbref) not in object_map ] if doc_type: references = doc_type._get_db()[collection].find( {"_id": { "$in": refs }}, session=doc_type._get_local_session()) for ref in references: doc = doc_type._from_son(ref) object_map[(collection, doc.id)] = doc else: references = get_db()[collection].find( {"_id": { "$in": refs }}, session=get_local_session()) for ref in references: if "_cls" in ref: doc = get_document(ref["_cls"])._from_son(ref) elif doc_type is None: doc = get_document("".join( x.capitalize() for x in collection.split("_")))._from_son(ref) else: doc = doc_type._from_son(ref) object_map[(collection, doc.id)] = doc return object_map
def count_documents( collection, filter, skip=None, limit=None, hint=None, collation=None ): """Pymongo>3.7 deprecates count in favour of count_documents""" if limit == 0: return 0 # Pymongo raises an OperationFailure if called with limit=0 kwargs = {} if skip is not None: kwargs["skip"] = skip if limit is not None: kwargs["limit"] = limit if hint not in (-1, None): kwargs["hint"] = hint if collation is not None: kwargs["collation"] = collation # count_documents appeared in pymongo 3.7 if IS_PYMONGO_GTE_37: try: return collection.count_documents( filter=filter, session=get_local_session(), **kwargs ) except OperationFailure: # OperationFailure - accounts for some operators that used to work # with .count but are no longer working with count_documents (i.e $geoNear, $near, and $nearSphere) # fallback to deprecated Cursor.count # Keeping this should be reevaluated the day pymongo removes .count entirely pass cursor = collection.find(filter, session=get_local_session()) for option, option_value in kwargs.items(): cursor_method = getattr(cursor, option) cursor = cursor_method(option_value) with_limit_and_skip = "skip" in kwargs or "limit" in kwargs return cursor.count(with_limit_and_skip=with_limit_and_skip)
def test_set_get_local_session(self): session = {"db": "1"} sessions.set_local_session("test", session) self.assertEqual(session, sessions.get_local_session("test")) session2 = {"db": "2"} sessions.set_local_session("test2", session2) self.assertEqual(session2, sessions.get_local_session("test2")) self.assertNotEqual(sessions.get_local_session("test2"), sessions.get_local_session("test")) sessions.clear_local_session("test") self.assertIsNone(sessions.get_local_session("test")) sessions.clear_local_session("test2") self.assertIsNone(sessions.get_local_session("test2"))
def _get_local_session(cls): return get_local_session(cls._get_db_alias())
def new_session(i): db_alias = "test" session = {"db": i} sessions.set_local_session(db_alias, session) self.assertEqual(i, sessions.get_local_session(db_alias)["db"]) sessions.clear_local_session(db_alias)