Fixed #25534, Fixed #31639 -- Added support for transform references in expressions.

Thanks Mariusz Felisiak and Simon Charette for reviews.

2 files changed, 63 insertions, 4 deletions

diff --git a/tests/model_fields/models.py b/tests/model_fields/models.py
index a11eb0ba44..c8867834da 100644
--- a/tests/model_fields/models.py
+++ b/tests/model_fields/models.py
@@ -363,6 +363,14 @@ class NullableJSONModel(models.Model):
         required_db_features = {'supports_json_field'}
 
 
+class RelatedJSONModel(models.Model):
+    value = models.JSONField()
+    json_model = models.ForeignKey(NullableJSONModel, models.CASCADE)
+
+    class Meta:
+        required_db_features = {'supports_json_field'}
+
+
 class AllFieldsModel(models.Model):
     big_integer = models.BigIntegerField()
     binary = models.BinaryField()
diff --git a/tests/model_fields/test_jsonfield.py b/tests/model_fields/test_jsonfield.py
index c1bacbe750..321826814b 100644
--- a/tests/model_fields/test_jsonfield.py
+++ b/tests/model_fields/test_jsonfield.py
@@ -25,7 +25,9 @@ from django.test import (
 )
 from django.test.utils import CaptureQueriesContext
 
-from .models import CustomJSONDecoder, JSONModel, NullableJSONModel
+from .models import (
+    CustomJSONDecoder, JSONModel, NullableJSONModel, RelatedJSONModel,
+)
 
 
 @skipUnlessDBFeature('supports_json_field')
@@ -357,12 +359,11 @@ class TestQuerying(TestCase):
             operator.itemgetter('key', 'count'),
         )
 
-    @skipUnlessDBFeature('allows_group_by_lob')
     def test_ordering_grouping_by_count(self):
         qs = NullableJSONModel.objects.filter(
             value__isnull=False,
         ).values('value__d__0').annotate(count=Count('value__d__0')).order_by('count')
-        self.assertQuerysetEqual(qs, [1, 11], operator.itemgetter('count'))
+        self.assertQuerysetEqual(qs, [0, 1], operator.itemgetter('count'))
 
     def test_order_grouping_custom_decoder(self):
         NullableJSONModel.objects.create(value_custom={'a': 'b'})
@@ -400,6 +401,17 @@ class TestQuerying(TestCase):
             [self.objs[4]],
         )
 
+    def test_key_transform_annotation_expression(self):
+        obj = NullableJSONModel.objects.create(value={'d': ['e', 'e']})
+        self.assertSequenceEqual(
+            NullableJSONModel.objects.filter(value__d__0__isnull=False).annotate(
+                key=F('value__d'),
+                chain=F('key__0'),
+                expr=Cast('key', models.JSONField()),
+            ).filter(chain=F('expr__1')),
+            [obj],
+        )
+
     def test_nested_key_transform_expression(self):
         self.assertSequenceEqual(
             NullableJSONModel.objects.filter(value__d__0__isnull=False).annotate(
@@ -410,6 +422,19 @@ class TestQuerying(TestCase):
             [self.objs[4]],
         )
 
+    def test_nested_key_transform_annotation_expression(self):
+        obj = NullableJSONModel.objects.create(
+            value={'d': ['e', {'f': 'g'}, {'f': 'g'}]},
+        )
+        self.assertSequenceEqual(
+            NullableJSONModel.objects.filter(value__d__0__isnull=False).annotate(
+                key=F('value__d'),
+                chain=F('key__1__f'),
+                expr=Cast('key', models.JSONField()),
+            ).filter(chain=F('expr__2__f')),
+            [obj],
+        )
+
     def test_nested_key_transform_on_subquery(self):
         self.assertSequenceEqual(
             NullableJSONModel.objects.filter(value__d__0__isnull=False).annotate(
@@ -449,12 +474,15 @@ class TestQuerying(TestCase):
         tests = [
             (Q(value__baz__has_key='a'), self.objs[7]),
             (Q(value__has_key=KeyTransform('a', KeyTransform('baz', 'value'))), self.objs[7]),
+            (Q(value__has_key=F('value__baz__a')), self.objs[7]),
             (Q(value__has_key=KeyTransform('c', KeyTransform('baz', 'value'))), self.objs[7]),
+            (Q(value__has_key=F('value__baz__c')), self.objs[7]),
             (Q(value__d__1__has_key='f'), self.objs[4]),
             (
                 Q(value__has_key=KeyTransform('f', KeyTransform('1', KeyTransform('d', 'value')))),
                 self.objs[4],
-            )
+            ),
+            (Q(value__has_key=F('value__d__1__f')), self.objs[4]),
         ]
         for condition, expected in tests:
             with self.subTest(condition=condition):
@@ -469,6 +497,7 @@ class TestQuerying(TestCase):
             Q(value__1__has_key='b'),
             Q(value__has_key=KeyTransform('b', KeyTransform(1, 'value'))),
             Q(value__has_key=KeyTransform('b', KeyTransform('1', 'value'))),
+            Q(value__has_key=F('value__1__b')),
         ]
         for condition in tests:
             with self.subTest(condition=condition):
@@ -733,11 +762,13 @@ class TestQuerying(TestCase):
                 [KeyTransform('foo', KeyTransform('bax', 'value'))],
                 [self.objs[7]],
             ),
+            ('value__foo__in', [F('value__bax__foo')], [self.objs[7]]),
             (
                 'value__foo__in',
                 [KeyTransform('foo', KeyTransform('bax', 'value')), 'baz'],
                 [self.objs[7]],
             ),
+            ('value__foo__in', [F('value__bax__foo'), 'baz'], [self.objs[7]]),
             ('value__foo__in', ['bar', 'baz'], [self.objs[7]]),
             ('value__bar__in', [['foo', 'bar']], [self.objs[7]]),
             ('value__bar__in', [['foo', 'bar'], ['a']], [self.objs[7]]),
@@ -850,6 +881,7 @@ class TestQuerying(TestCase):
             ('value__d__contains', 'e'),
             ('value__d__contains', [{'f': 'g'}]),
             ('value__contains', KeyTransform('bax', 'value')),
+            ('value__contains', F('value__bax')),
             ('value__baz__contains', {'a': 'b'}),
             ('value__baz__contained_by', {'a': 'b', 'c': 'd', 'e': 'f'}),
             (
@@ -869,3 +901,22 @@ class TestQuerying(TestCase):
                 self.assertIs(NullableJSONModel.objects.filter(
                     **{lookup: value},
                 ).exists(), True)
+
+    def test_join_key_transform_annotation_expression(self):
+        related_obj = RelatedJSONModel.objects.create(
+            value={'d': ['f', 'e']},
+            json_model=self.objs[4],
+        )
+        RelatedJSONModel.objects.create(
+            value={'d': ['e', 'f']},
+            json_model=self.objs[4],
+        )
+        self.assertSequenceEqual(
+            RelatedJSONModel.objects.annotate(
+                key=F('value__d'),
+                related_key=F('json_model__value__d'),
+                chain=F('key__1'),
+                expr=Cast('key', models.JSONField()),
+            ).filter(chain=F('related_key__0')),
+            [related_obj],
+        )