Experimental support for custom JOIN conditions

Author: Photonios

psqlextra/datastructures.py

@@ -0,0 +1,76 @@
+from typing import List, Tuple, Any
+
+from django.db.models.sql.datastructures import Join
+
+
+class ConditionalJoin(Join):
+    """A custom JOIN statement that allows attaching
+    extra conditions."""
+
+    def __init__(self, *args, **kwargs):
+        """Initializes a new instance of :see:ConditionalJoin."""
+
+        super().__init__(*args, **kwargs)
+        self.join_type = 'LEFT OUTER JOIN'
+        self.extra_conditions = []
+
+    def add_condition(self, field, value: Any) -> None:
+        """Adds an extra condition to this join.
+
+        Arguments:
+            field:
+                The field that the condition will apply to.
+
+            value:
+                The value to compare.
+        """
+
+        self.extra_conditions.append((field, value))
+
+    def as_sql(self, compiler, connection) -> Tuple[str, List[Any]]:
+        """Compiles this JOIN into a SQL string."""
+
+        sql, params = super().as_sql(compiler, connection)
+        qn = compiler.quote_name_unless_alias
+
+        # generate the extra conditions
+        extra_conditions = ' AND '.join([
+            '{}.{} = %s'.format(
+                qn(self.table_name),
+                qn(field.column)
+            )
+            for field, value in self.extra_conditions
+        ])
+
+        # add to the existing params, so the connector will
+        # actually nicely format the value for us
+        for _, value in self.extra_conditions:
+            params.append(value)
+
+        # rewrite the sql to include the extra conditions
+        rewritten_sql = sql.replace(')', ' AND {})'.format(extra_conditions))
+        return rewritten_sql, params
+
+    @classmethod
+    def from_join(cls, join: Join) -> 'ConditionalJoin':
+        """Creates a new :see:ConditionalJoin from the
+        specified :see:Join object.
+
+        Arguments:
+            join:
+                The :see:Join object to create the
+                :see:ConditionalJoin object from.
+
+        Returns:
+            A :see:ConditionalJoin object created from
+            the :see:Join object.
+        """
+
+        return cls(
+            join.table_name,
+            join.parent_alias,
+            join.table_alias,
+            join.join_type,
+            join.join_field,
+            join.nullable
+        )

psqlextra/expressions.py

@@ -87,6 +87,7 @@ def resolve_expression(self, *args, **kwargs) -> HStoreColumn:
         )
         return expression
 
+
 class NonGroupableFunc(expressions.Func):
     """A version of Django's :see:Func expression that
     is _never_ included in the GROUP BY clause."""

psqlextra/manager.py

@@ -26,6 +26,16 @@ def __init__(self, model=None, query=None, using=None, hints=None):
         self.conflict_target = None
         self.conflict_action = None
 
+    def join(self, **conditions):
+        """Adds extra conditions to existing joins.
+
+        WARNING: This is an extremely experimental feature.
+                 DO NOT USE unless you know what you're doing.
+        """
+
+        self.query.add_join_conditions(conditions)
+        return self
+
     def update(self, **fields):
         """Updates all rows that match the filter."""
 
@@ -35,7 +45,7 @@ def update(self, **fields):
         query._annotations = None
         query.add_update_values(fields)
 
-        # build the compiler for form the query
+        # build the compiler for for the query
         connection = django.db.connections[self.db]
         compiler = PostgresReturningUpdateCompiler(query, connection, self.db)
 

psqlextra/query.py

@@ -1,13 +1,14 @@
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Dict, Any
 from enum import Enum
 
 from django.db import models
 from django.db.models import sql
-
 from django.db.models.constants import LOOKUP_SEP
+from django.core.exceptions import SuspiciousOperation
 
 from .fields import HStoreField
 from .expressions import HStoreColumn
+from .datastructures import ConditionalJoin
 
 
 class ConflictAction(Enum):
@@ -18,6 +19,43 @@ class ConflictAction(Enum):
 
 
 class PostgresQuery(sql.Query):
+    def add_join_conditions(self, conditions: Dict[str, Any]) -> None:
+        """Adds an extra condition to an existing JOIN.
+
+        This allows you to for example do:
+
+            INNER JOIN othertable ON (mytable.id = othertable.other_id AND [extra conditions])
+
+        This does not work if nothing else in your query doesn't already generate the
+        initial join in the first place.
+        """
+
+        alias = self.get_initial_alias()
+        opts = self.get_meta()
+
+        for name, value in conditions.items():
+            parts = name.split(LOOKUP_SEP)
+            _, targets, _, joins, path = self.setup_joins(parts, opts, alias, allow_many=True)
+            self.trim_joins(targets, joins, path)
+
+            target_table = joins[-1]
+            field = targets[-1]
+            join = self.alias_map.get(target_table)
+
+            if not join:
+                raise SuspiciousOperation((
+                    'Cannot add an extra join condition for "%s", there\'s no'
+                    'existing join to add it to.'
+                ) % target_table)
+
+            # convert the Join object into a ConditionalJoin object, which
+            # allows us to add the extra condition
+            if not isinstance(join, ConditionalJoin):
+                self.alias_map[target_table] = ConditionalJoin.from_join(join)
+                join = self.alias_map[target_table]
+
+            join.add_condition(field, value)
+
     def add_fields(self, field_names: List[str], allow_m2m: bool=True) -> bool:
         """
         Adds the given (model) fields to the select set. The field names are

tests/test_hstore_field.py

@@ -50,6 +50,7 @@ def test_values():
 
     result = list(model.objects.values_list('title__en', 'title__ar'))[0]
 
+
 def test_annotate_ref():
     """Tests whether annotating using a :see:HStoreRef expression
     works correctly.
@@ -65,7 +66,7 @@ def test_annotate_ref():
     })
 
     fk = model_fk.objects.create(title={'en': 'english', 'ar': 'arabic'})
-    obj = model.objects.create(fk=fk)
+    model.objects.create(fk=fk)
 
     queryset = (
         model.objects