Add support to join extras.

Author: WaVEV

django_mongodb/query.py

@@ -3,7 +3,7 @@
 
 from django.core.exceptions import EmptyResultSet, FullResultSet
 from django.db import DatabaseError, IntegrityError
-from django.db.models.expressions import Case, When
+from django.db.models.expressions import Case, Col, When
 from django.db.models.functions import Mod
 from django.db.models.lookups import Exact
 from django.db.models.sql.constants import INNER
@@ -97,6 +97,7 @@ def join(self, compiler, connection):
     lhs_fields = []
     rhs_fields = []
     # Add a join condition for each pair of joining fields.
+    parent_template = "parent__field__"
     for lhs, rhs in self.join_fields:
         lhs, rhs = connection.ops.prepare_join_on_clause(
             self.parent_alias, lhs, self.table_name, rhs
@@ -105,8 +106,28 @@ def join(self, compiler, connection):
         # In the lookup stage, the reference to this column doesn't include
         # the collection name.
         rhs_fields.append(rhs.as_mql(compiler, connection).replace(f"{self.table_name}.", "", 1))
+    extra = self.join_field.get_extra_restriction(self.table_alias, self.parent_alias)
+    if extra:
+        columns = []
+        for expr in extra.leaves():
+            if hasattr(expr, "lhs") and isinstance(expr.lhs, Col):
+                columns.append((expr.lhs, len(lhs_fields)))
+                lhs_fields.append(expr.lhs.as_mql(compiler, connection))
+            if hasattr(expr, "rhs") and isinstance(expr.rhs, Col):
+                columns.append((expr.rhs, None))
+        replacements = {}
+        for col, parent_pos in columns:
+            column_target = col.copy()
+            if column_target.alias == self.table_name:
+                column_target.alias = compiler.collection_name
+            else:
+                column_target.target.db_column = f"{parent_template}{parent_pos}"
+                column_target.target.set_attributes_from_name(f"{parent_template}{len(lhs_fields)}")
+            replacements[col] = column_target
+        extra_condition = [extra.replace_expressions(replacements).as_mql(compiler, connection)]
+    else:
+        extra_condition = []
 
-    parent_template = "parent__field__"
     lookup_pipeline = [
         {
             "$lookup": {
@@ -132,6 +153,7 @@ def join(self, compiler, connection):
                                     {"$eq": [f"$${parent_template}{i}", field]}
                                     for i, field in enumerate(rhs_fields)
                                 ]
+                                + extra_condition
                             }
                         }
                     }