Repository: MongoEngine/mongoengine

Issue:
Flaky test in test suite makes CI randomly failing due to query_counter context manager
The test `def test_no_cached_queryset(self):` Is failing from time to time which is not handy as it makes the CI failing for unrelated reasons in PRs (e.g: #1766). https://github.com/MongoEngine/mongoengine/blob/42bbe6392729ee12ee4461452ac3614814933dcd/tests/queryset/queryset.py#L4709-L4719

The issue occurs because the query_counter() context was catching queries that aren't related to this particular test.

I managed to reproduce the issue after adding some debugging code and running travis in my fork (https://travis-ci.com/bagerard/mongoengine/jobs/143421029). The additional query being counted is actually a 'killcursors' query, probably being issued after the garbage collector closes another pymongo cursors.

The query_counter context manager is very handy for debugging but since it is not only counting the queries being issued within the context (other threads or processes can pollute it), it is not very robust.

I'll push a PR that makes the query_counter context ignoring 'killcursors' queries to fix this (and add a parameter `ignore_query` to the query_counter to let users customize the behavior if needed). Let me know if you have other ideas


Discussion:

Resolving patch:
diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py
--- a/mongoengine/context_managers.py
+++ b/mongoengine/context_managers.py
@@ -159,51 +159,69 @@ def __exit__(self, t, value, traceback):
 
 
 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()
-
-    def __enter__(self):
-        """On every with block we need to drop the profile collection."""
+        self.initial_profiling_level = None
+        self._ctx_query_counter = 0             # number of queries issued by the context
+
+        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 @@ def __repr__(self):
         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
 
 


Test patch:
diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py
--- a/tests/queryset/queryset.py
+++ b/tests/queryset/queryset.py
@@ -4714,18 +4714,27 @@ class Person(Document):
         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):
 
@@ -5053,7 +5062,7 @@ class Person(Document):
                                            {"$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
--- a/tests/test_context_managers.py
+++ b/tests/test_context_managers.py
@@ -209,18 +209,99 @@ class User(Document):
             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(q, 0)
+            list(collection.find())
+            self.assertEqual(q, 2)  # 1st select + 1 getmore
+
+    def test_query_counter_ignores_particular_queries(self):
+        connect('mongoenginetest')
+        db = get_db()
+
+        collection = db.query_counter
+        collection.insert_many([{'test': 'garbage %s' % i} for i in range(10)])
 
         with query_counter() as q:
-            self.assertEqual(0, 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)
 
-            for i in range(1, 51):
-                db.test.find({}).count()
+            cursor.close()              # issues a `killcursors` query that is ignored by the context
+            self.assertEqual(q, 1)
 
-            self.assertEqual(50, q)
+            _ = db.system.indexes.find_one()    # queries on db.system.indexes are ignored as well
+            self.assertEqual(q, 1)
 
 if __name__ == '__main__':
     unittest.main()