Support lookup from related collections (#1)

Add support to select_related.

Author: WaVEV

.github/workflows/test-python.yml

@@ -93,6 +93,7 @@ jobs:
           ordering
           or_lookups
           queries.tests.Ticket12807Tests.test_ticket_12807
+          select_related
           sessions_tests
           timezones
           update

README.md

@@ -114,7 +114,6 @@ Migrations for 'admin':
   - `datetimes()`
   - `distinct()`
   - `extra()`
-  - `select_related()`
 
 - `Subquery`, `Exists`, and using a `QuerySet` in `QuerySet.annotate()` aren't
   supported.

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

@@ -44,16 +44,33 @@ def results_iter(
         """
         columns = self.get_columns()
 
+        related_columns = []
         if results is None:
             # QuerySet.values() or values_list()
             try:
                 results = self.build_query(columns).fetch()
             except EmptyResultSet:
                 results = []
+        else:
+            index = len(columns)
+            while index < self.col_count:
+                foreign_columns = []
+                foreign_relation = self.select[index][0].alias
+                while index < self.col_count and foreign_relation == self.select[index][0].alias:
+                    foreign_columns.append(self.select[index][0])
+                    index += 1
+                related_columns.append(
+                    (
+                        foreign_relation,
+                        [(column.target.column, column) for column in foreign_columns],
+                    )
+                )
 
         converters = self.get_converters(columns)
         for entity in results:
-            yield self._make_result(entity, columns, converters, tuple_expected=tuple_expected)
+            yield self._make_result(
+                entity, columns, related_columns, converters, tuple_expected=tuple_expected
+            )
 
     def has_results(self):
         return bool(self.get_count(check_exists=True))
@@ -72,13 +89,22 @@ def get_converters(self, expressions):
                 converters[name] = backend_converters + field_converters
         return converters
 
-    def _make_result(self, entity, columns, converters, tuple_expected=False):
+    def _make_result(self, entity, columns, related_columns, converters, tuple_expected=False):
         """
         Decode values for the given fields from the database entity.
 
         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:
             field = col.field
@@ -90,8 +116,6 @@ def _make_result(self, entity, columns, converters, tuple_expected=False):
                 for converter in converters.get(name, ()):
                     value = converter(value, col, self.connection)
             result.append(value)
-        if tuple_expected:
-            result = tuple(result)
         return result
 
     def check_query(self):
@@ -111,9 +135,11 @@ def check_query(self):
         if self.query.extra:
             raise NotSupportedError("QuerySet.extra() is not supported on MongoDB.")
         if self.query.select_related:
-            raise NotSupportedError("QuerySet.select_related() is not supported on MongoDB.")
+            pass
+            # raise NotSupportedError("QuerySet.select_related() is not supported on MongoDB.")
         if len([a for a in self.query.alias_map if self.query.alias_refcount[a]]) > 1:
-            raise NotSupportedError("Queries with multiple tables are not supported on MongoDB.")
+            pass
+            # raise NotSupportedError("Queries with multiple tables are not supported on MongoDB.")
         if any(
             isinstance(a, Aggregate) and not isinstance(a, Count)
             for a in self.query.annotations.values()
@@ -145,8 +171,8 @@ def get_count(self, check_exists=False):
     def build_query(self, columns=None):
         """Check if the query is supported and prepare a MongoQuery."""
         self.check_query()
-        self.setup_query()
         query = self.query_class(self, columns)
+        query.mongo_lookups = self.get_lookup_clauses()
         try:
             query.mongo_query = {"$expr": self.query.where.as_mql(self, self.connection)}
         except FullResultSet:
@@ -212,8 +238,36 @@ def _get_ordering(self):
             field_ordering.append((opts.get_field(name), ascending))
         return field_ordering
 
+    @property
+    def collection_name(self):
+        return self.query.get_meta().db_table
+
     def get_collection(self):
-        return self.connection.get_collection(self.query.get_meta().db_table)
+        return self.connection.get_collection(self.collection_name)
+
+    def get_lookup_clauses(self):
+        result = []
+        for alias in tuple(self.query.alias_map):
+            if not self.query.alias_refcount[alias] or self.collection_name == alias:
+                continue
+
+            from_clause = self.query.alias_map[alias]
+            clause_mql = from_clause.as_mql(self, self.connection)
+            result += clause_mql
+
+        """
+        for t in self.query.extra_tables:
+            alias, _ = self.query.table_alias(t)
+            # Only add the alias if it's not already present (the table_alias()
+            # call increments the refcount, so an alias refcount of one means
+            # this is the only reference).
+            if (
+                alias not in self.query.alias_map
+                or self.query.alias_refcount[alias] == 1
+            ):
+                result.append(", %s" % self.quote_name_unless_alias(alias))
+        """
+        return result
 
 
 class SQLInsertCompiler(SQLCompiler):

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

@@ -43,7 +43,8 @@ def case(self, compiler, connection):
 
 
 def col(self, compiler, connection):  # noqa: ARG001
-    return f"${self.target.column}"
+    prefix = f"{self.alias}." if self.alias != compiler.collection_name else ""
+    return f"${prefix}{self.target.column}"
 
 
 def combined_expression(self, compiler, connection):

django_mongodb/features.py

@@ -297,6 +297,7 @@ def django_test_expected_failures(self):
             "ordering.tests.OrderingTests.test_extra_ordering",
             "ordering.tests.OrderingTests.test_extra_ordering_quoting",
             "ordering.tests.OrderingTests.test_extra_ordering_with_table_name",
+            "select_related.tests.SelectRelatedTests.test_select_related_with_extra",
         },
         "Queries with multiple tables are not supported.": {
             "annotations.tests.AliasTests.test_alias_default_alias_expression",

django_mongodb/operations.py

@@ -7,7 +7,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
 
@@ -159,6 +159,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

@@ -40,8 +40,13 @@ 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.mongo_query = getattr(compiler.query, "raw_query", {})
+        # maybe I have to create a new object or named tuple.
+        # it will save lookups, some filters (in case of inner) and project to rename field
+        # don't know if the rename is needed
+        self.mongo_lookups = None
 
     def __repr__(self):
         return f"<MongoQuery: {self.mongo_query!r} ORDER {self.ordering!r}>"
@@ -106,11 +111,22 @@ def get_cursor(self):
                 # If name != column, then this is an annotatation referencing
                 # another column.
                 fields[name] = 1 if name == column else f"${column}"
+
+        # 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
+
         pipeline = []
+        if self.mongo_lookups:
+            lookups = self.mongo_lookups
+            pipeline.extend(lookups)
         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: