Merge pull request #1871 from bagerard/improve_query_counter
Fix few things related to query_counter context manager
This commit is contained in:
commit
5b1693a908
@ -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
|
||||
|
||||
|
||||
|
@ -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(),
|
||||
|
@ -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()
|
||||
|
Loading…
x
Reference in New Issue
Block a user