diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py index 67c83581..ee1f5e01 100644 --- a/mongoengine/context_managers.py +++ b/mongoengine/context_managers.py @@ -159,51 +159,69 @@ class no_sub_classes(object): class query_counter(object): - """Query_counter context manager to get the number of queries.""" + """Query_counter context manager to get the number of queries. + This works by updating the `profiling_level` of the database so that all queries get logged, + resetting the db.system.profile collection at the beginnig of the context and counting the new entries. + + This was designed for debugging purpose. In fact it is a global counter so queries issued by other threads/processes + can interfere with it + + Be aware that: + - Iterating over large amount of documents (>101) makes pymongo issue `getmore` queries to fetch the next batch of + documents (https://docs.mongodb.com/manual/tutorial/iterate-a-cursor/#cursor-batches) + - Some queries are ignored by default by the counter (killcursors, db.system.indexes) + """ def __init__(self): - """Construct the query_counter.""" - self.counter = 0 + """Construct the query_counter + """ self.db = get_db() + self.initial_profiling_level = None + self._ctx_query_counter = 0 # number of queries issued by the context - def __enter__(self): - """On every with block we need to drop the profile collection.""" + self._ignored_query = { + 'ns': + {'$ne': '%s.system.indexes' % self.db.name}, + 'op': + {'$ne': 'killcursors'} + } + + def _turn_on_profiling(self): + self.initial_profiling_level = self.db.profiling_level() self.db.set_profiling_level(0) self.db.system.profile.drop() self.db.set_profiling_level(2) + + def _resets_profiling(self): + self.db.set_profiling_level(self.initial_profiling_level) + + def __enter__(self): + self._turn_on_profiling() return self def __exit__(self, t, value, traceback): - """Reset the profiling level.""" - self.db.set_profiling_level(0) + self._resets_profiling() def __eq__(self, value): - """== Compare querycounter.""" counter = self._get_count() return value == counter def __ne__(self, value): - """!= Compare querycounter.""" return not self.__eq__(value) def __lt__(self, value): - """< Compare querycounter.""" return self._get_count() < value def __le__(self, value): - """<= Compare querycounter.""" return self._get_count() <= value def __gt__(self, value): - """> Compare querycounter.""" return self._get_count() > value def __ge__(self, value): - """>= Compare querycounter.""" return self._get_count() >= value def __int__(self): - """int representation.""" return self._get_count() def __repr__(self): @@ -211,10 +229,12 @@ class query_counter(object): return u"%s" % self._get_count() def _get_count(self): - """Get the number of queries.""" - ignore_query = {'ns': {'$ne': '%s.system.indexes' % self.db.name}} - count = self.db.system.profile.find(ignore_query).count() - self.counter - self.counter += 1 # Account for the query we just fired + """Get the number of queries by counting the current number of entries in db.system.profile + and substracting the queries issued by this context. In fact everytime this is called, 1 query is + issued so we need to balance that + """ + count = self.db.system.profile.find(self._ignored_query).count() - self._ctx_query_counter + self._ctx_query_counter += 1 # Account for the query we just issued to gather the information return count diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index c6aa3594..1da79864 100644 --- a/tests/queryset/queryset.py +++ b/tests/queryset/queryset.py @@ -4750,18 +4750,27 @@ class QuerySetTest(unittest.TestCase): for i in range(100): Person(name="No: %s" % i).save() - with query_counter() as q: - self.assertEqual(q, 0) - people = Person.objects.no_cache() + with query_counter() as q: + try: + self.assertEqual(q, 0) + people = Person.objects.no_cache() - [x for x in people] - self.assertEqual(q, 1) + [x for x in people] + self.assertEqual(q, 1) - list(people) - self.assertEqual(q, 2) + list(people) + self.assertEqual(q, 2) + + people.count() + self.assertEqual(q, 3) + except AssertionError as exc: + db = get_db() + msg = '' + for q in list(db.system.profile.find())[-50:]: + msg += str([q['ts'], q['ns'], q.get('query'), q['op']])+'\n' + msg += str(q) + raise AssertionError(str(exc) + '\n'+msg) - people.count() - self.assertEqual(q, 3) def test_cache_not_cloned(self): @@ -5089,7 +5098,7 @@ class QuerySetTest(unittest.TestCase): {"$ne": "%s.system.indexes" % q.db.name}})[0] self.assertFalse('$orderby' in op['query'], - 'BaseQuerySet must remove orderby from meta in boolen test') + 'BaseQuerySet must remove orderby from meta in boolean test') self.assertEqual(Person.objects.first().name, 'A') self.assertTrue(Person.objects._has_data(), diff --git a/tests/test_context_managers.py b/tests/test_context_managers.py index 8c96016c..2b265cbd 100644 --- a/tests/test_context_managers.py +++ b/tests/test_context_managers.py @@ -209,18 +209,99 @@ class ContextManagersTest(unittest.TestCase): with no_sub_classes(User): raise TypeError() + def test_query_counter_does_not_swallow_exception(self): + + with self.assertRaises(TypeError): + with query_counter() as q: + raise TypeError() + + def test_query_counter_temporarily_modifies_profiling_level(self): + connect('mongoenginetest') + db = get_db() + + initial_profiling_level = db.profiling_level() + + try: + NEW_LEVEL = 1 + db.set_profiling_level(NEW_LEVEL) + self.assertEqual(db.profiling_level(), NEW_LEVEL) + with query_counter() as q: + self.assertEqual(db.profiling_level(), 2) + self.assertEqual(db.profiling_level(), NEW_LEVEL) + except Exception: + db.set_profiling_level(initial_profiling_level) # Ensures it gets reseted no matter the outcome of the test + raise + def test_query_counter(self): connect('mongoenginetest') db = get_db() - db.test.find({}) + + collection = db.query_counter + collection.drop() + + def issue_1_count_query(): + collection.find({}).count() + + def issue_1_insert_query(): + collection.insert_one({'test': 'garbage'}) + + def issue_1_find_query(): + collection.find_one() + + counter = 0 + with query_counter() as q: + self.assertEqual(q, counter) + self.assertEqual(q, counter) # Ensures previous count query did not get counted + + for _ in range(10): + issue_1_insert_query() + counter += 1 + self.assertEqual(q, counter) + + for _ in range(4): + issue_1_find_query() + counter += 1 + self.assertEqual(q, counter) + + for _ in range(3): + issue_1_count_query() + counter += 1 + self.assertEqual(q, counter) + + def test_query_counter_counts_getmore_queries(self): + connect('mongoenginetest') + db = get_db() + + collection = db.query_counter + collection.drop() + + many_docs = [{'test': 'garbage %s' % i} for i in range(150)] + collection.insert_many(many_docs) # first batch of documents contains 101 documents with query_counter() as q: - self.assertEqual(0, q) + self.assertEqual(q, 0) + list(collection.find()) + self.assertEqual(q, 2) # 1st select + 1 getmore - for i in range(1, 51): - db.test.find({}).count() + def test_query_counter_ignores_particular_queries(self): + connect('mongoenginetest') + db = get_db() - self.assertEqual(50, q) + collection = db.query_counter + collection.insert_many([{'test': 'garbage %s' % i} for i in range(10)]) + + with query_counter() as q: + self.assertEqual(q, 0) + cursor = collection.find() + self.assertEqual(q, 0) # cursor wasn't opened yet + _ = next(cursor) # opens the cursor and fires the find query + self.assertEqual(q, 1) + + cursor.close() # issues a `killcursors` query that is ignored by the context + self.assertEqual(q, 1) + + _ = db.system.indexes.find_one() # queries on db.system.indexes are ignored as well + self.assertEqual(q, 1) if __name__ == '__main__': unittest.main()