From 59fe0b85414f5030499075d95cad3cd73789cfe8 Mon Sep 17 00:00:00 2001
From: Mariusz Felisiak <felisiak.mariusz@gmail.com>
Date: Tue, 13 Oct 2020 11:35:55 +0200
Subject: [PATCH] [3.1.x] Refs #32096 -- Fixed __in lookup crash against key
 transforms for JSONField.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Regression in 6789ded0a6ab797f0dcdfa6ad5d1cfa46e23abcd and
1251772cb83aa4106f526fe00738e51c0eb59122.

Thanks Simon Charette and Igor Jerosimić for the report.

Backport of 7e1e198494d4fc72cf6e153f9d24fe2493c17dc1 from master
---
 django/db/models/fields/json.py      | 37 ++++++++++++++--------------
 django/db/models/lookups.py          |  2 +-
 docs/releases/3.1.3.txt              |  4 +++
 tests/model_fields/test_jsonfield.py | 10 ++++++++
 4 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/django/db/models/fields/json.py b/django/db/models/fields/json.py
index a249f4cdbf..d0a55ac1e7 100644
--- a/django/db/models/fields/json.py
+++ b/django/db/models/fields/json.py
@@ -369,27 +369,26 @@ class KeyTransformIsNull(lookups.IsNull):
 
 
 class KeyTransformIn(lookups.In):
-    def process_rhs(self, compiler, connection):
-        rhs, rhs_params = super().process_rhs(compiler, connection)
-        if not connection.features.has_native_json_field:
-            func = ()
+    def resolve_expression_parameter(self, compiler, connection, sql, param):
+        sql, params = super().resolve_expression_parameter(
+            compiler, connection, sql, param,
+        )
+        if (
+            not hasattr(param, 'as_sql') and
+            not connection.features.has_native_json_field
+        ):
             if connection.vendor == 'oracle':
-                func = []
-                for value in rhs_params:
-                    value = json.loads(value)
-                    function = 'JSON_QUERY' if isinstance(value, (list, dict)) else 'JSON_VALUE'
-                    func.append("%s('%s', '$.value')" % (
-                        function,
-                        json.dumps({'value': value}),
-                    ))
-                func = tuple(func)
-                rhs_params = ()
-            elif connection.vendor == 'mysql' and connection.mysql_is_mariadb:
-                func = ("JSON_UNQUOTE(JSON_EXTRACT(%s, '$'))",) * len(rhs_params)
+                value = json.loads(param)
+                if isinstance(value, (list, dict)):
+                    sql = "JSON_QUERY(%s, '$.value')"
+                else:
+                    sql = "JSON_VALUE(%s, '$.value')"
+                params = (json.dumps({'value': value}),)
             elif connection.vendor in {'sqlite', 'mysql'}:
-                func = ("JSON_EXTRACT(%s, '$')",) * len(rhs_params)
-            rhs = rhs % func
-        return rhs, rhs_params
+                sql = "JSON_EXTRACT(%s, '$')"
+        if connection.vendor == 'mysql' and connection.mysql_is_mariadb:
+            sql = 'JSON_UNQUOTE(%s)' % sql
+        return sql, params
 
 
 class KeyTransformExact(JSONExact):
diff --git a/django/db/models/lookups.py b/django/db/models/lookups.py
index 2874578801..1cc441776f 100644
--- a/django/db/models/lookups.py
+++ b/django/db/models/lookups.py
@@ -241,7 +241,7 @@ class FieldGetDbPrepValueIterableMixin(FieldGetDbPrepValueMixin):
         if hasattr(param, 'resolve_expression'):
             param = param.resolve_expression(compiler.query)
         if hasattr(param, 'as_sql'):
-            sql, params = param.as_sql(compiler, connection)
+            sql, params = compiler.compile(param)
         return sql, params
 
     def batch_process_rhs(self, compiler, connection, rhs=None):
diff --git a/docs/releases/3.1.3.txt b/docs/releases/3.1.3.txt
index 8bc4877a0c..e2e764a06d 100644
--- a/docs/releases/3.1.3.txt
+++ b/docs/releases/3.1.3.txt
@@ -25,3 +25,7 @@ Bugfixes
   :class:`~django.contrib.postgres.aggregates.ArrayAgg` and
   :class:`~django.contrib.postgres.aggregates.StringAgg` with ``ordering``
   on key transforms for :class:`~django.db.models.JSONField` (:ticket:`32096`).
+
+* Fixed a regression in Django 3.1 that caused a crash of ``__in`` lookup when
+  using key transforms for :class:`~django.db.models.JSONField` in the lookup
+  value (:ticket:`32096`).
diff --git a/tests/model_fields/test_jsonfield.py b/tests/model_fields/test_jsonfield.py
index a56b7e74df..381636bb2d 100644
--- a/tests/model_fields/test_jsonfield.py
+++ b/tests/model_fields/test_jsonfield.py
@@ -631,6 +631,16 @@ class TestQuerying(TestCase):
             ('value__0__in', [1], [self.objs[5]]),
             ('value__0__in', [1, 3], [self.objs[5]]),
             ('value__foo__in', ['bar'], [self.objs[7]]),
+            (
+                'value__foo__in',
+                [KeyTransform('foo', KeyTransform('bax', 'value'))],
+                [self.objs[7]],
+            ),
+            (
+                'value__foo__in',
+                [KeyTransform('foo', KeyTransform('bax', 'value')), '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]]),