Auto comment - WIP

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)

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`"""

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.
         """