[2.1.x] Fixed #30672 -- Fixed crash of JSONField/HStoreField key transforms on expressions with params.

Regression in 4f5b58f5cd3c57fee9972ab074f8dc6895d8f387.

Thanks Florian Apolloner for the report and helping with tests.

Backport of 1f8382d34d54061eddc41df6994e20ee38c60907 from master.

2 files changed, 30 insertions, 1 deletions

diff --git a/tests/postgres_tests/test_hstore.py b/tests/postgres_tests/test_hstore.py
index f322231a4e..65f0b3616b 100644
--- a/tests/postgres_tests/test_hstore.py
+++ b/tests/postgres_tests/test_hstore.py
@@ -2,6 +2,7 @@ import json
 
 from django.core import checks, exceptions, serializers
 from django.db import connection
+from django.db.models.expressions import RawSQL
 from django.forms import Form
 from django.test.utils import (
     CaptureQueriesContext, isolate_apps, modify_settings,
@@ -13,6 +14,7 @@ from .models import HStoreModel, PostgreSQLModel
 try:
     from django.contrib.postgres import forms
     from django.contrib.postgres.fields import HStoreField
+    from django.contrib.postgres.fields.hstore import KeyTransform
     from django.contrib.postgres.validators import KeysValidator
 except ImportError:
     pass
@@ -133,6 +135,13 @@ class TestQuerying(HStoreTestCase):
             self.objs[:2]
         )
 
+    def test_key_transform_raw_expression(self):
+        expr = RawSQL('%s::hstore', ['x => b, y => c'])
+        self.assertSequenceEqual(
+            HStoreModel.objects.filter(field__a=KeyTransform('x', expr)),
+            self.objs[:2]
+        )
+
     def test_keys(self):
         self.assertSequenceEqual(
             HStoreModel.objects.filter(field__keys=['a']),
diff --git a/tests/postgres_tests/test_json.py b/tests/postgres_tests/test_json.py
index e58c7a2e7f..01215aa954 100644
--- a/tests/postgres_tests/test_json.py
+++ b/tests/postgres_tests/test_json.py
@@ -5,7 +5,9 @@ from decimal import Decimal
 from django.core import checks, exceptions, serializers
 from django.core.serializers.json import DjangoJSONEncoder
 from django.db import connection
-from django.db.models import Q
+from django.db.models import F, Q
+from django.db.models.expressions import RawSQL
+from django.db.models.functions import Cast
 from django.forms import CharField, Form, widgets
 from django.test.utils import CaptureQueriesContext, isolate_apps
 from django.utils.html import escape
@@ -16,6 +18,7 @@ from .models import JSONModel, PostgreSQLModel
 try:
     from django.contrib.postgres import forms
     from django.contrib.postgres.fields import JSONField
+    from django.contrib.postgres.fields.jsonb import KeyTransform
 except ImportError:
     pass
 
@@ -154,6 +157,23 @@ class TestQuerying(PostgreSQLTestCase):
         query = JSONModel.objects.filter(field__name__isnull=False).order_by('field__ord')
         self.assertSequenceEqual(query, [objs[4], objs[2], objs[3], objs[1], objs[0]])
 
+    def test_key_transform_raw_expression(self):
+        expr = RawSQL('%s::jsonb', ['{"x": "bar"}'])
+        self.assertSequenceEqual(
+            JSONModel.objects.filter(field__foo=KeyTransform('x', expr)),
+            [self.objs[-1]],
+        )
+
+    def test_key_transform_expression(self):
+        self.assertSequenceEqual(
+            JSONModel.objects.filter(field__d__0__isnull=False).annotate(
+                key=KeyTransform('d', 'field'),
+                chain=KeyTransform('0', 'key'),
+                expr=KeyTransform('0', Cast('key', JSONField())),
+            ).filter(chain=F('expr')),
+            [self.objs[8]],
+        )
+
     def test_deep_values(self):
         query = JSONModel.objects.values_list('field__k__l')
         self.assertSequenceEqual(