diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index fae2aabf..6e546564 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -7,6 +7,7 @@ import pymongo.dbref import pymongo.objectid import re import copy +import itertools __all__ = ['queryset_manager', 'Q', 'InvalidQueryError', 'InvalidCollectionError'] @@ -18,6 +19,7 @@ REPR_OUTPUT_SIZE = 20 class DoesNotExist(Exception): pass + class MultipleObjectsReturned(Exception): pass @@ -29,13 +31,244 @@ class InvalidQueryError(Exception): class OperationError(Exception): pass + class InvalidCollectionError(Exception): pass + RE_TYPE = type(re.compile('')) -class Q(object): +class QNodeVisitor(object): + """Base visitor class for visiting Q-object nodes in a query tree. + """ + + def visit_combination(self, combination): + """Called by QCombination objects. + """ + return combination + + def visit_query(self, query): + """Called by (New)Q objects. + """ + return query + + +class SimplificationVisitor(QNodeVisitor): + """Simplifies query trees by combinging unnecessary 'and' connection nodes + into a single Q-object. + """ + + def visit_combination(self, combination): + if combination.operation == combination.AND: + # The simplification only applies to 'simple' queries + if all(isinstance(node, Q) for node in combination.children): + queries = [node.query for node in combination.children] + return Q(**self._query_conjunction(queries)) + return combination + + def _query_conjunction(self, queries): + """Merges query dicts - effectively &ing them together. + """ + query_ops = set() + combined_query = {} + for query in queries: + ops = set(query.keys()) + # Make sure that the same operation isn't applied more than once + # to a single field + intersection = ops.intersection(query_ops) + if intersection: + msg = 'Duplicate query contitions: ' + raise InvalidQueryError(msg + ', '.join(intersection)) + + query_ops.update(ops) + combined_query.update(copy.deepcopy(query)) + return combined_query + + +class QueryTreeTransformerVisitor(QNodeVisitor): + """Transforms the query tree in to a form that may be used with MongoDB. + """ + + def visit_combination(self, combination): + if combination.operation == combination.AND: + # MongoDB doesn't allow us to have too many $or operations in our + # queries, so the aim is to move the ORs up the tree to one + # 'master' $or. Firstly, we must find all the necessary parts (part + # of an AND combination or just standard Q object), and store them + # separately from the OR parts. + or_groups = [] + and_parts = [] + for node in combination.children: + if isinstance(node, QCombination): + if node.operation == node.OR: + # Any of the children in an $or component may cause + # the query to succeed + or_groups.append(node.children) + elif node.operation == node.AND: + and_parts.append(node) + elif isinstance(node, Q): + and_parts.append(node) + + # Now we combine the parts into a usable query. AND together all of + # the necessary parts. Then for each $or part, create a new query + # that ANDs the necessary part with the $or part. + clauses = [] + for or_group in itertools.product(*or_groups): + q_object = reduce(lambda a, b: a & b, and_parts, Q()) + q_object = reduce(lambda a, b: a & b, or_group, q_object) + clauses.append(q_object) + + # Finally, $or the generated clauses in to one query. Each of the + # clauses is sufficient for the query to succeed. + return reduce(lambda a, b: a | b, clauses, Q()) + + if combination.operation == combination.OR: + children = [] + # Crush any nested ORs in to this combination as MongoDB doesn't + # support nested $or operations + for node in combination.children: + if (isinstance(node, QCombination) and + node.operation == combination.OR): + children += node.children + else: + children.append(node) + combination.children = children + + return combination + + +class QueryCompilerVisitor(QNodeVisitor): + """Compiles the nodes in a query tree to a PyMongo-compatible query + dictionary. + """ + + def __init__(self, document): + self.document = document + + def visit_combination(self, combination): + if combination.operation == combination.OR: + return {'$or': combination.children} + elif combination.operation == combination.AND: + return self._mongo_query_conjunction(combination.children) + return combination + + def visit_query(self, query): + return QuerySet._transform_query(self.document, **query.query) + + def _mongo_query_conjunction(self, queries): + """Merges Mongo query dicts - effectively &ing them together. + """ + combined_query = {} + for query in queries: + for field, ops in query.items(): + if field not in combined_query: + combined_query[field] = ops + else: + # The field is already present in the query the only way + # we can merge is if both the existing value and the new + # value are operation dicts, reject anything else + if (not isinstance(combined_query[field], dict) or + not isinstance(ops, dict)): + message = 'Conflicting values for ' + field + raise InvalidQueryError(message) + + current_ops = set(combined_query[field].keys()) + new_ops = set(ops.keys()) + # Make sure that the same operation isn't applied more than + # once to a single field + intersection = current_ops.intersection(new_ops) + if intersection: + msg = 'Duplicate query contitions: ' + raise InvalidQueryError(msg + ', '.join(intersection)) + + # Right! We've got two non-overlapping dicts of operations! + combined_query[field].update(copy.deepcopy(ops)) + return combined_query + + +class QNode(object): + """Base class for nodes in query trees. + """ + + AND = 0 + OR = 1 + + def to_query(self, document): + query = self.accept(SimplificationVisitor()) + query = query.accept(QueryTreeTransformerVisitor()) + query = query.accept(QueryCompilerVisitor(document)) + return query + + def accept(self, visitor): + raise NotImplementedError + + def _combine(self, other, operation): + """Combine this node with another node into a QCombination object. + """ + if other.empty: + return self + + if self.empty: + return other + + return QCombination(operation, [self, other]) + + @property + def empty(self): + return False + + def __or__(self, other): + return self._combine(other, self.OR) + + def __and__(self, other): + return self._combine(other, self.AND) + + +class QCombination(QNode): + """Represents the combination of several conditions by a given logical + operator. + """ + + def __init__(self, operation, children): + self.operation = operation + self.children = [] + for node in children: + # If the child is a combination of the same type, we can merge its + # children directly into this combinations children + if isinstance(node, QCombination) and node.operation == operation: + self.children += node.children + else: + self.children.append(node) + + def accept(self, visitor): + for i in range(len(self.children)): + self.children[i] = self.children[i].accept(visitor) + + return visitor.visit_combination(self) + + @property + def empty(self): + return not bool(self.children) + + +class Q(QNode): + """A simple query object, used in a query tree to build up more complex + query structures. + """ + + def __init__(self, **query): + self.query = query + + def accept(self, visitor): + return visitor.visit_query(self) + + @property + def empty(self): + return not bool(self.query) + + +class OldQ(object): OR = '||' AND = '&&' @@ -162,7 +395,9 @@ class QuerySet(object): self._document = document self._collection_obj = collection self._accessed_collection = False - self._query = {} + self._mongo_query = None + self._query_obj = Q() + self._initial_query = {} self._where_clause = None self._loaded_fields = [] self._ordering = [] @@ -170,11 +405,18 @@ class QuerySet(object): # If inheritance is allowed, only return instances and instances of # subclasses of the class being used if document._meta.get('allow_inheritance'): - self._query = {'_types': self._document._class_name} + self._initial_query = {'_types': self._document._class_name} self._cursor_obj = None self._limit = None self._skip = None + @property + def _query(self): + if self._mongo_query is None: + self._mongo_query = self._query_obj.to_query(self._document) + self._mongo_query.update(self._initial_query) + return self._mongo_query + def ensure_index(self, key_or_list, drop_dups=False, background=False, **kwargs): """Ensure that the given indexes are in place. @@ -233,10 +475,14 @@ class QuerySet(object): objects, only the last one will be used :param query: Django-style query keyword arguments """ + #if q_obj: + #self._where_clause = q_obj.as_js(self._document) + query = Q(**query) if q_obj: - self._where_clause = q_obj.as_js(self._document) - query = QuerySet._transform_query(_doc_cls=self._document, **query) - self._query.update(query) + query &= q_obj + self._query_obj &= query + self._mongo_query = None + self._cursor_obj = None return self def filter(self, *q_objs, **query): @@ -338,7 +584,7 @@ class QuerySet(object): """Transform a query from Django-style format to Mongo format. """ operators = ['ne', 'gt', 'gte', 'lt', 'lte', 'in', 'nin', 'mod', - 'all', 'size', 'exists'] + 'all', 'size', 'exists', 'not'] geo_operators = ['within_distance', 'within_box', 'near'] match_operators = ['contains', 'icontains', 'startswith', 'istartswith', 'endswith', 'iendswith', @@ -354,6 +600,11 @@ class QuerySet(object): if parts[-1] in operators + match_operators + geo_operators: op = parts.pop() + negate = False + if parts[-1] == 'not': + parts.pop() + negate = True + if _doc_cls: # Switch field names to proper names [set in Field(name='foo')] fields = QuerySet._lookup_field(_doc_cls, parts) @@ -361,7 +612,7 @@ class QuerySet(object): # Convert value to proper value field = fields[-1] - singular_ops = [None, 'ne', 'gt', 'gte', 'lt', 'lte'] + singular_ops = [None, 'ne', 'gt', 'gte', 'lt', 'lte', 'not'] singular_ops += match_operators if op in singular_ops: value = field.prepare_query_value(op, value) @@ -386,6 +637,9 @@ class QuerySet(object): "been implemented" % op) elif op not in match_operators: value = {'$' + op: value} + + if negate: + value = {'$not': value} for i, part in indices: parts.insert(i, part) diff --git a/tests/queryset.py b/tests/queryset.py index 2271c366..32d1902e 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -53,9 +53,6 @@ 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) @@ -156,7 +153,8 @@ class QuerySetTest(unittest.TestCase): # Retrieve the first person from the database self.assertRaises(MultipleObjectsReturned, self.Person.objects.get) - self.assertRaises(self.Person.MultipleObjectsReturned, self.Person.objects.get) + self.assertRaises(self.Person.MultipleObjectsReturned, + self.Person.objects.get) # Use a query to filter the people found to just person2 person = self.Person.objects.get(age=30) @@ -234,7 +232,8 @@ class QuerySetTest(unittest.TestCase): self.assertEqual(created, False) # Try retrieving when no objects exists - new doc should be created - person, created = self.Person.objects.get_or_create(age=50, defaults={'name': 'User C'}) + kwargs = dict(age=50, defaults={'name': 'User C'}) + person, created = self.Person.objects.get_or_create(**kwargs) self.assertEqual(created, True) person = self.Person.objects.get(age=50) @@ -337,6 +336,18 @@ class QuerySetTest(unittest.TestCase): obj = self.Person.objects(Q(name__icontains='[.\'Geek')).first() self.assertEqual(obj, person) + def test_not(self): + """Ensure that the __not operator works as expected. + """ + alice = self.Person(name='Alice', age=25) + alice.save() + + obj = self.Person.objects(name__iexact='alice').first() + self.assertEqual(obj, alice) + + obj = self.Person.objects(name__not__iexact='alice').first() + self.assertEqual(obj, None) + def test_filter_chaining(self): """Ensure filters can be chained together. """ @@ -546,9 +557,10 @@ class QuerySetTest(unittest.TestCase): obj = self.Person.objects(Q(name=re.compile('^gui', re.I))).first() self.assertEqual(obj, person) - obj = self.Person.objects(Q(name__ne=re.compile('^bob'))).first() + obj = self.Person.objects(Q(name__not=re.compile('^bob'))).first() self.assertEqual(obj, person) - obj = self.Person.objects(Q(name__ne=re.compile('^Gui'))).first() + + obj = self.Person.objects(Q(name__not=re.compile('^Gui'))).first() self.assertEqual(obj, None) def test_q_lists(self): @@ -1346,43 +1358,6 @@ class QuerySetTest(unittest.TestCase): class QTest(unittest.TestCase): - def test_or_and(self): - """Ensure that Q objects may be combined correctly. - """ - 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 instanceof Array) && ' - 'this.name.indexOf(i0f0) != -1) || 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 instanceof Array) && ' - 'this.name.indexOf(i0f1) != -1) || 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) - def test_empty_q(self): """Ensure that empty Q objects won't hurt. """ @@ -1392,11 +1367,15 @@ class QTest(unittest.TestCase): q4 = Q(name='test') q5 = Q() - query = ['(', {'age__gte': 18}, '||', {'name': 'test'}, ')'] - self.assertEqual((q1 | q2 | q3 | q4 | q5).query, query) + class Person(Document): + name = StringField() + age = IntField() - query = ['(', {'age__gte': 18}, '&&', {'name': 'test'}, ')'] - self.assertEqual((q1 & q2 & q3 & q4 & q5).query, query) + query = {'$or': [{'age': {'$gte': 18}}, {'name': 'test'}]} + self.assertEqual((q1 | q2 | q3 | q4 | q5).to_query(Person), query) + + query = {'age': {'$gte': 18}, 'name': 'test'} + self.assertEqual((q1 & q2 & q3 & q4 & q5).to_query(Person), query) def test_q_with_dbref(self): """Ensure Q objects handle DBRefs correctly""" @@ -1415,5 +1394,107 @@ class QTest(unittest.TestCase): self.assertEqual(Post.objects.filter(Q(created_user=user)).count(), 1) +class NewQTest(unittest.TestCase): + + def test_and_combination(self): + """Ensure that Q-objects correctly AND together. + """ + class TestDoc(Document): + x = IntField() + y = StringField() + + # Check than an error is raised when conflicting queries are anded + def invalid_combination(): + query = Q(x__lt=7) & Q(x__lt=3) + query.to_query(TestDoc) + self.assertRaises(InvalidQueryError, invalid_combination) + + # Check normal cases work without an error + query = Q(x__lt=7) & Q(x__gt=3) + + q1 = Q(x__lt=7) + q2 = Q(x__gt=3) + query = (q1 & q2).to_query(TestDoc) + self.assertEqual(query, {'x': {'$lt': 7, '$gt': 3}}) + + # More complex nested example + query = Q(x__lt=100) & Q(y__ne='NotMyString') + query &= Q(y__in=['a', 'b', 'c']) & Q(x__gt=-100) + mongo_query = { + 'x': {'$lt': 100, '$gt': -100}, + 'y': {'$ne': 'NotMyString', '$in': ['a', 'b', 'c']}, + } + self.assertEqual(query.to_query(TestDoc), mongo_query) + + def test_or_combination(self): + """Ensure that Q-objects correctly OR together. + """ + class TestDoc(Document): + x = IntField() + + q1 = Q(x__lt=3) + q2 = Q(x__gt=7) + query = (q1 | q2).to_query(TestDoc) + self.assertEqual(query, { + '$or': [ + {'x': {'$lt': 3}}, + {'x': {'$gt': 7}}, + ] + }) + + def test_and_or_combination(self): + """Ensure that Q-objects handle ANDing ORed components. + """ + class TestDoc(Document): + x = IntField() + y = BooleanField() + + query = (Q(x__gt=0) | Q(x__exists=False)) + query &= Q(x__lt=100) + self.assertEqual(query.to_query(TestDoc), { + '$or': [ + {'x': {'$lt': 100, '$gt': 0}}, + {'x': {'$lt': 100, '$exists': False}}, + ] + }) + + q1 = (Q(x__gt=0) | Q(x__exists=False)) + q2 = (Q(x__lt=100) | Q(y=True)) + query = (q1 & q2).to_query(TestDoc) + + self.assertEqual(['$or'], query.keys()) + conditions = [ + {'x': {'$lt': 100, '$gt': 0}}, + {'x': {'$lt': 100, '$exists': False}}, + {'x': {'$gt': 0}, 'y': True}, + {'x': {'$exists': False}, 'y': True}, + ] + self.assertEqual(len(conditions), len(query['$or'])) + for condition in conditions: + self.assertTrue(condition in query['$or']) + + def test_or_and_or_combination(self): + """Ensure that Q-objects handle ORing ANDed ORed components. :) + """ + class TestDoc(Document): + x = IntField() + y = BooleanField() + + q1 = (Q(x__gt=0) & (Q(y=True) | Q(y__exists=False))) + q2 = (Q(x__lt=100) & (Q(y=False) | Q(y__exists=False))) + query = (q1 | q2).to_query(TestDoc) + + self.assertEqual(['$or'], query.keys()) + conditions = [ + {'x': {'$gt': 0}, 'y': True}, + {'x': {'$gt': 0}, 'y': {'$exists': False}}, + {'x': {'$lt': 100}, 'y':False}, + {'x': {'$lt': 100}, 'y': {'$exists': False}}, + ] + self.assertEqual(len(conditions), len(query['$or'])) + for condition in conditions: + self.assertTrue(condition in query['$or']) + + if __name__ == '__main__': unittest.main()