Refactor: lookup stages to use localField and foreignField when applicable

Author: WaVEV

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)
+            and isinstance(lhs_prepared, Col)
+            and isinstance(rhs_prepared, Col)
+        ):
+            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