From 4a2b782ba75a8379b1bc73e2df99349a8601101e Mon Sep 17 00:00:00 2001 From: James Satterfield Date: Mon, 21 May 2018 14:13:21 -0500 Subject: [PATCH 1/3] Mongo cursor support --- mongoengine/queryset/queryset.py | 27 +++++++++++++++++++++++++++ tests/queryset/queryset.py | 15 +++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py index 29537793c..a11dcdf30 100644 --- a/mongoengine/queryset/queryset.py +++ b/mongoengine/queryset/queryset.py @@ -903,6 +903,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. @@ -1637,6 +1645,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. """ From b303578ee4e463057900f323f0db802affa836cd Mon Sep 17 00:00:00 2001 From: James Satterfield Date: Wed, 6 Jun 2018 13:45:59 -0500 Subject: [PATCH 2/3] Auto comment with trace. Set query_trace=True and optionally trace_depth=N in MONGODB_SETTINGS to enable --- mongoengine/connection.py | 2 + mongoengine/queryset/queryset.py | 68 +++++++++++++++++++++++++++++++- 2 files changed, 68 insertions(+), 2 deletions(-) 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 a11dcdf30..fc430b0c1 100644 --- a/mongoengine/queryset/queryset.py +++ b/mongoengine/queryset/queryset.py @@ -1,10 +1,13 @@ from __future__ import absolute_import import copy +import io import itertools import operator +import os import pprint import re +import sys import warnings from bson.code import Code @@ -20,7 +23,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 +41,59 @@ 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)) + # rv.append([co.co_filename, f.f_lineno, co.co_name]) + 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 +123,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 @@ -1342,6 +1399,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) @@ -1663,7 +1727,7 @@ def _chainable_method(self, method_name, val): return queryset - + # Deprecated def ensure_index(self, **kwargs): """Deprecated use :func:`Document.ensure_index`""" From d385333a94237569bcbf6744f5580348696b4797 Mon Sep 17 00:00:00 2001 From: James Satterfield Date: Wed, 6 Jun 2018 13:48:36 -0500 Subject: [PATCH 3/3] Removed io module and unneeded comment --- mongoengine/queryset/queryset.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py index fc430b0c1..86eb25c03 100644 --- a/mongoengine/queryset/queryset.py +++ b/mongoengine/queryset/queryset.py @@ -1,7 +1,6 @@ from __future__ import absolute_import import copy -import io import itertools import operator import os @@ -88,7 +87,6 @@ def find_callers(): f = f.f_back continue trace_comment.append('{}({})'.format(co.co_filename, f.f_lineno)) - # rv.append([co.co_filename, f.f_lineno, co.co_name]) f = f.f_back frame += 1 return trace_comment