def test_fetch_related(self):
with custom_query_counter() as q:
objs = list(self.B.objects.all())
fetch_related(objs, {
'ref': True
})
# make sure A objs are fetched
for obj in objs:
self.assertTrue(obj.ref.txt in ('a1', 'a2'))
# one query for B, one query for A
self.assertEqual(q, 2)
def test_batch_size_2(self):
"""
Ensure we batch requests properly, if a batch size is given.
"""
objs = list(self.B.objects.all())
with custom_query_counter() as q:
fetch_related(objs, {
'ref': True,
}, batch_size=3)
# make sure A objs are fetched
for obj in objs:
self.assertTrue(obj.ref.txt in ('a1', 'a2', 'a3'))
# All 3 objects are fetched in one query.
self.assertEqual(q, 1)
def test_fetch_related_multiple_objs(self):
with custom_query_counter() as q:
objs = list(self.B.objects.all()) + list(self.C.objects.all())
fetch_related(objs, {
'ref': True,
'ref_a': True
})
# make sure A objs are fetched
for obj in objs:
if isinstance(obj, self.B):
self.assertTrue(obj.ref.txt in ('a1', 'a2'))
else:
self.assertEqual(obj.ref_a.txt, 'a3')
# one query for B, one for C, one for A
self.assertEqual(q, 3)
def test_extra_filters(self):
"""
Ensure we apply extra filters by collection.
"""
objs = list(self.E.objects.all())
with custom_query_counter() as q:
fetch_related(objs, {
'refs_a': True,
'ref_b': True,
}, extra_filters={
self.A: {'shard_a': self.shard},
self.B: {'shard_b': self.shard},
})
ops = list(q.db.system.profile.find({'op': 'query'}))
assert len(ops) == 2
filters = {op['query']['find']: op['query']['filter'] for op in ops}
assert filters['a']['shard_a'] == self.shard.pk
assert filters['b']['shard_b'] == self.shard.pk
def test_fetch_related_subdict(self):
"""
Make sure fetching related references works with subfields and that
it uses caching properly.
"""
with custom_query_counter() as q:
objs = list(self.D.objects.all())
fetch_related(objs, {
'ref_a': True,
'ref_c': {
'ref_a': True
}
})
# make sure A objs are fetched
for obj in objs:
self.assertEqual(obj.ref_a.txt, 'a3')
self.assertEqual(obj.ref_c.ref_a.txt, 'a3')
# one query for D, one query for C, one query for A
self.assertEqual(q, 3)
def test_safe_reference_fields(self):
"""
Make sure SafeReferenceField and SafeReferenceListField don't fetch
the entire objects if we use a partial fetch_related on them.
"""
objs = list(self.E.objects.all())
with custom_query_counter() as q:
fetch_related(objs, {
'refs_a': ["id"],
'ref_b': ["id"]
})
# make sure the IDs match
self.assertEqual(
[a.pk for a in objs[0].refs_a],
[self.a1.pk, self.a2.pk, self.a3.pk]
)
self.assertEqual(objs[0].ref_b.pk, self.b1.pk)
# make sure other fields are empty
self.assertEqual(set([a.txt for a in objs[0].refs_a]), set([None]))
self.assertEqual(objs[0].ref_b.ref, None)
# make sure the queries to MongoDB only fetched the IDs
queries = list(q.db.system.profile.find({ 'op': 'query' }, { 'ns': 1, 'execStats': 1 }))
self.assertEqual(
set([ q['ns'].split('.')[1] for q in queries ]),
set([ 'a', 'b' ])
)
self.assertEqual(
set([ q['execStats']['stage'] for q in queries ]),
set([ 'PROJECTION' ]),
)
self.assertEqual(
set([ tuple(q['execStats']['transformBy'].keys()) for q in queries ]),
set([ ('_id',) ]),
)
def test_fetch_same_doc_class_multiple_times_with_cache_map(self):
"""
Make sure that the right documents are fetched when we reuse a cache
map for the same document type and the second fetch_related is a
partial fetch.
"""
self.b1.reload()
self.c1.reload()
cache_map = {}
objs = [self.b1, self.c1]
with custom_query_counter() as q:
fetch_related(objs, {
'ref': True
}, cache_map=cache_map)
fetch_related(objs, {
'ref_a': ['id']
}, cache_map=cache_map)
self.assertEqual(q, 2)
self.assertEqual(
[op['query']['_id']['$in'][0] for op in q.db.system.profile.find({ 'op': 'query' })],
[self.a1.pk, self.a3.pk]
)