Permalink
Browse files

Refactor tracker, add stack traces to queries

  • Loading branch information...
1 parent 330f117 commit af131a3ceb106dfbc0500af9ca90b62ab3165570 @hmarr hmarr committed May 26, 2011
@@ -1,173 +1,220 @@
import functools
+import traceback
import time
+import inspect
+import os
+import SocketServer
+import django
+from django.conf import settings
+
+import pymongo
import pymongo.collection
import pymongo.cursor
+__all__ = ['queries', 'inserts', 'updates', 'removes', 'install_tracker',
+ 'uninstall_tracker', 'reset']
+
+
+_original_methods = {
+ 'insert': pymongo.collection.Collection.insert,
+ 'update': pymongo.collection.Collection.update,
+ 'remove': pymongo.collection.Collection.remove,
+ 'refresh': pymongo.cursor.Cursor._refresh,
+}
+
+queries = []
+inserts = []
+updates = []
+removes = []
+
+
+# Wrap Cursor._refresh for getting queries
+@functools.wraps(_original_methods['insert'])
+def _insert(collection_self, doc_or_docs, manipulate=True,
+ safe=False, check_keys=True, **kwargs):
+ start_time = time.time()
+ result = _original_methods['insert'](
+ collection_self,
+ doc_or_docs,
+ safe=safe,
+ **kwargs
+ )
+ total_time = (time.time() - start_time) * 1000
+
+ inserts.append({
+ 'document': doc_or_docs,
+ 'safe': safe,
+ 'time': total_time,
+ })
+ return result
+
+# Wrap Cursor._refresh for getting queries
+@functools.wraps(_original_methods['update'])
+def _update(collection_self, spec, document, upsert=False,
+ maniuplate=False, safe=False, multi=False, **kwargs):
+ start_time = time.time()
+ result = _original_methods['update'](
+ collection_self,
+ spec,
+ document,
+ upsert=False,
+ safe=False,
+ multi=False,
+ **kwargs
+ )
+ total_time = (time.time() - start_time) * 1000
+
+ updates.append({
+ 'document': document,
+ 'upsert': upsert,
+ 'multi': multi,
+ 'spec': spec,
+ 'safe': safe,
+ 'time': total_time,
+ })
+ return result
+
+# Wrap Cursor._refresh for getting queries
+@functools.wraps(_original_methods['remove'])
+def _remove(collection_self, spec_or_id, safe=False, **kwargs):
+ start_time = time.time()
+ result = _original_methods['insert'](
+ collection_self,
+ spec_or_id,
+ safe=safe,
+ **kwargs
+ )
+ total_time = (time.time() - start_time) * 1000
+
+ removes.append({
+ 'spec_or_id': spec_or_id,
+ 'safe': safe,
+ 'time': total_time,
+ })
+ return result
+
+# Wrap Cursor._refresh for getting queries
+@functools.wraps(_original_methods['refresh'])
+def _cursor_refresh(cursor_self):
+ # Look up __ private instance variables
+ def privar(name):
+ return getattr(cursor_self, '_Cursor__{0}'.format(name))
+
+ if privar('id') is not None:
+ # getMore not query - move on
+ return _original_methods['refresh'](cursor_self)
+
+ # NOTE: See pymongo/cursor.py+557 [_refresh()] and
+ # pymongo/message.py for where information is stored
+
+ # Time the actual query
+ start_time = time.time()
+ result = _original_methods['refresh'](cursor_self)
+ total_time = (time.time() - start_time) * 1000
+
+ query_son = privar('query_spec')()
+
+ query_data = {
+ 'time': total_time,
+ 'operation': 'query',
+ 'stack_trace': _tidy_stacktrace(reversed(inspect.stack())),
+ }
+
+ # Collection in format <db_name>.<collection_name>
+ collection_name = privar('collection')
+ query_data['collection'] = collection_name.full_name.split('.')[1]
+
+ if query_data['collection'] == '$cmd':
+ query_data['operation'] = 'command'
+ # Handle count as a special case
+ if 'count' in query_son:
+ # Information is in a different format to a standar query
+ query_data['collection'] = query_son['count']
+ query_data['operation'] = 'count'
+ query_data['skip'] = query_son.get('skip')
+ query_data['limit'] = query_son.get('limit')
+ query_data['query'] = query_son['query']
+ else:
+ # Normal Query
+ query_data['skip'] = privar('skip')
+ query_data['limit'] = privar('limit')
+ query_data['query'] = query_son['$query']
+ query_data['ordering'] = _get_ordering(query_son)
+
+ queries.append(query_data)
+
+ return result
+
+def install_tracker():
+ if pymongo.collection.Collection.insert != _insert:
+ pymongo.collection.Collection.insert = _insert
+ if pymongo.collection.Collection.update != _update:
+ pymongo.collection.Collection.update = _update
+ if pymongo.collection.Collection.remove != _remove:
+ pymongo.collection.Collection.remove = _remove
+ if pymongo.cursor.Cursor._refresh != _cursor_refresh:
+ pymongo.cursor.Cursor._refresh = _cursor_refresh
+
+def uninstall_tracker():
+ if pymongo.collection.Collection.insert == _insert:
+ pymongo.collection.Collection.insert = _original_methods['insert']
+ if pymongo.collection.Collection.update == _update:
+ pymongo.collection.Collection.update = _original_methods['update']
+ if pymongo.collection.Collection.remove == _remove:
+ pymongo.collection.Collection.remove = _original_methods['remove']
+ if pymongo.cursor.Cursor._refresh == _cursor_refresh:
+ pymongo.cursor.Cursor._refresh = _original_methods['cursor_refresh']
+
+def reset():
+ global queries, inserts, updates, removes
+ queries = []
+ inserts = []
+ updates = []
+ removes = []
+
+def _get_ordering(son):
+ """Helper function to extract formatted ordering from dict.
+ """
+ def fmt(field, direction):
+ return '{0}{1}'.format({-1: '-', 1: '+'}[direction], field)
+
+ if '$orderby' in son:
+ return ', '.join(fmt(f, d) for f, d in son['$orderby'].items())
+
+# Taken from Django Debug Toolbar 0.8.6
+def _tidy_stacktrace(stack):
+ """
+ Clean up stacktrace and remove all entries that:
+ 1. Are part of Django (except contrib apps)
+ 2. Are part of SocketServer (used by Django's dev server)
+ 3. Are the last entry (which is part of our stacktracing code)
-class MongoOperationTracker(object):
- """Track operations sent to MongoDB via PyMongo.
+ ``stack`` should be a list of frame tuples from ``inspect.stack()``
"""
+ django_path = os.path.realpath(os.path.dirname(django.__file__))
+ django_path = os.path.normpath(os.path.join(django_path, '..'))
+ socketserver_path = os.path.realpath(os.path.dirname(SocketServer.__file__))
+ pymongo_path = os.path.realpath(os.path.dirname(pymongo.__file__))
+
+ trace = []
+ for frame, path, line_no, func_name, text in (f[:5] for f in stack):
+ s_path = os.path.realpath(path)
+ # Support hiding of frames -- used in various utilities that provide
+ # inspection.
+ if '__traceback_hide__' in frame.f_locals:
+ continue
+ if getattr(settings, 'DEBUG_TOOLBAR_CONFIG', {}).get('HIDE_DJANGO_SQL', True) \
+ and django_path in s_path and not 'django/contrib' in s_path:
+ continue
+ if socketserver_path in s_path:
+ continue
+ if pymongo_path in s_path:
+ continue
+ if not text:
+ text = ''
+ else:
+ text = (''.join(text)).strip()
+ trace.append((path, line_no, func_name, text))
+ return trace
- def __init__(self):
- self.queries = []
- self.inserts = []
- self.updates = []
- self.removes = []
- self.active = False
-
- def install(self):
- self.original_methods = {
- 'insert': pymongo.collection.Collection.insert,
- 'update': pymongo.collection.Collection.update,
- 'remove': pymongo.collection.Collection.remove,
- 'refresh': pymongo.cursor.Cursor._refresh,
- }
-
- # Wrap Cursor._refresh for getting queries
- @functools.wraps(self.original_methods['insert'])
- def insert(collection_self, doc_or_docs, manipulate=True,
- safe=False, check_keys=True, **kwargs):
- start_time = time.time()
- result = self.original_methods['insert'](
- collection_self,
- doc_or_docs,
- safe=safe,
- **kwargs
- )
- total_time = (time.time() - start_time) * 1000
-
- if self.active:
- self.inserts.append({
- 'document': doc_or_docs,
- 'safe': safe,
- 'time': total_time,
- })
- return result
-
- pymongo.collection.Collection.insert = insert
-
- # Wrap Cursor._refresh for getting queries
- @functools.wraps(self.original_methods['update'])
- def update(collection_self, spec, document, upsert=False,
- maniuplate=False, safe=False, multi=False, **kwargs):
- start_time = time.time()
- result = self.original_methods['update'](
- collection_self,
- spec,
- document,
- upsert=False,
- safe=False,
- multi=False,
- **kwargs
- )
- total_time = (time.time() - start_time) * 1000
-
- if self.active:
- self.updates.append({
- 'document': document,
- 'upsert': upsert,
- 'multi': multi,
- 'spec': spec,
- 'safe': safe,
- 'time': total_time,
- })
- return result
-
- pymongo.collection.Collection.update = update
-
- # Wrap Cursor._refresh for getting queries
- @functools.wraps(self.original_methods['remove'])
- def remove(collection_self, spec_or_id, safe=False, **kwargs):
- start_time = time.time()
- result = self.original_methods['insert'](
- collection_self,
- spec_or_id,
- safe=safe,
- **kwargs
- )
- total_time = (time.time() - start_time) * 1000
-
- if self.active:
- self.removes.append({
- 'spec_or_id': spec_or_id,
- 'safe': safe,
- 'time': total_time,
- })
- return result
-
- pymongo.collection.Collection.remove = remove
-
- # Wrap Cursor._refresh for getting queries
- @functools.wraps(self.original_methods['refresh'])
- def cursor_refresh(cursor_self):
- # Look up __ private instance variables
- def privar(name):
- return getattr(cursor_self, '_Cursor__{0}'.format(name))
-
- if not self.active or privar('id') is not None:
- # Inactive or getMore - move on
- return self.original_methods['refresh'](cursor_self)
-
- # NOTE: See pymongo/cursor.py+557 [_refresh()] and
- # pymongo/message.py for where information is stored
-
- # Time the actual query
- start_time = time.time()
- result = self.original_methods['refresh'](cursor_self)
- total_time = (time.time() - start_time) * 1000
-
- query_son = privar('query_spec')()
-
- query_data = {
- 'time': total_time,
- 'operation': 'query',
- }
-
- # Collection in format <db_name>.<collection_name>
- collection_name = privar('collection')
- query_data['collection'] = collection_name.full_name.split('.')[1]
-
- if query_data['collection'] == '$cmd':
- query_data['operation'] = 'command'
- # Handle count as a special case
- if 'count' in query_son:
- # Information is in a different format to a standar query
- query_data['collection'] = query_son['count']
- query_data['operation'] = 'count'
- query_data['skip'] = query_son.get('skip')
- query_data['limit'] = query_son.get('limit')
- query_data['query'] = query_son['query']
- else:
- # Normal Query
- query_data['skip'] = privar('skip')
- query_data['limit'] = privar('limit')
- query_data['query'] = query_son['$query']
- query_data['ordering'] = self._get_ordering(query_son)
-
- self.queries.append(query_data)
-
- return result
-
- pymongo.cursor.Cursor._refresh = cursor_refresh
-
- def _get_ordering(self, son):
- def fmt(field, direction):
- return '{0}{1}'.format({-1: '-', 1: '+'}[direction], field)
-
- if '$orderby' in son:
- return ', '.join(fmt(f, d) for f, d in son['$orderby'].items())
-
- def uninstall(self):
- pymongo.cursor.Cursor._refresh = self.original_methods['refresh']
- pymongo.cursor.Collection.insert = self.original_methods['insert']
- pymongo.cursor.Collection.update = self.original_methods['update']
-
- def reset(self):
- self.queries = []
-
- def start(self):
- self.active = True
-
- def stop(self):
- self.active = False
Oops, something went wrong.

0 comments on commit af131a3

Please sign in to comment.