diff --git a/mongoengine/connection.py b/mongoengine/connection.py index abab269f1..a2e19e008 100644 --- a/mongoengine/connection.py +++ b/mongoengine/connection.py @@ -96,6 +96,8 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False): msg = 'You have not defined a default connection' raise ConnectionError(msg) conn_settings = _connection_settings[alias].copy() + conn_settings.pop('query_trace', None) + conn_settings.pop('trace_depth', None) if hasattr(pymongo, 'version_tuple'): # Support for 2.1+ conn_settings.pop('name', None) diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py index 29537793c..86eb25c03 100644 --- a/mongoengine/queryset/queryset.py +++ b/mongoengine/queryset/queryset.py @@ -3,8 +3,10 @@ import copy import itertools import operator +import os import pprint import re +import sys import warnings from bson.code import Code @@ -20,7 +22,7 @@ from mongoengine.queryset import transform from mongoengine.queryset.field_list import QueryFieldList from mongoengine.queryset.visitor import Q, QNode - +from mongoengine.connection import _connection_settings __all__ = ('QuerySet', 'DO_NOTHING', 'NULLIFY', 'CASCADE', 'DENY', 'PULL') @@ -38,6 +40,58 @@ RE_TYPE = type(re.compile('')) +# Borrowed from CPython's stdlib logging +# https://github.com/python/cpython/blob/master/Lib/logging/__init__.py +if hasattr(sys, '_getframe'): + currentframe = lambda: sys._getframe(1) +else: #pragma: no cover + def currentframe(): + """Return the frame object for the caller's stack frame.""" + try: + raise Exception + except Exception: + return sys.exc_info()[2].tb_frame.f_back + + +def dummy(): + pass + + +_srcfile = os.path.normcase(dummy.__code__.co_filename) + +def find_callers(): + """ + Find the stack frame of the caller so that we can note the source + file name, line number and function name. + """ + try: + if _connection_settings['default']['query_trace'] is True: + trace_depth = 3 + try: + trace_depth = _connection_settings['default']['trace_depth'] + except KeyError: + pass + except KeyError: + return None + trace_comment = [] + f = currentframe() + #On some versions of IronPython, currentframe() returns None if + #IronPython isn't run with -X:Frames. + if f is not None: + f = f.f_back + frame = 0 + while hasattr(f, "f_code") and frame < trace_depth: + co = f.f_code + filename = os.path.normcase(co.co_filename) + if filename == _srcfile: + f = f.f_back + continue + trace_comment.append('{}({})'.format(co.co_filename, f.f_lineno)) + f = f.f_back + frame += 1 + return trace_comment + + class QuerySet(object): """A set of results returned from a query. Wraps a MongoDB cursor, providing :class:`~mongoengine.Document` objects as the results. @@ -67,6 +121,7 @@ def __init__(self, document, collection): self._result_cache = [] self._has_more = True self._len = None + self._comment = find_callers() # If inheritance is allowed, only return instances and instances of # subclasses of the class being used @@ -903,6 +958,14 @@ def clear_initial_query(self): queryset._initial_query = {} return queryset + def comment(self, text): + """Add a comment to the query. + + See https://docs.mongodb.com/manual/reference/method/cursor.comment/#cursor.comment + for details. + """ + return self._chainable_method("comment", text) + def explain(self, format=False): """Return an explain plan record for the :class:`~mongoengine.queryset.QuerySet`\ 's cursor. @@ -1334,6 +1397,13 @@ def _cursor(self): self._cursor_obj = self._collection.find(self._query, **self._cursor_args) + + # Auto-omment with stack trace. Set query_trace=True and optionally + # set depth with trace_depth in MONGODB_SETTINGS + our_comment = find_callers() + if our_comment is not None: + self._cursor_obj.comment(our_comment) + # Apply where clauses to cursor if self._where_clause: where_clause = self._sub_js_fields(self._where_clause) @@ -1637,6 +1707,25 @@ def field_path_sub(match): code) return code + def _chainable_method(self, method_name, val): + """Call a particular method on the PyMongo cursor call a particular chainable method + with the provided value. + """ + queryset = self.clone() + + # Get an existing cursor object or create a new one + cursor = queryset._cursor + + # Find the requested method on the cursor and call it with the + # provided value + getattr(cursor, method_name)(val) + + # Cache the value on the queryset._{method_name} + setattr(queryset, '_' + method_name, val) + + return queryset + + # Deprecated def ensure_index(self, **kwargs): """Deprecated use :func:`Document.ensure_index`""" diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index 681e257b8..a8bca9888 100644 --- a/tests/queryset/queryset.py +++ b/tests/queryset/queryset.py @@ -1760,6 +1760,21 @@ class Author(Document): names = [a.author.name for a in Author.objects.order_by('-author__age')] self.assertEqual(names, ['User A', 'User B', 'User C']) + def test_comment(self): + """Make sure adding a comment to the query works.""" + class User(Document): + age = IntField() + + with db_ops_tracker() as q: + adult = (User.objects.filter(age__gte=18) + .comment('looking for an adult') + .first()) + ops = q.get_ops() + self.assertEqual(len(ops), 1) + op = ops[0] + self.assertEqual(op['query']['$query'], {'age': {'$gte': 18}}) + self.assertEqual(op['query']['$comment'], 'looking for an adult') + def test_map_reduce(self): """Ensure map/reduce is both mapping and reducing. """