Fix JSONField has key lookup

Author: absci

mssql/functions.py

@@ -12,7 +12,9 @@
 from django.db.models import lookups
 
 if VERSION >= (3, 1):
-    from django.db.models.fields.json import KeyTransform, KeyTransformIn, KeyTransformExact
+    from django.db.models.fields.json import (
+        KeyTransform, KeyTransformIn, KeyTransformExact,
+        HasKeyLookup, compile_json_path)
 
 DJANGO3 = VERSION[0] >= 3
 
@@ -134,11 +136,42 @@ def json_KeyTransformIn(self, compiler, connection):
 
     return (lhs + ' IN ' + rhs, unquote_json_rhs(rhs_params))
 
+def json_HasKeyLookup(self, compiler, connection):
+    # Process JSON path from the left-hand side.
+    if isinstance(self.lhs, KeyTransform):
+        lhs, _, lhs_key_transforms = self.lhs.preprocess_lhs(compiler, connection)
+        lhs_json_path = compile_json_path(lhs_key_transforms)
+    else:
+        lhs, _ = self.process_lhs(compiler, connection)
+        lhs_json_path = '$'
+    sql = lhs + ' IN (SELECT ' + lhs + ' FROM ' + self.lhs.output_field.model._meta.db_table + \
+    ' CROSS APPLY OPENJSON(' + lhs + ') WITH ( [json_path_value] char(1) \'%s\') WHERE [json_path_value] IS NOT NULL)'
+    # Process JSON path from the right-hand side.
+    rhs = self.rhs
+    rhs_params = []
+    if not isinstance(rhs, (list, tuple)):
+        rhs = [rhs]
+    for key in rhs:
+        if isinstance(key, KeyTransform):
+            *_, rhs_key_transforms = key.preprocess_lhs(compiler, connection)
+        else:
+            rhs_key_transforms = [key]
+        rhs_params.append('%s%s' % (
+            lhs_json_path,
+            compile_json_path(rhs_key_transforms, include_root=False),
+        ))
+    # Add condition for each key.
+    if self.logical_operator:
+        sql = '(%s)' % self.logical_operator.join([sql] * len(rhs_params))
+
+    return sql % tuple(rhs_params), []
+
 ATan2.as_microsoft = sqlserver_atan2
 In.split_parameter_list_as_sql = split_parameter_list_as_sql
 if VERSION >= (3, 1):
     KeyTransformIn.as_microsoft = json_KeyTransformIn
     KeyTransformExact.process_rhs = json_KeyTransformExact_process_rhs
+    HasKeyLookup.as_microsoft = json_HasKeyLookup
 Ln.as_microsoft = sqlserver_ln
 Log.as_microsoft = sqlserver_log
 Mod.as_microsoft = sqlserver_mod

testapp/settings.py

@@ -191,28 +191,25 @@
                   'datetimes.tests.DateTimesTests.test_21432',
 
                   #JSONFields
-                  'model_fields.test_jsonfield.TestQuerying.test_has_any_keys',
-                  'model_fields.test_jsonfield.TestQuerying.test_has_key',
-                  'model_fields.test_jsonfield.TestQuerying.test_has_key_deep',
                   'model_fields.test_jsonfield.TestQuerying.test_has_key_list',
                   'model_fields.test_jsonfield.TestQuerying.test_has_key_null_value',
-                  'model_fields.test_jsonfield.TestQuerying.test_has_keys',
-                  'model_fields.test_jsonfield.TestQuerying.test_key_iregex',
                   '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_ordering_grouping_by_count',
                   'model_fields.test_jsonfield.JSONFieldTests.test_db_check_constraints',
                   'model_fields.test_jsonfield.TestQuerying.test_isnull_key',
                   '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_ordering_by_transform',
                   'model_fields.test_jsonfield.TestQuerying.test_key_escape',
+
+                  'model_fields.test_jsonfield.TestQuerying.test_ordering_by_transform',
                   ]
 
 REGEX_TESTS = ['lookup.tests.LookupTests.test_regex',
                'lookup.tests.LookupTests.test_regex_backreferencing',
                'lookup.tests.LookupTests.test_regex_non_ascii',
                'lookup.tests.LookupTests.test_regex_non_string',
-               'lookup.tests.LookupTests.test_regex_null'
+               'lookup.tests.LookupTests.test_regex_null',
+               'model_fields.test_jsonfield.TestQuerying.test_key_iregex',
+               'model_fields.test_jsonfield.TestQuerying.test_key_regex',
                ]