Fix few things related to query_counter context manager

mongoengine/context_managers.py

@@ -159,62 +159,82 @@ 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):
         """repr query_counter as the number of queries."""
         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
 
 

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(),

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()