diff --git a/django_mongodb_backend/query.py b/django_mongodb_backend/query.py index c743ca8bf..1e9a30630 100644 --- a/django_mongodb_backend/query.py +++ b/django_mongodb_backend/query.py @@ -4,6 +4,7 @@ from django.core.exceptions import EmptyResultSet, FullResultSet from django.db import DatabaseError, IntegrityError, NotSupportedError from django.db.models.expressions import Case, Col, When +from django.db.models.fields.related import ForeignKey from django.db.models.functions import Mod from django.db.models.lookups import Exact from django.db.models.sql.constants import INNER @@ -181,14 +182,23 @@ def _get_reroot_replacements(expression): lhs_fields = [] rhs_fields = [] # Add a join condition for each pair of joining fields. + local_field = foreign_field = None for lhs, rhs in self.join_fields: - lhs, rhs = connection.ops.prepare_join_on_clause( + lhs_prepared, rhs_prepared = connection.ops.prepare_join_on_clause( self.parent_alias, lhs, compiler.collection_name, rhs ) - lhs_fields.append(lhs.as_mql(compiler, connection, as_expr=True)) - # In the lookup stage, the reference to this column doesn't include the - # collection name. - rhs_fields.append(rhs.as_mql(compiler, connection, as_expr=True)) + if ( + (isinstance(lhs, ForeignKey) or isinstance(rhs, ForeignKey)) + and lhs_prepared.is_simple_column + and rhs_prepared.is_simple_column + ): + local_field = lhs_prepared.as_mql(compiler, connection) + foreign_field = rhs_prepared.as_mql(compiler, connection) + else: + lhs_fields.append(lhs_prepared.as_mql(compiler, connection, as_expr=True)) + # In the lookup stage, the reference to this column doesn't include the + # collection name. + rhs_fields.append(rhs_prepared.as_mql(compiler, connection, as_expr=True)) # Handle any join conditions besides matching field pairs. extra = self.join_field.get_extra_restriction(self.table_alias, self.parent_alias) extra_conditions = [] @@ -218,32 +228,47 @@ def _get_reroot_replacements(expression): # self.table_name.field2 = parent_table.field2 # AND # ... - condition = { - "$expr": { - "$and": [ - {"$eq": [f"$${parent_template}{i}", field]} for i, field in enumerate(rhs_fields) - ] - } - } + all_conditions = [] + if rhs_fields: + all_conditions.append( + { + "$expr": { + "$and": [ + {"$eq": [f"$${parent_template}{i}", field]} + for i, field in enumerate(rhs_fields) + ] + } + } + ) if extra_conditions: - condition = {"$and": [condition, *extra_conditions]} - lookup_pipeline = [ - { - "$lookup": { - # The right-hand table to join. - "from": self.table_name, - # The pipeline variables to be matched in the pipeline's - # expression. - "let": { - f"{parent_template}{i}": parent_field - for i, parent_field in enumerate(lhs_fields) - }, - "pipeline": [{"$match": condition}], - # Rename the output as table_alias. - "as": self.table_alias, + all_conditions.extend(extra_conditions) + # Build matching pipeline + if len(all_conditions) == 0: + pipeline = [] + elif len(all_conditions) == 1: + pipeline = [{"$match": all_conditions[0]}] + else: + pipeline = [{"$match": {"$and": all_conditions}}] + + lookup = { + # The right-hand table to join. + "from": self.table_name, + "pipeline": pipeline, + # Rename the output as table_alias. + "as": self.table_alias, + } + if local_field and foreign_field: + lookup.update( + { + "localField": local_field, + "foreignField": foreign_field, } - }, - ] + ) + if lhs_fields: + lookup["let"] = { + f"{parent_template}{i}": parent_field for i, parent_field in enumerate(lhs_fields) + } + lookup_pipeline = [{"$lookup": lookup}] # To avoid missing data when using $unwind, an empty collection is added if # the join isn't an inner join. For inner joins, rows with empty arrays are # removed, as $unwind unrolls or unnests the array and removes the row if diff --git a/tests/queries_/test_mql.py b/tests/queries_/test_mql.py index e8837bf8a..2858b32be 100644 --- a/tests/queries_/test_mql.py +++ b/tests/queries_/test_mql.py @@ -22,23 +22,11 @@ def test_join(self): [ { "$lookup": { - "from": "queries__author", - "let": {"parent__field__0": "$author_id"}, - "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$_id"]}] - } - }, - {"name": "Bob"}, - ] - } - } - ], "as": "queries__author", + "foreignField": "_id", + "from": "queries__author", + "localField": "author_id", + "pipeline": [{"$match": {"name": "Bob"}}], } }, {"$unwind": "$queries__author"}, @@ -58,22 +46,10 @@ def test_filter_on_local_and_related_fields(self): { "$lookup": { "from": "queries__author", - "let": {"parent__field__0": "$author_id"}, - "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$_id"]}] - } - }, - {"name": "John"}, - ] - } - } - ], + "pipeline": [{"$match": {"name": "John"}}], "as": "queries__author", + "localField": "author_id", + "foreignField": "_id", } }, {"$unwind": "$queries__author"}, @@ -91,15 +67,10 @@ def test_or_mixing_local_and_related_fields_is_not_pushable(self): { "$lookup": { "from": "queries__author", - "let": {"parent__field__0": "$author_id"}, - "pipeline": [ - { - "$match": { - "$expr": {"$and": [{"$eq": ["$$parent__field__0", "$_id"]}]} - } - } - ], + "pipeline": [], "as": "queries__author", + "localField": "author_id", + "foreignField": "_id", } }, {"$unwind": "$queries__author"}, @@ -121,27 +92,19 @@ def test_filter_on_self_join_fields(self): { "$lookup": { "from": "queries__tag", - "let": {"parent__field__0": "$parent_id"}, "pipeline": [ { "$match": { "$and": [ - { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$_id"]}] - } - }, - { - "$and": [ - {"group_id": ObjectId("6891ff7822e475eddc20f159")}, - {"name": "parent"}, - ] - }, + {"group_id": ObjectId("6891ff7822e475eddc20f159")}, + {"name": "parent"}, ] } } ], "as": "T2", + "localField": "parent_id", + "foreignField": "_id", } }, {"$unwind": "$T2"}, @@ -166,24 +129,10 @@ def test_filter_on_reverse_foreignkey_relation(self): { "$lookup": { "from": "queries__orderitem", - "let": {"parent__field__0": "$_id"}, - "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [ - {"$eq": ["$$parent__field__0", "$order_id"]} - ] - } - }, - {"status": ObjectId("6891ff7822e475eddc20f159")}, - ] - } - } - ], + "localField": "_id", + "foreignField": "order_id", "as": "queries__orderitem", + "pipeline": [{"$match": {"status": ObjectId("6891ff7822e475eddc20f159")}}], } }, {"$unwind": "$queries__orderitem"}, @@ -209,23 +158,9 @@ def test_filter_on_local_and_nested_join_fields(self): { "$lookup": { "from": "queries__orderitem", - "let": {"parent__field__0": "$_id"}, - "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [ - {"$eq": ["$$parent__field__0", "$order_id"]} - ] - } - }, - {"status": ObjectId("6891ff7822e475eddc20f159")}, - ] - } - } - ], + "localField": "_id", + "foreignField": "order_id", + "pipeline": [{"$match": {"status": ObjectId("6891ff7822e475eddc20f159")}}], "as": "queries__orderitem", } }, @@ -233,22 +168,10 @@ def test_filter_on_local_and_nested_join_fields(self): { "$lookup": { "from": "queries__order", - "let": {"parent__field__0": "$queries__orderitem.order_id"}, - "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$_id"]}] - } - }, - {"name": "My Order"}, - ] - } - } - ], + "pipeline": [{"$match": {"name": "My Order"}}], "as": "T3", + "localField": "queries__orderitem.order_id", + "foreignField": "_id", } }, {"$unwind": "$T3"}, @@ -275,16 +198,11 @@ def test_negated_related_filter_is_not_pushable(self): [ { "$lookup": { - "as": "queries__author", "from": "queries__author", - "let": {"parent__field__0": "$author_id"}, - "pipeline": [ - { - "$match": { - "$expr": {"$and": [{"$eq": ["$$parent__field__0", "$_id"]}]} - } - } - ], + "pipeline": [], + "as": "queries__author", + "localField": "author_id", + "foreignField": "_id", } }, {"$unwind": "$queries__author"}, @@ -315,15 +233,10 @@ def test_or_with_mixed_pushable_and_non_pushable_fields(self): { "$lookup": { "from": "queries__author", - "let": {"parent__field__0": "$author_id"}, - "pipeline": [ - { - "$match": { - "$expr": {"$and": [{"$eq": ["$$parent__field__0", "$_id"]}]} - } - } - ], + "pipeline": [], "as": "queries__author", + "localField": "author_id", + "foreignField": "_id", } }, {"$unwind": "$queries__author"}, @@ -340,25 +253,14 @@ def test_push_equality_between_parent_and_child_fields(self): [ { "$lookup": { - "as": "queries__orderitem", "from": "queries__orderitem", - "let": {"parent__field__0": "$_id", "parent__field__1": "$_id"}, + "let": {"parent__field__0": "$_id"}, "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [ - {"$eq": ["$$parent__field__0", "$order_id"]} - ] - } - }, - {"$expr": {"$eq": ["$status", "$$parent__field__1"]}}, - ] - } - } + {"$match": {"$expr": {"$eq": ["$status", "$$parent__field__0"]}}} ], + "as": "queries__orderitem", + "localField": "_id", + "foreignField": "order_id", } }, {"$unwind": "$queries__orderitem"}, @@ -380,39 +282,20 @@ def test_simple_related_filter_is_pushed(self): { "$lookup": { "from": "queries__library_readers", - "let": {"parent__field__0": "$_id"}, - "pipeline": [ - { - "$match": { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$library_id"]}] - } - } - } - ], + "pipeline": [], "as": "queries__library_readers", + "localField": "_id", + "foreignField": "library_id", } }, {"$unwind": "$queries__library_readers"}, { "$lookup": { "from": "queries__reader", - "let": {"parent__field__0": "$queries__library_readers.reader_id"}, - "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$_id"]}] - } - }, - {"name": "Alice"}, - ] - } - } - ], + "pipeline": [{"$match": {"name": "Alice"}}], "as": "queries__reader", + "localField": "queries__library_readers.reader_id", + "foreignField": "_id", } }, {"$unwind": "$queries__reader"}, @@ -421,10 +304,8 @@ def test_simple_related_filter_is_pushed(self): ) def test_subquery_join_is_pushed(self): - # TODO; isn't fully OPTIMIZED with self.assertNumQueries(1) as ctx: list(Library.objects.filter(~models.Q(readers__name="Alice"))) - self.assertAggregateQuery( ctx.captured_queries[0]["sql"], "queries__library", @@ -437,30 +318,11 @@ def test_subquery_join_is_pushed(self): "pipeline": [ { "$lookup": { - "from": "queries__reader", - "let": {"parent__field__0": "$reader_id"}, - "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [ - { - "$eq": [ - "$$parent__field__0", - "$_id", - ] - } - ] - } - }, - {"name": "Alice"}, - ] - } - } - ], "as": "U2", + "foreignField": "_id", + "from": "queries__reader", + "localField": "reader_id", + "pipeline": [{"$match": {"name": "Alice"}}], } }, {"$unwind": "$U2"}, @@ -481,6 +343,7 @@ def test_subquery_join_is_pushed(self): "$set": { "__subquery0": { "$cond": { + "else": {"$arrayElemAt": ["$__subquery0", 0]}, "if": { "$or": [ {"$eq": [{"$type": "$__subquery0"}, "missing"]}, @@ -488,7 +351,6 @@ def test_subquery_join_is_pushed(self): ] }, "then": {}, - "else": {"$arrayElemAt": ["$__subquery0", 0]}, } } } @@ -532,16 +394,9 @@ def test_filter_on_local_and_related_fields(self): { "$lookup": { "from": "queries__library_readers", - "let": {"parent__field__0": "$_id"}, - "pipeline": [ - { - "$match": { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$library_id"]}] - } - } - } - ], + "localField": "_id", + "foreignField": "library_id", + "pipeline": [], "as": "queries__library_readers", } }, @@ -549,22 +404,10 @@ def test_filter_on_local_and_related_fields(self): { "$lookup": { "from": "queries__reader", - "let": {"parent__field__0": "$queries__library_readers.reader_id"}, - "pipeline": [ - { - "$match": { - "$and": [ - { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$_id"]}] - } - }, - {"name": "Alice"}, - ] - } - } - ], + "pipeline": [{"$match": {"name": "Alice"}}], "as": "queries__reader", + "localField": "queries__library_readers.reader_id", + "foreignField": "_id", } }, {"$unwind": "$queries__reader"}, @@ -572,7 +415,7 @@ def test_filter_on_local_and_related_fields(self): ], ) - def test_or_on_local_fields_only(self): + def test_annotate_foreign_field(self): with self.assertNumQueries(1) as ctx: list( Library.objects.annotate(foreing_field=models.F("readers__name")).filter( @@ -586,16 +429,9 @@ def test_or_on_local_fields_only(self): { "$lookup": { "from": "queries__library_readers", - "let": {"parent__field__0": "$_id"}, - "pipeline": [ - { - "$match": { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$library_id"]}] - } - } - } - ], + "localField": "_id", + "foreignField": "library_id", + "pipeline": [], "as": "queries__library_readers", } }, @@ -624,15 +460,10 @@ def test_or_on_local_fields_only(self): { "$lookup": { "from": "queries__reader", - "let": {"parent__field__0": "$queries__library_readers.reader_id"}, - "pipeline": [ - { - "$match": { - "$expr": {"$and": [{"$eq": ["$$parent__field__0", "$_id"]}]} - } - } - ], + "pipeline": [], "as": "queries__reader", + "localField": "queries__library_readers.reader_id", + "foreignField": "_id", } }, { @@ -673,16 +504,9 @@ def test_or_with_mixed_pushable_and_non_pushable_fields(self): { "$lookup": { "from": "queries__library_readers", - "let": {"parent__field__0": "$_id"}, - "pipeline": [ - { - "$match": { - "$expr": { - "$and": [{"$eq": ["$$parent__field__0", "$library_id"]}] - } - } - } - ], + "localField": "_id", + "foreignField": "library_id", + "pipeline": [], "as": "queries__library_readers", } }, @@ -711,15 +535,10 @@ def test_or_with_mixed_pushable_and_non_pushable_fields(self): { "$lookup": { "from": "queries__reader", - "let": {"parent__field__0": "$queries__library_readers.reader_id"}, - "pipeline": [ - { - "$match": { - "$expr": {"$and": [{"$eq": ["$$parent__field__0", "$_id"]}]} - } - } - ], + "pipeline": [], "as": "queries__reader", + "localField": "queries__library_readers.reader_id", + "foreignField": "_id", } }, {