Fixed #35972 -- Fixed lookup crashes after subquery annotations.

3 files changed, 86 insertions, 4 deletions

diff --git a/tests/custom_lookups/tests.py b/tests/custom_lookups/tests.py
index 6728c49afd..db985240d7 100644
--- a/tests/custom_lookups/tests.py
+++ b/tests/custom_lookups/tests.py
@@ -5,6 +5,7 @@ from datetime import date, datetime
 from django.core.exceptions import FieldError
 from django.db import connection, models
 from django.db.models.fields.related_lookups import RelatedGreaterThan
+from django.db.models.functions import Lower
 from django.db.models.lookups import EndsWith, StartsWith
 from django.test import SimpleTestCase, TestCase, override_settings
 from django.test.utils import register_lookup
@@ -17,15 +18,15 @@ class Div3Lookup(models.Lookup):
     lookup_name = "div3"
 
     def as_sql(self, compiler, connection):
-        lhs, params = self.process_lhs(compiler, connection)
+        lhs, lhs_params = self.process_lhs(compiler, connection)
         rhs, rhs_params = self.process_rhs(compiler, connection)
-        params.extend(rhs_params)
+        params = (*lhs_params, *rhs_params)
         return "(%s) %%%% 3 = %s" % (lhs, rhs), params
 
     def as_oracle(self, compiler, connection):
-        lhs, params = self.process_lhs(compiler, connection)
+        lhs, lhs_params = self.process_lhs(compiler, connection)
         rhs, rhs_params = self.process_rhs(compiler, connection)
-        params.extend(rhs_params)
+        params = (*lhs_params, *rhs_params)
         return "mod(%s, 3) = %s" % (lhs, rhs), params
 
 
@@ -249,6 +250,39 @@ class LookupTests(TestCase):
             self.assertSequenceEqual(qs1, [a1])
             self.assertSequenceEqual(qs2, [a1])
 
+    def test_custom_lookup_with_subquery(self):
+        class NotEqual(models.Lookup):
+            lookup_name = "ne"
+
+            def as_sql(self, compiler, connection):
+                lhs, lhs_params = self.process_lhs(compiler, connection)
+                rhs, rhs_params = self.process_rhs(compiler, connection)
+                # Although combining via (*lhs_params, *rhs_params) would be
+                # more resilient, the "simple" way works too.
+                params = lhs_params + rhs_params
+                return "%s <> %s" % (lhs, rhs), params
+
+        author = Author.objects.create(name="Isabella")
+
+        with register_lookup(models.Field, NotEqual):
+            qs = Author.objects.annotate(
+                unknown_age=models.Subquery(
+                    Author.objects.filter(age__isnull=True)
+                    .order_by("name")
+                    .values("name")[:1]
+                )
+            ).filter(unknown_age__ne="Plato")
+            self.assertSequenceEqual(qs, [author])
+
+            qs = Author.objects.annotate(
+                unknown_age=Lower(
+                    Author.objects.filter(age__isnull=True)
+                    .order_by("name")
+                    .values("name")[:1]
+                )
+            ).filter(unknown_age__ne="plato")
+            self.assertSequenceEqual(qs, [author])
+
     def test_custom_exact_lookup_none_rhs(self):
         """
         __exact=None is transformed to __isnull=True if a custom lookup class
diff --git a/tests/expressions/tests.py b/tests/expressions/tests.py
index 5f61f65ac0..ac06c212ea 100644
--- a/tests/expressions/tests.py
+++ b/tests/expressions/tests.py
@@ -749,6 +749,24 @@ class BasicExpressionsTests(TestCase):
         )
         self.assertCountEqual(subquery_test2, [self.foobar_ltd])
 
+    def test_lookups_subquery(self):
+        smallest_company = Company.objects.order_by("num_employees").values("name")[:1]
+        for lookup in CharField.get_lookups():
+            if lookup == "isnull":
+                continue  # not allowed, rhs must be a literal boolean.
+            if (
+                lookup == "in"
+                and not connection.features.allow_sliced_subqueries_with_in
+            ):
+                continue
+            if lookup == "range":
+                rhs = (Subquery(smallest_company), Subquery(smallest_company))
+            else:
+                rhs = Subquery(smallest_company)
+            with self.subTest(lookup=lookup):
+                qs = Company.objects.filter(**{f"name__{lookup}": rhs})
+                self.assertGreater(len(qs), 0)
+
     def test_uuid_pk_subquery(self):
         u = UUIDPK.objects.create()
         UUID.objects.create(uuid_fk=u)
diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index 0bf16cb481..ec4b772c13 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -2340,6 +2340,36 @@ class OperationTests(OperationTestBase):
         pony = project_state.apps.get_model(app_label, "pony").objects.create(weight=1)
         self.assertEqual(pony.pink, 3)
 
+    @skipUnlessDBFeature("supports_expression_defaults")
+    def test_alter_field_add_database_default_func(self):
+        app_label = "test_alfladdf"
+        project_state = self.set_up_test_model(app_label)
+        operation = migrations.AlterField(
+            "Pony", "weight", models.FloatField(db_default=Pi())
+        )
+        new_state = project_state.clone()
+        operation.state_forwards(app_label, new_state)
+        old_weight = project_state.models[app_label, "pony"].fields["weight"]
+        self.assertIs(old_weight.default, models.NOT_PROVIDED)
+        self.assertIs(old_weight.db_default, models.NOT_PROVIDED)
+        new_weight = new_state.models[app_label, "pony"].fields["weight"]
+        self.assertIs(new_weight.default, models.NOT_PROVIDED)
+        self.assertIsInstance(new_weight.db_default, Pi)
+        pony = project_state.apps.get_model(app_label, "pony").objects.create(weight=1)
+        self.assertEqual(pony.weight, 1)
+        # Alter field.
+        with connection.schema_editor() as editor:
+            operation.database_forwards(app_label, editor, project_state, new_state)
+        pony = new_state.apps.get_model(app_label, "pony").objects.create()
+        if not connection.features.can_return_columns_from_insert:
+            pony.refresh_from_db()
+        self.assertAlmostEqual(pony.weight, math.pi)
+        # Reversal.
+        with connection.schema_editor() as editor:
+            operation.database_backwards(app_label, editor, new_state, project_state)
+        pony = project_state.apps.get_model(app_label, "pony").objects.create(weight=1)
+        self.assertEqual(pony.weight, 1)
+
     def test_alter_field_change_nullable_to_database_default_not_null(self):
         """
         The AlterField operation changing a null field to db_default.