We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mongoengine/context_managers.py`
Content:
```
1 from contextlib import contextmanager
2 from pymongo.write_concern import WriteConcern
3 from mongoengine.common import _import_class
4 from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
5 
6 
7 __all__ = ('switch_db', 'switch_collection', 'no_dereference',
8            'no_sub_classes', 'query_counter', 'set_write_concern')
9 
10 
11 class switch_db(object):
12     """switch_db alias context manager.
13 
14     Example ::
15 
16         # Register connections
17         register_connection('default', 'mongoenginetest')
18         register_connection('testdb-1', 'mongoenginetest2')
19 
20         class Group(Document):
21             name = StringField()
22 
23         Group(name='test').save()  # Saves in the default db
24 
25         with switch_db(Group, 'testdb-1') as Group:
26             Group(name='hello testdb!').save()  # Saves in testdb-1
27     """
28 
29     def __init__(self, cls, db_alias):
30         """Construct the switch_db context manager
31 
32         :param cls: the class to change the registered db
33         :param db_alias: the name of the specific database to use
34         """
35         self.cls = cls
36         self.collection = cls._get_collection()
37         self.db_alias = db_alias
38         self.ori_db_alias = cls._meta.get('db_alias', DEFAULT_CONNECTION_NAME)
39 
40     def __enter__(self):
41         """Change the db_alias and clear the cached collection."""
42         self.cls._meta['db_alias'] = self.db_alias
43         self.cls._collection = None
44         return self.cls
45 
46     def __exit__(self, t, value, traceback):
47         """Reset the db_alias and collection."""
48         self.cls._meta['db_alias'] = self.ori_db_alias
49         self.cls._collection = self.collection
50 
51 
52 class switch_collection(object):
53     """switch_collection alias context manager.
54 
55     Example ::
56 
57         class Group(Document):
58             name = StringField()
59 
60         Group(name='test').save()  # Saves in the default db
61 
62         with switch_collection(Group, 'group1') as Group:
63             Group(name='hello testdb!').save()  # Saves in group1 collection
64     """
65 
66     def __init__(self, cls, collection_name):
67         """Construct the switch_collection context manager.
68 
69         :param cls: the class to change the registered db
70         :param collection_name: the name of the collection to use
71         """
72         self.cls = cls
73         self.ori_collection = cls._get_collection()
74         self.ori_get_collection_name = cls._get_collection_name
75         self.collection_name = collection_name
76 
77     def __enter__(self):
78         """Change the _get_collection_name and clear the cached collection."""
79 
80         @classmethod
81         def _get_collection_name(cls):
82             return self.collection_name
83 
84         self.cls._get_collection_name = _get_collection_name
85         self.cls._collection = None
86         return self.cls
87 
88     def __exit__(self, t, value, traceback):
89         """Reset the collection."""
90         self.cls._collection = self.ori_collection
91         self.cls._get_collection_name = self.ori_get_collection_name
92 
93 
94 class no_dereference(object):
95     """no_dereference context manager.
96 
97     Turns off all dereferencing in Documents for the duration of the context
98     manager::
99 
100         with no_dereference(Group) as Group:
101             Group.objects.find()
102     """
103 
104     def __init__(self, cls):
105         """Construct the no_dereference context manager.
106 
107         :param cls: the class to turn dereferencing off on
108         """
109         self.cls = cls
110 
111         ReferenceField = _import_class('ReferenceField')
112         GenericReferenceField = _import_class('GenericReferenceField')
113         ComplexBaseField = _import_class('ComplexBaseField')
114 
115         self.deref_fields = [k for k, v in self.cls._fields.iteritems()
116                              if isinstance(v, (ReferenceField,
117                                                GenericReferenceField,
118                                                ComplexBaseField))]
119 
120     def __enter__(self):
121         """Change the objects default and _auto_dereference values."""
122         for field in self.deref_fields:
123             self.cls._fields[field]._auto_dereference = False
124         return self.cls
125 
126     def __exit__(self, t, value, traceback):
127         """Reset the default and _auto_dereference values."""
128         for field in self.deref_fields:
129             self.cls._fields[field]._auto_dereference = True
130         return self.cls
131 
132 
133 class no_sub_classes(object):
134     """no_sub_classes context manager.
135 
136     Only returns instances of this class and no sub (inherited) classes::
137 
138         with no_sub_classes(Group) as Group:
139             Group.objects.find()
140     """
141 
142     def __init__(self, cls):
143         """Construct the no_sub_classes context manager.
144 
145         :param cls: the class to turn querying sub classes on
146         """
147         self.cls = cls
148         self.cls_initial_subclasses = None
149 
150     def __enter__(self):
151         """Change the objects default and _auto_dereference values."""
152         self.cls_initial_subclasses = self.cls._subclasses
153         self.cls._subclasses = (self.cls._class_name,)
154         return self.cls
155 
156     def __exit__(self, t, value, traceback):
157         """Reset the default and _auto_dereference values."""
158         self.cls._subclasses = self.cls_initial_subclasses
159 
160 
161 class query_counter(object):
162     """Query_counter context manager to get the number of queries."""
163 
164     def __init__(self):
165         """Construct the query_counter."""
166         self.counter = 0
167         self.db = get_db()
168 
169     def __enter__(self):
170         """On every with block we need to drop the profile collection."""
171         self.db.set_profiling_level(0)
172         self.db.system.profile.drop()
173         self.db.set_profiling_level(2)
174         return self
175 
176     def __exit__(self, t, value, traceback):
177         """Reset the profiling level."""
178         self.db.set_profiling_level(0)
179 
180     def __eq__(self, value):
181         """== Compare querycounter."""
182         counter = self._get_count()
183         return value == counter
184 
185     def __ne__(self, value):
186         """!= Compare querycounter."""
187         return not self.__eq__(value)
188 
189     def __lt__(self, value):
190         """< Compare querycounter."""
191         return self._get_count() < value
192 
193     def __le__(self, value):
194         """<= Compare querycounter."""
195         return self._get_count() <= value
196 
197     def __gt__(self, value):
198         """> Compare querycounter."""
199         return self._get_count() > value
200 
201     def __ge__(self, value):
202         """>= Compare querycounter."""
203         return self._get_count() >= value
204 
205     def __int__(self):
206         """int representation."""
207         return self._get_count()
208 
209     def __repr__(self):
210         """repr query_counter as the number of queries."""
211         return u"%s" % self._get_count()
212 
213     def _get_count(self):
214         """Get the number of queries."""
215         ignore_query = {'ns': {'$ne': '%s.system.indexes' % self.db.name}}
216         count = self.db.system.profile.find(ignore_query).count() - self.counter
217         self.counter += 1       # Account for the query we just fired
218         return count
219 
220 
221 @contextmanager
222 def set_write_concern(collection, write_concerns):
223     combined_concerns = dict(collection.write_concern.document.items())
224     combined_concerns.update(write_concerns)
225     yield collection.with_options(write_concern=WriteConcern(**combined_concerns))
226
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
 
 
 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 @@
         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

{"golden_diff": "diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py\n--- a/mongoengine/context_managers.py\n+++ b/mongoengine/context_managers.py\n@@ -159,51 +159,69 @@\n \n \n class query_counter(object):\n-    \"\"\"Query_counter context manager to get the number of queries.\"\"\"\n+    \"\"\"Query_counter context manager to get the number of queries.\n+    This works by updating the `profiling_level` of the database so that all queries get logged,\n+    resetting the db.system.profile collection at the beginnig of the context and counting the new entries.\n+\n+    This was designed for debugging purpose. In fact it is a global counter so queries issued by other threads/processes\n+    can interfere with it\n+\n+    Be aware that:\n+    - Iterating over large amount of documents (>101) makes pymongo issue `getmore` queries to fetch the next batch of\n+        documents (https://docs.mongodb.com/manual/tutorial/iterate-a-cursor/#cursor-batches)\n+    - Some queries are ignored by default by the counter (killcursors, db.system.indexes)\n+    \"\"\"\n \n     def __init__(self):\n-        \"\"\"Construct the query_counter.\"\"\"\n-        self.counter = 0\n+        \"\"\"Construct the query_counter\n+        \"\"\"\n         self.db = get_db()\n-\n-    def __enter__(self):\n-        \"\"\"On every with block we need to drop the profile collection.\"\"\"\n+        self.initial_profiling_level = None\n+        self._ctx_query_counter = 0             # number of queries issued by the context\n+\n+        self._ignored_query = {\n+            'ns':\n+                {'$ne': '%s.system.indexes' % self.db.name},\n+            'op':\n+                {'$ne': 'killcursors'}\n+        }\n+\n+    def _turn_on_profiling(self):\n+        self.initial_profiling_level = self.db.profiling_level()\n         self.db.set_profiling_level(0)\n         self.db.system.profile.drop()\n         self.db.set_profiling_level(2)\n+\n+    def _resets_profiling(self):\n+        self.db.set_profiling_level(self.initial_profiling_level)\n+\n+    def __enter__(self):\n+        self._turn_on_profiling()\n         return self\n \n     def __exit__(self, t, value, traceback):\n-        \"\"\"Reset the profiling level.\"\"\"\n-        self.db.set_profiling_level(0)\n+        self._resets_profiling()\n \n     def __eq__(self, value):\n-        \"\"\"== Compare querycounter.\"\"\"\n         counter = self._get_count()\n         return value == counter\n \n     def __ne__(self, value):\n-        \"\"\"!= Compare querycounter.\"\"\"\n         return not self.__eq__(value)\n \n     def __lt__(self, value):\n-        \"\"\"< Compare querycounter.\"\"\"\n         return self._get_count() < value\n \n     def __le__(self, value):\n-        \"\"\"<= Compare querycounter.\"\"\"\n         return self._get_count() <= value\n \n     def __gt__(self, value):\n-        \"\"\"> Compare querycounter.\"\"\"\n         return self._get_count() > value\n \n     def __ge__(self, value):\n-        \"\"\">= Compare querycounter.\"\"\"\n         return self._get_count() >= value\n \n     def __int__(self):\n-        \"\"\"int representation.\"\"\"\n         return self._get_count()\n \n     def __repr__(self):\n@@ -211,10 +229,12 @@\n         return u\"%s\" % self._get_count()\n \n     def _get_count(self):\n-        \"\"\"Get the number of queries.\"\"\"\n-        ignore_query = {'ns': {'$ne': '%s.system.indexes' % self.db.name}}\n-        count = self.db.system.profile.find(ignore_query).count() - self.counter\n-        self.counter += 1       # Account for the query we just fired\n+        \"\"\"Get the number of queries by counting the current number of entries in db.system.profile\n+        and substracting the queries issued by this context. In fact everytime this is called, 1 query is\n+        issued so we need to balance that\n+        \"\"\"\n+        count = self.db.system.profile.find(self._ignored_query).count() - self._ctx_query_counter\n+        self._ctx_query_counter += 1    # Account for the query we just issued to gather the information\n         return count\n", "issue": "Flaky test in test suite makes CI randomly failing due to query_counter context manager\nThe 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\r\n\r\nThe issue occurs because the query_counter() context was catching queries that aren't related to this particular test.\r\n\r\nI 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.\r\n\r\nThe 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.\r\n\r\nI'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\n", "before_files": [{"content": "from contextlib import contextmanager\nfrom pymongo.write_concern import WriteConcern\nfrom mongoengine.common import _import_class\nfrom mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db\n\n\n__all__ = ('switch_db', 'switch_collection', 'no_dereference',\n           'no_sub_classes', 'query_counter', 'set_write_concern')\n\n\nclass switch_db(object):\n    \"\"\"switch_db alias context manager.\n\n    Example ::\n\n        # Register connections\n        register_connection('default', 'mongoenginetest')\n        register_connection('testdb-1', 'mongoenginetest2')\n\n        class Group(Document):\n            name = StringField()\n\n        Group(name='test').save()  # Saves in the default db\n\n        with switch_db(Group, 'testdb-1') as Group:\n            Group(name='hello testdb!').save()  # Saves in testdb-1\n    \"\"\"\n\n    def __init__(self, cls, db_alias):\n        \"\"\"Construct the switch_db context manager\n\n        :param cls: the class to change the registered db\n        :param db_alias: the name of the specific database to use\n        \"\"\"\n        self.cls = cls\n        self.collection = cls._get_collection()\n        self.db_alias = db_alias\n        self.ori_db_alias = cls._meta.get('db_alias', DEFAULT_CONNECTION_NAME)\n\n    def __enter__(self):\n        \"\"\"Change the db_alias and clear the cached collection.\"\"\"\n        self.cls._meta['db_alias'] = self.db_alias\n        self.cls._collection = None\n        return self.cls\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the db_alias and collection.\"\"\"\n        self.cls._meta['db_alias'] = self.ori_db_alias\n        self.cls._collection = self.collection\n\n\nclass switch_collection(object):\n    \"\"\"switch_collection alias context manager.\n\n    Example ::\n\n        class Group(Document):\n            name = StringField()\n\n        Group(name='test').save()  # Saves in the default db\n\n        with switch_collection(Group, 'group1') as Group:\n            Group(name='hello testdb!').save()  # Saves in group1 collection\n    \"\"\"\n\n    def __init__(self, cls, collection_name):\n        \"\"\"Construct the switch_collection context manager.\n\n        :param cls: the class to change the registered db\n        :param collection_name: the name of the collection to use\n        \"\"\"\n        self.cls = cls\n        self.ori_collection = cls._get_collection()\n        self.ori_get_collection_name = cls._get_collection_name\n        self.collection_name = collection_name\n\n    def __enter__(self):\n        \"\"\"Change the _get_collection_name and clear the cached collection.\"\"\"\n\n        @classmethod\n        def _get_collection_name(cls):\n            return self.collection_name\n\n        self.cls._get_collection_name = _get_collection_name\n        self.cls._collection = None\n        return self.cls\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the collection.\"\"\"\n        self.cls._collection = self.ori_collection\n        self.cls._get_collection_name = self.ori_get_collection_name\n\n\nclass no_dereference(object):\n    \"\"\"no_dereference context manager.\n\n    Turns off all dereferencing in Documents for the duration of the context\n    manager::\n\n        with no_dereference(Group) as Group:\n            Group.objects.find()\n    \"\"\"\n\n    def __init__(self, cls):\n        \"\"\"Construct the no_dereference context manager.\n\n        :param cls: the class to turn dereferencing off on\n        \"\"\"\n        self.cls = cls\n\n        ReferenceField = _import_class('ReferenceField')\n        GenericReferenceField = _import_class('GenericReferenceField')\n        ComplexBaseField = _import_class('ComplexBaseField')\n\n        self.deref_fields = [k for k, v in self.cls._fields.iteritems()\n                             if isinstance(v, (ReferenceField,\n                                               GenericReferenceField,\n                                               ComplexBaseField))]\n\n    def __enter__(self):\n        \"\"\"Change the objects default and _auto_dereference values.\"\"\"\n        for field in self.deref_fields:\n            self.cls._fields[field]._auto_dereference = False\n        return self.cls\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the default and _auto_dereference values.\"\"\"\n        for field in self.deref_fields:\n            self.cls._fields[field]._auto_dereference = True\n        return self.cls\n\n\nclass no_sub_classes(object):\n    \"\"\"no_sub_classes context manager.\n\n    Only returns instances of this class and no sub (inherited) classes::\n\n        with no_sub_classes(Group) as Group:\n            Group.objects.find()\n    \"\"\"\n\n    def __init__(self, cls):\n        \"\"\"Construct the no_sub_classes context manager.\n\n        :param cls: the class to turn querying sub classes on\n        \"\"\"\n        self.cls = cls\n        self.cls_initial_subclasses = None\n\n    def __enter__(self):\n        \"\"\"Change the objects default and _auto_dereference values.\"\"\"\n        self.cls_initial_subclasses = self.cls._subclasses\n        self.cls._subclasses = (self.cls._class_name,)\n        return self.cls\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the default and _auto_dereference values.\"\"\"\n        self.cls._subclasses = self.cls_initial_subclasses\n\n\nclass query_counter(object):\n    \"\"\"Query_counter context manager to get the number of queries.\"\"\"\n\n    def __init__(self):\n        \"\"\"Construct the query_counter.\"\"\"\n        self.counter = 0\n        self.db = get_db()\n\n    def __enter__(self):\n        \"\"\"On every with block we need to drop the profile collection.\"\"\"\n        self.db.set_profiling_level(0)\n        self.db.system.profile.drop()\n        self.db.set_profiling_level(2)\n        return self\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the profiling level.\"\"\"\n        self.db.set_profiling_level(0)\n\n    def __eq__(self, value):\n        \"\"\"== Compare querycounter.\"\"\"\n        counter = self._get_count()\n        return value == counter\n\n    def __ne__(self, value):\n        \"\"\"!= Compare querycounter.\"\"\"\n        return not self.__eq__(value)\n\n    def __lt__(self, value):\n        \"\"\"< Compare querycounter.\"\"\"\n        return self._get_count() < value\n\n    def __le__(self, value):\n        \"\"\"<= Compare querycounter.\"\"\"\n        return self._get_count() <= value\n\n    def __gt__(self, value):\n        \"\"\"> Compare querycounter.\"\"\"\n        return self._get_count() > value\n\n    def __ge__(self, value):\n        \"\"\">= Compare querycounter.\"\"\"\n        return self._get_count() >= value\n\n    def __int__(self):\n        \"\"\"int representation.\"\"\"\n        return self._get_count()\n\n    def __repr__(self):\n        \"\"\"repr query_counter as the number of queries.\"\"\"\n        return u\"%s\" % self._get_count()\n\n    def _get_count(self):\n        \"\"\"Get the number of queries.\"\"\"\n        ignore_query = {'ns': {'$ne': '%s.system.indexes' % self.db.name}}\n        count = self.db.system.profile.find(ignore_query).count() - self.counter\n        self.counter += 1       # Account for the query we just fired\n        return count\n\n\n@contextmanager\ndef set_write_concern(collection, write_concerns):\n    combined_concerns = dict(collection.write_concern.document.items())\n    combined_concerns.update(write_concerns)\n    yield collection.with_options(write_concern=WriteConcern(**combined_concerns))\n", "path": "mongoengine/context_managers.py"}], "after_files": [{"content": "from contextlib import contextmanager\nfrom pymongo.write_concern import WriteConcern\nfrom mongoengine.common import _import_class\nfrom mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db\n\n\n__all__ = ('switch_db', 'switch_collection', 'no_dereference',\n           'no_sub_classes', 'query_counter', 'set_write_concern')\n\n\nclass switch_db(object):\n    \"\"\"switch_db alias context manager.\n\n    Example ::\n\n        # Register connections\n        register_connection('default', 'mongoenginetest')\n        register_connection('testdb-1', 'mongoenginetest2')\n\n        class Group(Document):\n            name = StringField()\n\n        Group(name='test').save()  # Saves in the default db\n\n        with switch_db(Group, 'testdb-1') as Group:\n            Group(name='hello testdb!').save()  # Saves in testdb-1\n    \"\"\"\n\n    def __init__(self, cls, db_alias):\n        \"\"\"Construct the switch_db context manager\n\n        :param cls: the class to change the registered db\n        :param db_alias: the name of the specific database to use\n        \"\"\"\n        self.cls = cls\n        self.collection = cls._get_collection()\n        self.db_alias = db_alias\n        self.ori_db_alias = cls._meta.get('db_alias', DEFAULT_CONNECTION_NAME)\n\n    def __enter__(self):\n        \"\"\"Change the db_alias and clear the cached collection.\"\"\"\n        self.cls._meta['db_alias'] = self.db_alias\n        self.cls._collection = None\n        return self.cls\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the db_alias and collection.\"\"\"\n        self.cls._meta['db_alias'] = self.ori_db_alias\n        self.cls._collection = self.collection\n\n\nclass switch_collection(object):\n    \"\"\"switch_collection alias context manager.\n\n    Example ::\n\n        class Group(Document):\n            name = StringField()\n\n        Group(name='test').save()  # Saves in the default db\n\n        with switch_collection(Group, 'group1') as Group:\n            Group(name='hello testdb!').save()  # Saves in group1 collection\n    \"\"\"\n\n    def __init__(self, cls, collection_name):\n        \"\"\"Construct the switch_collection context manager.\n\n        :param cls: the class to change the registered db\n        :param collection_name: the name of the collection to use\n        \"\"\"\n        self.cls = cls\n        self.ori_collection = cls._get_collection()\n        self.ori_get_collection_name = cls._get_collection_name\n        self.collection_name = collection_name\n\n    def __enter__(self):\n        \"\"\"Change the _get_collection_name and clear the cached collection.\"\"\"\n\n        @classmethod\n        def _get_collection_name(cls):\n            return self.collection_name\n\n        self.cls._get_collection_name = _get_collection_name\n        self.cls._collection = None\n        return self.cls\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the collection.\"\"\"\n        self.cls._collection = self.ori_collection\n        self.cls._get_collection_name = self.ori_get_collection_name\n\n\nclass no_dereference(object):\n    \"\"\"no_dereference context manager.\n\n    Turns off all dereferencing in Documents for the duration of the context\n    manager::\n\n        with no_dereference(Group) as Group:\n            Group.objects.find()\n    \"\"\"\n\n    def __init__(self, cls):\n        \"\"\"Construct the no_dereference context manager.\n\n        :param cls: the class to turn dereferencing off on\n        \"\"\"\n        self.cls = cls\n\n        ReferenceField = _import_class('ReferenceField')\n        GenericReferenceField = _import_class('GenericReferenceField')\n        ComplexBaseField = _import_class('ComplexBaseField')\n\n        self.deref_fields = [k for k, v in self.cls._fields.iteritems()\n                             if isinstance(v, (ReferenceField,\n                                               GenericReferenceField,\n                                               ComplexBaseField))]\n\n    def __enter__(self):\n        \"\"\"Change the objects default and _auto_dereference values.\"\"\"\n        for field in self.deref_fields:\n            self.cls._fields[field]._auto_dereference = False\n        return self.cls\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the default and _auto_dereference values.\"\"\"\n        for field in self.deref_fields:\n            self.cls._fields[field]._auto_dereference = True\n        return self.cls\n\n\nclass no_sub_classes(object):\n    \"\"\"no_sub_classes context manager.\n\n    Only returns instances of this class and no sub (inherited) classes::\n\n        with no_sub_classes(Group) as Group:\n            Group.objects.find()\n    \"\"\"\n\n    def __init__(self, cls):\n        \"\"\"Construct the no_sub_classes context manager.\n\n        :param cls: the class to turn querying sub classes on\n        \"\"\"\n        self.cls = cls\n        self.cls_initial_subclasses = None\n\n    def __enter__(self):\n        \"\"\"Change the objects default and _auto_dereference values.\"\"\"\n        self.cls_initial_subclasses = self.cls._subclasses\n        self.cls._subclasses = (self.cls._class_name,)\n        return self.cls\n\n    def __exit__(self, t, value, traceback):\n        \"\"\"Reset the default and _auto_dereference values.\"\"\"\n        self.cls._subclasses = self.cls_initial_subclasses\n\n\nclass query_counter(object):\n    \"\"\"Query_counter context manager to get the number of queries.\n    This works by updating the `profiling_level` of the database so that all queries get logged,\n    resetting the db.system.profile collection at the beginnig of the context and counting the new entries.\n\n    This was designed for debugging purpose. In fact it is a global counter so queries issued by other threads/processes\n    can interfere with it\n\n    Be aware that:\n    - Iterating over large amount of documents (>101) makes pymongo issue `getmore` queries to fetch the next batch of\n        documents (https://docs.mongodb.com/manual/tutorial/iterate-a-cursor/#cursor-batches)\n    - Some queries are ignored by default by the counter (killcursors, db.system.indexes)\n    \"\"\"\n\n    def __init__(self):\n        \"\"\"Construct the query_counter\n        \"\"\"\n        self.db = get_db()\n        self.initial_profiling_level = None\n        self._ctx_query_counter = 0             # number of queries issued by the context\n\n        self._ignored_query = {\n            'ns':\n                {'$ne': '%s.system.indexes' % self.db.name},\n            'op':\n                {'$ne': 'killcursors'}\n        }\n\n    def _turn_on_profiling(self):\n        self.initial_profiling_level = self.db.profiling_level()\n        self.db.set_profiling_level(0)\n        self.db.system.profile.drop()\n        self.db.set_profiling_level(2)\n\n    def _resets_profiling(self):\n        self.db.set_profiling_level(self.initial_profiling_level)\n\n    def __enter__(self):\n        self._turn_on_profiling()\n        return self\n\n    def __exit__(self, t, value, traceback):\n        self._resets_profiling()\n\n    def __eq__(self, value):\n        counter = self._get_count()\n        return value == counter\n\n    def __ne__(self, value):\n        return not self.__eq__(value)\n\n    def __lt__(self, value):\n        return self._get_count() < value\n\n    def __le__(self, value):\n        return self._get_count() <= value\n\n    def __gt__(self, value):\n        return self._get_count() > value\n\n    def __ge__(self, value):\n        return self._get_count() >= value\n\n    def __int__(self):\n        return self._get_count()\n\n    def __repr__(self):\n        \"\"\"repr query_counter as the number of queries.\"\"\"\n        return u\"%s\" % self._get_count()\n\n    def _get_count(self):\n        \"\"\"Get the number of queries by counting the current number of entries in db.system.profile\n        and substracting the queries issued by this context. In fact everytime this is called, 1 query is\n        issued so we need to balance that\n        \"\"\"\n        count = self.db.system.profile.find(self._ignored_query).count() - self._ctx_query_counter\n        self._ctx_query_counter += 1    # Account for the query we just issued to gather the information\n        return count\n\n\n@contextmanager\ndef set_write_concern(collection, write_concerns):\n    combined_concerns = dict(collection.write_concern.document.items())\n    combined_concerns.update(write_concerns)\n    yield collection.with_options(write_concern=WriteConcern(**combined_concerns))\n", "path": "mongoengine/context_managers.py"}]}