Fix JSONField KeyTransformIn, KeyTransformExact

Author: absci

azure-pipelines.yml

@@ -17,6 +17,7 @@ jobs:
   - job: Windows
     pool:
       Django-agent-pool
+    timeoutInMinutes: 120
 
     strategy:
       matrix:

mssql/compiler.py

@@ -46,8 +46,10 @@ def _as_sql_greatest(self, compiler, connection):
 def _as_sql_json_keytransform(self, compiler, connection):
     lhs, params, key_transforms = self.preprocess_lhs(compiler, connection)
     json_path = compile_json_path(key_transforms)
-
-    return 'JSON_VALUE(%s, %%s)' % lhs, tuple(params) + (json_path,)
+    return (
+        "COALESCE(JSON_QUERY(%s, '%s'), JSON_VALUE(%s, '%s'))" %
+        ((lhs, json_path) * 2)
+    ), tuple(params) * 2
 
 def _as_sql_least(self, compiler, connection):
     # SQL Server does not provide LEAST function,

mssql/functions.py

@@ -1,14 +1,18 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import json
+
 from django import VERSION
 from django.db.models import BooleanField
 from django.db.models.functions import Cast
 from django.db.models.functions.math import ATan2, Log, Ln, Mod, Round
 from django.db.models.expressions import Case, Exists, OrderBy, When
-from django.db.models.lookups import Lookup, In, Exact
+from django.db.models.lookups import Lookup, In
+from django.db.models import lookups
+
 if VERSION >= (3, 1):
-    from django.db.models.fields.json import KeyTransform, KeyTransformExact
+    from django.db.models.fields.json import KeyTransform, KeyTransformIn, KeyTransformExact
 
 DJANGO3 = VERSION[0] >= 3
 
@@ -110,19 +114,31 @@ def split_parameter_list_as_sql(self, compiler, connection):
 
     return in_clause, ()
 
-def KeyTransformExact_process_rhs(self, compiler, connection):
+def unquote_json_rhs(rhs_params):
+    for value in rhs_params:
+        value = json.loads(value)
+        if not isinstance(value, (list, dict)):
+            rhs_params = [param.replace('"', '') for param in rhs_params]
+    return rhs_params
+
+def json_KeyTransformExact_process_rhs(self, compiler, connection):
     if isinstance(self.rhs, KeyTransform):
-        return super(Exact, self).process_rhs(compiler, connection)
-    rhs, rhs_params = super(Exact, self).process_rhs(compiler, connection)
-    if connection.vendor == 'microsoft':
-        if rhs_params != [None]:
-            rhs_params = [params.strip('"') for params in rhs_params]
-    return rhs, rhs_params
+        return super(lookups.Exact, self).process_rhs(compiler, connection)
+    rhs, rhs_params = super(KeyTransformExact, self).process_rhs(compiler, connection)
+
+    return rhs, unquote_json_rhs(rhs_params)
+
+def json_KeyTransformIn(self, compiler, connection):
+    lhs, _ = super(KeyTransformIn, self).process_lhs(compiler, connection)
+    rhs, rhs_params = super(KeyTransformIn, self).process_rhs(compiler, connection)
+
+    return (lhs + ' IN ' + rhs, unquote_json_rhs(rhs_params))
 
 ATan2.as_microsoft = sqlserver_atan2
 In.split_parameter_list_as_sql = split_parameter_list_as_sql
 if VERSION >= (3, 1):
-    KeyTransformExact.process_rhs = KeyTransformExact_process_rhs
+    KeyTransformIn.as_microsoft = json_KeyTransformIn
+    KeyTransformExact.process_rhs = json_KeyTransformExact_process_rhs
 Ln.as_microsoft = sqlserver_ln
 Log.as_microsoft = sqlserver_log
 Mod.as_microsoft = sqlserver_mod

testapp/settings.py

@@ -201,20 +201,13 @@
                   'model_fields.test_jsonfield.TestQuerying.test_key_quoted_string',
                   'model_fields.test_jsonfield.TestQuerying.test_key_regex',
                   'model_fields.test_jsonfield.TestQuerying.test_lookups_with_key_transform',
-                  'model_fields.test_jsonfield.TestQuerying.test_order_grouping_custom_decoder',
                   'model_fields.test_jsonfield.TestQuerying.test_ordering_grouping_by_count',
-                  'model_fields.test_jsonfield.TestQuerying.test_ordering_grouping_by_key_transform',
                   'model_fields.test_jsonfield.JSONFieldTests.test_db_check_constraints',
                   'model_fields.test_jsonfield.TestQuerying.test_isnull_key',
-                  'model_fields.test_jsonfield.TestQuerying.test_key_in',
-                  'model_fields.test_jsonfield.TestQuerying.test_key_transform_expression',
-                  'model_fields.test_jsonfield.TestQuerying.test_key_values',
-                  'model_fields.test_jsonfield.TestQuerying.test_nested_key_transform_expression',
                   'model_fields.test_jsonfield.TestQuerying.test_none_key',
                   'model_fields.test_jsonfield.TestQuerying.test_none_key_and_exact_lookup',
-                  'model_fields.test_jsonfield.TestQuerying.test_none_key_exclude',
                   'model_fields.test_jsonfield.TestQuerying.test_ordering_by_transform',
-                  'model_fields.test_jsonfield.TestQuerying.test_shallow_lookup_obj_target'
+                  'model_fields.test_jsonfield.TestQuerying.test_key_escape',
                   ]
 
 REGEX_TESTS = ['lookup.tests.LookupTests.test_regex',