From da2d282cf645d234208d2556f9cb2b6d3d2248dd Mon Sep 17 00:00:00 2001 From: Harry Marr Date: Sat, 9 Jan 2010 22:19:33 +0000 Subject: [PATCH] Added Q class for building advanced queries --- docs/userguide.rst | 26 ++++++++++++ mongoengine/queryset.py | 94 ++++++++++++++++++++++++++++++++++++++++- tests/queryset.py | 80 +++++++++++++++++++++++++++++++++-- 3 files changed, 195 insertions(+), 5 deletions(-) diff --git a/docs/userguide.rst b/docs/userguide.rst index bba6ffbb..c030ea61 100644 --- a/docs/userguide.rst +++ b/docs/userguide.rst @@ -454,6 +454,32 @@ would be generating "tag-clouds":: from operator import itemgetter top_tags = sorted(tag_freqs.items(), key=itemgetter(1), reverse=True)[:10] +Advanced queries +---------------- +Sometimes calling a :class:`~mongoengine.queryset.QuerySet` object with keyword +arguments can't fully express the query you want to use -- for example if you +need to combine a number of constraints using *and* and *or*. This is made +possible in MongoEngine through the :class:`~mongoengine.queryset.Q` class. +A :class:`~mongoengine.queryset.Q` object represents part of a query, and +can be initialised using the same keyword-argument syntax you use to query +documents. To build a complex query, you may combine +:class:`~mongoengine.queryset.Q` objects using the ``&`` (and) and ``|`` (or) +operators. To use :class:`~mongoengine.queryset.Q` objects, pass them in +as positional arguments to :attr:`Document.objects` when you filter it by +calling it with keyword arguments:: + + # Get published posts + Post.objects(Q(published=True) | Q(publish_date__lte=datetime.now())) + + # Get top posts + Post.objects((Q(featured=True) & Q(hits__gte=1000)) | Q(hits__gte=5000)) + +.. warning:: + Only use these advanced queries if absolutely necessary as they will execute + significantly slower than regular queries. This is because they are not + natively supported by MongoDB -- they are compiled to Javascript and sent + to the server for execution. + Atomic updates -------------- Documents may be updated atomically by using the diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 39cef283..2a4f5383 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -1,9 +1,11 @@ from connection import _get_db import pymongo +import copy -__all__ = ['queryset_manager', 'InvalidQueryError', 'InvalidCollectionError'] +__all__ = ['queryset_manager', 'Q', 'InvalidQueryError', + 'InvalidCollectionError'] class InvalidQueryError(Exception): @@ -14,6 +16,88 @@ class OperationError(Exception): pass +class Q(object): + + OR = '||' + AND = '&&' + OPERATORS = { + 'eq': 'this.%(field)s == %(value)s', + 'neq': 'this.%(field)s != %(value)s', + 'gt': 'this.%(field)s > %(value)s', + 'gte': 'this.%(field)s >= %(value)s', + 'lt': 'this.%(field)s < %(value)s', + 'lte': 'this.%(field)s <= %(value)s', + 'lte': 'this.%(field)s <= %(value)s', + 'in': 'this.%(field)s.indexOf(%(value)s) != -1', + 'nin': 'this.%(field)s.indexOf(%(value)s) == -1', + 'mod': '%(field)s %% %(value)s', + 'all': ('%(value)s.every(function(a){' + 'return this.%(field)s.indexOf(a) != -1 })'), + 'size': 'this.%(field)s.length == %(value)s', + 'exists': 'this.%(field)s != null', + } + + def __init__(self, **query): + self.query = [query] + + def _combine(self, other, op): + obj = Q() + obj.query = ['('] + copy.deepcopy(self.query) + [op] + obj.query += copy.deepcopy(other.query) + [')'] + return obj + + def __or__(self, other): + return self._combine(other, self.OR) + + def __and__(self, other): + return self._combine(other, self.AND) + + def as_js(self, document): + js = [] + js_scope = {} + for i, item in enumerate(self.query): + if isinstance(item, dict): + item_query = QuerySet._transform_query(document, **item) + # item_query will values will either be a value or a dict + js.append(self._item_query_as_js(item_query, js_scope, i)) + else: + js.append(item) + return pymongo.code.Code(' '.join(js), js_scope) + + def _item_query_as_js(self, item_query, js_scope, item_num): + # item_query will be in one of the following forms + # {'age': 25, 'name': 'Test'} + # {'age': {'$lt': 25}, 'name': {'$in': ['Test', 'Example']} + # {'age': {'$lt': 25, '$gt': 18}} + js = [] + for i, (key, value) in enumerate(item_query.items()): + op = 'eq' + # Construct a variable name for the value in the JS + value_name = 'i%sf%s' % (item_num, i) + if isinstance(value, dict): + # Multiple operators for this field + for j, (op, value) in enumerate(value.items()): + # Create a custom variable name for this operator + op_value_name = '%so%s' % (value_name, j) + # Update the js scope with the value for this op + js_scope[op_value_name] = value + # Construct the JS that uses this op + operation_js = Q.OPERATORS[op.strip('$')] % { + 'field': key, + 'value': op_value_name + } + js.append(operation_js) + else: + js_scope[value_name] = value + # Construct the JS for this field + field_js = Q.OPERATORS[op.strip('$')] % { + 'field': key, + 'value': value_name + } + js.append(field_js) + return ' && '.join(js) + + class QuerySet(object): """A set of results returned from a query. Wraps a MongoDB cursor, providing :class:`~mongoengine.Document` objects as the results. @@ -24,6 +108,7 @@ class QuerySet(object): self._collection_obj = collection self._accessed_collection = False self._query = {} + self._where_clauses = [] # If inheritance is allowed, only return instances and instances of # subclasses of the class being used @@ -55,10 +140,12 @@ class QuerySet(object): self._collection.ensure_index(index_list) return self - def __call__(self, **query): + def __call__(self, *q_objs, **query): """Filter the selected documents by calling the :class:`~mongoengine.QuerySet` with a query. """ + for q in q_objs: + self._where_clauses.append(q.as_js(self._document)) query = QuerySet._transform_query(_doc_cls=self._document, **query) self._query.update(query) return self @@ -89,6 +176,9 @@ class QuerySet(object): def _cursor(self): if not self._cursor_obj: self._cursor_obj = self._collection.find(self._query) + # Apply where clauses to cursor + for js in self._where_clauses: + self._cursor_obj.where(js) # apply default ordering if self._document._meta['ordering']: diff --git a/tests/queryset.py b/tests/queryset.py index 698ada9c..fe0f6483 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -1,5 +1,6 @@ import unittest import pymongo +from datetime import datetime from mongoengine.queryset import QuerySet from mongoengine import * @@ -16,7 +17,7 @@ class QuerySetTest(unittest.TestCase): self.Person = Person def test_initialisation(self): - """Ensure that CollectionManager is correctly initialised. + """Ensure that a QuerySet is correctly initialised by QuerySetManager. """ self.assertTrue(isinstance(self.Person.objects, QuerySet)) self.assertEqual(self.Person.objects._collection.name(), @@ -48,6 +49,9 @@ class QuerySetTest(unittest.TestCase): person2 = self.Person(name="User B", age=30) person2.save() + q1 = Q(name='test') + q2 = Q(age__gte=18) + # Find all people in the collection people = self.Person.objects self.assertEqual(len(people), 2) @@ -134,8 +138,6 @@ class QuerySetTest(unittest.TestCase): def test_ordering(self): """Ensure default ordering is applied and can be overridden. """ - from datetime import datetime - class BlogPost(Document): title = StringField() published_date = DateTimeField() @@ -144,6 +146,8 @@ class QuerySetTest(unittest.TestCase): 'ordering': ['-published_date'] } + BlogPost.drop_collection() + blog_post_1 = BlogPost(title="Blog Post #1", published_date=datetime(2010, 1, 5, 0, 0 ,0)) blog_post_2 = BlogPost(title="Blog Post #2", @@ -176,6 +180,8 @@ class QuerySetTest(unittest.TestCase): content = StringField() author = EmbeddedDocumentField(User) + BlogPost.drop_collection() + post = BlogPost(content='Had a good coffee today...') post.author = User(name='Test User') post.save() @@ -186,6 +192,42 @@ class QuerySetTest(unittest.TestCase): BlogPost.drop_collection() + def test_q(self): + class BlogPost(Document): + publish_date = DateTimeField() + published = BooleanField() + + BlogPost.drop_collection() + + post1 = BlogPost(publish_date=datetime(2010, 1, 8), published=False) + post1.save() + + post2 = BlogPost(publish_date=datetime(2010, 1, 15), published=True) + post2.save() + + post3 = BlogPost(published=True) + post3.save() + + post4 = BlogPost(publish_date=datetime(2010, 1, 8)) + post4.save() + + post5 = BlogPost(publish_date=datetime(2010, 1, 15)) + post5.save() + + post6 = BlogPost(published=False) + post6.save() + + date = datetime(2010, 1, 10) + q = BlogPost.objects(Q(publish_date__lte=date) | Q(published=True)) + posts = [post.id for post in q] + + published_posts = (post1, post2, post3, post4) + self.assertTrue(all(obj.id in posts for obj in published_posts)) + + self.assertFalse(any(obj.id in posts for obj in [post5, post6])) + + BlogPost.drop_collection() + def test_delete(self): """Ensure that documents are properly deleted from the database. """ @@ -428,5 +470,37 @@ class QuerySetTest(unittest.TestCase): self.Person.drop_collection() +class QTest(unittest.TestCase): + + def test_or_and(self): + q1 = Q(name='test') + q2 = Q(age__gte=18) + + query = ['(', {'name': 'test'}, '||', {'age__gte': 18}, ')'] + self.assertEqual((q1 | q2).query, query) + + query = ['(', {'name': 'test'}, '&&', {'age__gte': 18}, ')'] + self.assertEqual((q1 & q2).query, query) + + query = ['(', '(', {'name': 'test'}, '&&', {'age__gte': 18}, ')', '||', + {'name': 'example'}, ')'] + self.assertEqual((q1 & q2 | Q(name='example')).query, query) + + def test_item_query_as_js(self): + """Ensure that the _item_query_as_js utilitiy method works properly. + """ + q = Q() + examples = [ + ({'name': 'test'}, 'this.name == i0f0', {'i0f0': 'test'}), + ({'age': {'$gt': 18}}, 'this.age > i0f0o0', {'i0f0o0': 18}), + ({'name': 'test', 'age': {'$gt': 18, '$lte': 65}}, + 'this.age <= i0f0o0 && this.age > i0f0o1 && this.name == i0f1', + {'i0f0o0': 65, 'i0f0o1': 18, 'i0f1': 'test'}), + ] + for item, js, scope in examples: + test_scope = {} + self.assertEqual(q._item_query_as_js(item, test_scope, 0), js) + self.assertEqual(scope, test_scope) + if __name__ == '__main__': unittest.main()