Composite primary key fields cannot use __in lookup with explicit Subquery

[Jacob Walls]

In a recent PR, we ​surfaced some inconsistencies with the resolution of explicit Subquery() such that some lookups might fail:

tests/composite_pk/test_filter.py

@@ -437 +437 @@
         queryset = User.objects.filter(comments__in=subquery)
         self.assertSequenceEqual(queryset, (self.user_2,))
 
+    def test_explicit_subquery(self):
+        subquery = Subquery(User.objects.values("pk"))
+        self.assertEqual(User.objects.filter(pk__in=subquery).count(), 5)
+        self.assertEqual(Comment.objects.filter(user__in=subquery).count(), 5)
+
     def test_cannot_cast_pk(self):
         msg = "Cast expression does not support composite primary keys."
         with self.assertRaisesMessage(ValueError, msg):

gives

======================================================================
ERROR: test_explicit_subquery (composite_pk.test_filter.CompositePKFilterTests.test_explicit_subquery)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/.../django/tests/composite_pk/test_filter.py", line 418, in test_explicit_subquery
    self.assertEqual(Comment.objects.filter(user__in=subquery).count(), 5)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/query.py", line 603, in count
    return self.query.get_count(using=self.db)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/sql/query.py", line 644, in get_count
    return obj.get_aggregation(using, {"__count": Count("*")})["__count"]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/sql/query.py", line 626, in get_aggregation
    result = compiler.execute_sql(SINGLE)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/sql/compiler.py", line 1610, in execute_sql
    sql, params = self.as_sql()
                  ^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/sql/compiler.py", line 794, in as_sql
    self.compile(self.where) if self.where is not None else ("", [])
    ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/sql/compiler.py", line 577, in compile
    sql, params = node.as_sql(self, self.connection)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/sql/where.py", line 151, in as_sql
    sql, params = compiler.compile(child)
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/sql/compiler.py", line 577, in compile
    sql, params = node.as_sql(self, self.connection)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../django/django/db/models/fields/related_lookups.py", line 86, in as_sql
    SubqueryConstraint(
  File "/Users/.../django/django/db/models/sql/where.py", line 358, in __init__
    query_object.clear_ordering(clear_default=True)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'Subquery' object has no attribute 'clear_ordering'

[Simon Charette]

Well good news, I think the work on #36149 managed to make SubqueryConstraint completely irrelevant while solving this issue.

django/db/backends/mysql/compiler.py

@@ -1 +1 @@
 from django.core.exceptions import FieldError, FullResultSet
 from django.db.models.expressions import Col
-from django.db.models.sql import compiler
+from django.db.models.sql.compiler import SQLAggregateCompiler, SQLCompiler
+from django.db.models.sql.compiler import SQLDeleteCompiler as BaseSQLDeleteCompiler
+from django.db.models.sql.compiler import SQLInsertCompiler
+from django.db.models.sql.compiler import SQLUpdateCompiler as BaseSQLUpdateCompiler
 
+__all__ = [
+    "SQLAggregateCompiler",
+    "SQLCompiler",
+    "SQLDeleteCompiler",
+    "SQLInsertCompiler",
+    "SQLUpdateCompiler",
+]
 
-class SQLCompiler(compiler.SQLCompiler):
-    def as_subquery_condition(self, alias, columns, compiler):
-        qn = compiler.quote_name_unless_alias
-        qn2 = self.connection.ops.quote_name
-        sql, params = self.as_sql()
-        return (
-            "(%s) IN (%s)"
-            % (
-                ", ".join("%s.%s" % (qn(alias), qn2(column)) for column in columns),
-                sql,
-            ),
-            params,
-        )
-
-
-class SQLInsertCompiler(compiler.SQLInsertCompiler, SQLCompiler):
-    pass
 
-
-class SQLDeleteCompiler(compiler.SQLDeleteCompiler, SQLCompiler):
+class SQLDeleteCompiler(BaseSQLDeleteCompiler):
     def as_sql(self):
         # Prefer the non-standard DELETE FROM syntax over the SQL generated by
         # the SQLDeleteCompiler's default implementation when multiple tables
@@ -52 +44 @@
         return " ".join(result), tuple(params)
 
 
-class SQLUpdateCompiler(compiler.SQLUpdateCompiler, SQLCompiler):
+class SQLUpdateCompiler(BaseSQLUpdateCompiler):
     def as_sql(self):
         update_query, update_params = super().as_sql()
         # MySQL and MariaDB support UPDATE ... ORDER BY syntax.
@@ -78 +70 @@
                 # removed in .update() and cannot be resolved.
                 pass
         return update_query, update_params
-
-
-class SQLAggregateCompiler(compiler.SQLAggregateCompiler, SQLCompiler):
-    pass

django/db/models/fields/related_lookups.py

@@ -84 +84 @@
 
     def as_sql(self, compiler, connection):
         if isinstance(self.lhs, ColPairs):
-            from django.db.models.sql.where import SubqueryConstraint
-
             if self.rhs_is_direct_value():
                 values = [get_normalized_value(value, self.lhs) for value in self.rhs]
                 lookup = TupleIn(self.lhs, values)
-                return compiler.compile(lookup)
             else:
-                return compiler.compile(
-                    SubqueryConstraint(
-                        self.lhs.alias,
-                        [target.column for target in self.lhs.targets],
-                        [source.name for source in self.lhs.sources],
-                        self.rhs,
-                    ),
-                )
+                lookup = TupleIn(self.lhs, self.rhs)
+            return compiler.compile(lookup)
 
         return super().as_sql(compiler, connection)
 

django/db/models/fields/tuple_lookups.py

@@ -2 +2 @@
 
 from django.core.exceptions import EmptyResultSet
 from django.db.models import Field
-from django.db.models.expressions import ColPairs, Func, ResolvedOuterRef, Value
+from django.db.models.expressions import (
+    ColPairs,
+    Func,
+    ResolvedOuterRef,
+    Subquery,
+    Value,
+)
 from django.db.models.lookups import (
     Exact,
     GreaterThan,
@@ -301 +307 @@
             )
 
     def check_rhs_is_query(self):
-        if not isinstance(self.rhs, Query):
+        if not isinstance(self.rhs, (Query, Subquery)):
             lhs_str = self.get_lhs_str()
             rhs_cls = self.rhs.__class__.__name__
             raise ValueError(

django/db/models/sql/compiler.py

@@ -1661 +1661 @@
             return list(result)
         return result
 
-    def as_subquery_condition(self, alias, columns, compiler):
-        qn = compiler.quote_name_unless_alias
-        qn2 = self.connection.ops.quote_name
-        query = self.query.clone()
-
-        for index, select_col in enumerate(query.select):
-            lhs_sql, lhs_params = self.compile(select_col)
-            rhs = "%s.%s" % (qn(alias), qn2(columns[index]))
-            query.where.add(RawSQL("%s = %s" % (lhs_sql, rhs), lhs_params), AND)
-
-        sql, params = query.as_sql(compiler, self.connection)
-        return "EXISTS %s" % sql, params
-
     def explain_query(self):
         result = list(self.execute_sql())
         # Some backends return 1 item tuples with strings, and others return

django/db/models/sql/where.py

@@ -343 +343 @@
     def as_sql(self, compiler=None, connection=None):
         sqls = ["(%s)" % sql for sql in self.sqls]
         return " AND ".join(sqls), list(self.params or ())
-
-
-class SubqueryConstraint:
-    # Even if aggregates or windows would be used in a subquery,
-    # the outer query isn't interested about those.
-    contains_aggregate = False
-    contains_over_clause = False
-
-    def __init__(self, alias, columns, targets, query_object):
-        self.alias = alias
-        self.columns = columns
-        self.targets = targets
-        query_object.clear_ordering(clear_default=True)
-        self.query_object = query_object
-
-    def as_sql(self, compiler, connection):
-        query = self.query_object
-        query.set_values(self.targets)
-        query_compiler = query.get_compiler(connection=connection)
-        return query_compiler.as_subquery_condition(self.alias, self.columns, compiler)

tests/composite_pk/test_filter.py

@@ -442 +442 @@
         with self.assertRaisesMessage(ValueError, msg):
             Comment.objects.filter(text__gt=Cast(F("pk"), TextField())).count()
 
+    def test_explicit_subquery(self):
+        subquery = Subquery(User.objects.values("pk"))
+        self.assertEqual(User.objects.filter(pk__in=subquery).count(), 4)
+        self.assertEqual(Comment.objects.filter(user__in=subquery).count(), 5)
+
     def test_filter_case_when(self):
         msg = "When expression does not support composite primary keys."
         with self.assertRaisesMessage(ValueError, msg):

Getting rid of SuqueryConstraint and as_subquery_condition was a goal of #373 so that's a great win!