pub/mongoengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #1871 from bagerard/improve_query_counter

Browse Source 
Fix few things related to query_counter context manager
This commit is contained in:
erdenezul
2018-09-07 14:22:12 +08:00
committed by GitHub
parent fd7c00da49 8ff82996fb
commit 5b1693a908
3 changed files with 143 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
mongoengine/context_managers.py
View File

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

11
tests/queryset/queryset.py
View File

@@ -4751,6 +4751,7 @@ class QuerySetTest(unittest.TestCase):
Person(name="No: %s" % i).save()
with query_counter() as q:
try:
self.assertEqual(q, 0)
people = Person.objects.no_cache()
@@ -4762,6 +4763,14 @@ class QuerySetTest(unittest.TestCase):
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)
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(),

91
tests/test_context_managers.py
View File

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