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)