Added Q class for building advanced queries

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

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']:

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()