Support lookup from related collections (#1)

Add support to select_related.

Author: WaVEV

django_mongodb/__init__.py

@@ -6,12 +6,14 @@
 
 check_django_compatability()
 
+from .datastructures import register_datastructures  # noqa: E402
 from .expressions import register_expressions  # noqa: E402
 from .fields import register_fields  # noqa: E402
 from .functions import register_functions  # noqa: E402
 from .lookups import register_lookups  # noqa: E402
 from .query import register_nodes  # noqa: E402
 
+register_datastructures()
 register_expressions()
 register_fields()
 register_functions()

django_mongodb/compiler.py

@@ -73,6 +73,15 @@ def _make_result(self, entity, columns):
         The entity is assumed to be a dict using field database column
         names as keys.
         """
+        result = self._project_result(entity, columns, converters, tuple_expected)
+        # Related columns
+        for relation, columns in related_columns:
+            result += self._project_result(entity[relation], columns, converters, tuple_expected)
+        if tuple_expected:
+            result = tuple(result)
+        return result
+
+    def _project_result(self, entity, columns, converters, tuple_expected=False):
         result = []
         for name, col in columns:
             column_alias = getattr(col, "alias", None)

django_mongodb/datastructures.py

@@ -0,0 +1,76 @@
+from django.db.models.sql.constants import INNER
+from django.db.models.sql.datastructures import Join
+
+
+def join(self, compiler, connection):
+    lookups_pipeline = []
+    join_fields = self.join_fields or self.join_cols
+    lhs_fields = []
+    rhs_fields = []
+    for lhs, rhs in join_fields:
+        if isinstance(lhs, str):
+            lhs_mql = lhs
+            rhs_mql = rhs
+        else:
+            lhs, rhs = connection.ops.prepare_join_on_clause(
+                self.parent_alias, lhs, self.table_name, rhs
+            )
+            lhs_mql = lhs.as_mql(compiler, connection)
+            rhs_mql = rhs.as_mql(compiler, connection)
+            # replace prefix, in lookup stages the reference
+            # to this column is without the collection name.
+            rhs_mql = rhs_mql.replace(f"{self.table_name}.", "", 1)
+        lhs_fields.append(lhs_mql)
+        rhs_fields.append(rhs_mql)
+
+    # temp_table_name = f"{self.table_alias}__array"
+    parent_template = "parent__field__"
+    lookups_pipeline = [
+        {
+            "$lookup": {
+                "from": self.table_name,
+                "let": {
+                    f"{parent_template}{i}": parent_field
+                    for i, parent_field in enumerate(lhs_fields)
+                },
+                "pipeline": [
+                    {
+                        "$match": {
+                            "$expr": {
+                                "$and": [
+                                    {"$eq": [f"$${parent_template}{i}", field]}
+                                    for i, field in enumerate(rhs_fields)
+                                ]
+                            }
+                        }
+                    }
+                ],
+                "as": self.table_alias,
+            }
+        },
+    ]
+    if self.join_type != INNER:
+        lookups_pipeline.append(
+            {
+                "$project": {
+                    self.table_alias: {
+                        "$cond": {
+                            "if": {
+                                "$or": [
+                                    {"$eq": [{"$type": "$arrayField"}, "missing"]},
+                                    {"$eq": [{"$size": "$arrayField"}, 0]},
+                                ]
+                            },
+                            "then": [None],
+                            "else": "$arrayField",
+                        }
+                    }
+                }
+            }
+        )
+    lookups_pipeline.append({"$unwind": f"${self.table_alias}"})
+    return lookups_pipeline
+
+
+def register_datastructures():
+    Join.as_mql = join

django_mongodb/expressions.py

@@ -44,7 +44,6 @@ def case(self, compiler, connection):
 
 
 def col(self, compiler, connection):  # noqa: ARG001
-    # Add the column's collection's alias for columns in joined collections.
     prefix = f"{self.alias}." if self.alias != compiler.collection_name else ""
     return f"${prefix}{self.target.column}"
 

django_mongodb/operations.py

@@ -8,7 +8,7 @@
 from django.conf import settings
 from django.db import DataError
 from django.db.backends.base.operations import BaseDatabaseOperations
-from django.db.models.expressions import Combinable
+from django.db.models.expressions import Col, Combinable
 from django.utils import timezone
 from django.utils.regex_helper import _lazy_re_compile
 
@@ -176,6 +176,12 @@ def execute_sql_flush(self, tables):
             if not options.get("capped", False):
                 collection.drop()
 
+    def prepare_join_on_clause(self, lhs_table, lhs_field, rhs_table, rhs_field):
+        lhs_expr = Col(lhs_table, lhs_field)
+        rhs_expr = Col(rhs_table, rhs_field)
+
+        return lhs_expr, rhs_expr
+
     def prep_lookup_value(self, value, field, lookup):
         """
         Perform type-conversion on `value` before using as a filter parameter.

django_mongodb/query.py

@@ -45,6 +45,7 @@ def __init__(self, compiler, columns):
         self.columns = columns
         self._negated = False
         self.ordering = []
+        self.collection_name = self.compiler.collection_name
         self.collection = self.compiler.get_collection()
         self.collection_name = self.compiler.collection_name
         self.mongo_query = getattr(compiler.query, "raw_query", {})
@@ -123,10 +124,18 @@ def get_cursor(self):
         pipeline = []
         if self.lookup_pipeline:
             pipeline.extend(self.lookup_pipeline)
+
+        # add the subquery tables. if fields is defined
+        related_fields = {}
+        if fields:
+            for alias in self.query.alias_map:
+                if self.query.alias_refcount[alias] > 0 and self.collection_name != alias:
+                    related_fields[alias] = 1
+
         if self.mongo_query:
             pipeline.append({"$match": self.mongo_query})
         if fields:
-            pipeline.append({"$project": fields})
+            pipeline.append({"$project": {**fields, **related_fields}})
         if self.ordering:
             pipeline.append({"$sort": dict(self.ordering)})
         if self.query.low_mark > 0: