Skip to content

Commit

Permalink
fix incorrect GenericRelation joining
Browse files Browse the repository at this point in the history
By adding support for Field.get_extra_restriction().
  • Loading branch information
WaVEV authored and timgraham committed Sep 27, 2024
1 parent 5b92640 commit 78a8619
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 5 deletions.
4 changes: 1 addition & 3 deletions django_mongodb/features.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,9 +68,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
"aggregation.tests.AggregateTestCase.test_reverse_fkey_annotate",
"aggregation_regress.tests.AggregationTests.test_annotation_disjunction",
"aggregation_regress.tests.AggregationTests.test_decimal_aggregate_annotation_filter",
# Incorrect JOIN with GenericRelation gives incorrect results.
"aggregation_regress.tests.AggregationTests.test_aggregation_with_generic_reverse_relation",
"generic_relations.tests.GenericRelationsTests.test_queries_content_type_restriction",
# Wrong result for GenericRelation annotation.
"generic_relations_regress.tests.GenericRelationTests.test_annotate",
# subclasses of BaseDatabaseWrapper may require an is_usable() method
"backends.tests.BackendTestCase.test_is_usable_after_database_disconnects",
Expand Down
39 changes: 37 additions & 2 deletions django_mongodb/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

from django.core.exceptions import EmptyResultSet, FullResultSet
from django.db import DatabaseError, IntegrityError, NotSupportedError
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
Expand Down Expand Up @@ -105,6 +105,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, compiler.collection_name, rhs
Expand All @@ -113,8 +114,41 @@ 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))
# Handle any join conditions besides matching field pairs.
extra = self.join_field.get_extra_restriction(self.table_alias, self.parent_alias)
if extra:
columns = []
for expr in extra.leaves():
# Determine whether the column needs to be transformed or rerouted
# as part of the subquery.
for hand_side in ["lhs", "rhs"]:
hand_side_value = getattr(expr, hand_side, None)
if isinstance(hand_side_value, Col):
# If the column is not part of the joined table, add it to
# lhs_fields.
if hand_side_value.alias != self.table_name:
pos = len(lhs_fields)
lhs_fields.append(expr.lhs.as_mql(compiler, connection))
else:
pos = None
columns.append((hand_side_value, pos))
# Replace columns in the extra conditions with new column references
# based on their rerouted positions in the join pipeline.
replacements = {}
for col, parent_pos in columns:
column_target = Col(compiler.collection_name, expr.output_field.__class__())
if parent_pos is not None:
target_col = f"${parent_template}{parent_pos}"
column_target.target.db_column = target_col
column_target.target.set_attributes_from_name(target_col)
else:
column_target.target = col.target
replacements[col] = column_target
# Apply the transformed expressions in the extra condition.
extra_condition = [extra.replace_expressions(replacements).as_mql(compiler, connection)]
else:
extra_condition = []

parent_template = "parent__field__"
lookup_pipeline = [
{
"$lookup": {
Expand All @@ -140,6 +174,7 @@ def join(self, compiler, connection):
{"$eq": [f"$${parent_template}{i}", field]}
for i, field in enumerate(rhs_fields)
]
+ extra_condition
}
}
}
Expand Down

0 comments on commit 78a8619

Please sign in to comment.