From 3fcc0e97895d1f2d2a1bea150f15d1f8aa35b7e9 Mon Sep 17 00:00:00 2001 From: Harry Marr Date: Mon, 4 Oct 2010 02:10:37 +0100 Subject: [PATCH] Combining OR nodes works, fixed other Q-object bugs --- mongoengine/queryset.py | 97 +++++++++++++++++++++++++++++++---------- tests/queryset.py | 62 ++++++++++++++++++++++---- 2 files changed, 126 insertions(+), 33 deletions(-) diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 9650fe1f..06b67ab5 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -4,6 +4,7 @@ import pprint import pymongo import re import copy +import itertools __all__ = ['queryset_manager', 'Q', 'InvalidQueryError', 'InvalidCollectionError'] @@ -49,8 +50,22 @@ class QNodeVisitor(object): """ 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, NewQ) for node in combination.children): + queries = [node.query for node in combination.children] + return NewQ(**self._query_conjunction(queries)) + return combination + def _query_conjunction(self, queries): - """Merges two query dicts - effectively &ing them together. + """Merges query dicts - effectively &ing them together. """ query_ops = set() combined_query = {} @@ -68,23 +83,6 @@ class QNodeVisitor(object): return combined_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: - return combination - - # The simplification only applies to 'simple' queries - if any(not isinstance(node, NewQ) for node in combination.children): - return combination - - queries = [node.query for node in combination.children] - return NewQ(**self._query_conjunction(queries)) - - class QueryTreeTransformerVisitor(QNodeVisitor): """Transforms the query tree in to a form that may be used with MongoDB. """ @@ -96,14 +94,14 @@ class QueryTreeTransformerVisitor(QNodeVisitor): # '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_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_parts += node.children + or_groups.append(node.children) elif node.operation == node.AND: and_parts.append(node) elif isinstance(node, NewQ): @@ -113,13 +111,27 @@ class QueryTreeTransformerVisitor(QNodeVisitor): # the necessary parts. Then for each $or part, create a new query # that ANDs the necessary part with the $or part. clauses = [] - for or_part in or_parts: + for or_group in itertools.product(*or_groups): q_object = reduce(lambda a, b: a & b, and_parts, NewQ()) - clauses.append(q_object & or_part) + 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, NewQ()) + + 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 @@ -135,12 +147,42 @@ class QueryCompilerVisitor(QNodeVisitor): if combination.operation == combination.OR: return {'$or': combination.children} elif combination.operation == combination.AND: - return self._query_conjunction(combination.children) + 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. @@ -187,7 +229,14 @@ class QCombination(QNode): def __init__(self, operation, children): self.operation = operation - self.children = children + 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)): diff --git a/tests/queryset.py b/tests/queryset.py index f83a8d43..0f8c9a92 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -1415,6 +1415,8 @@ class QTest(unittest.TestCase): class NewQTest(unittest.TestCase): def test_and_combination(self): + """Ensure that Q-objects correctly AND together. + """ class TestDoc(Document): x = IntField() y = StringField() @@ -1443,6 +1445,8 @@ class NewQTest(unittest.TestCase): 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() @@ -1457,18 +1461,58 @@ class NewQTest(unittest.TestCase): }) def test_and_or_combination(self): + """Ensure that Q-objects handle ANDing ORed components. + """ class TestDoc(Document): x = IntField() + y = BooleanField() + + query = (NewQ(x__gt=0) | NewQ(x__exists=False)) + query &= NewQ(x__lt=100) + self.assertEqual(query.to_query(TestDoc), { + '$or': [ + {'x': {'$lt': 100, '$gt': 0}}, + {'x': {'$lt': 100, '$exists': False}}, + ] + }) + + q1 = (NewQ(x__gt=0) | NewQ(x__exists=False)) + q2 = (NewQ(x__lt=100) | NewQ(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 = (NewQ(x__gt=0) & (NewQ(y=True) | NewQ(y__exists=False))) + q2 = (NewQ(x__lt=100) & (NewQ(y=False) | NewQ(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']) - query = NewQ(x__gt=0) | NewQ(x__exists=False) - query &= NewQ(x__lt=100) | NewQ(x__in=[100, 200, 3000]) - print query.to_query(TestDoc) -# self.assertEqual(query.to_query(TestDoc, { -# '$or': [ -# {'x': {'$lt': 3}}, -# {'x': {'$gt': 7}}, -# ] -# }) if __name__ == '__main__': unittest.main()