Fix KeyTransformExact applied to all databases (#98)

* Fixed issue #82

Author: jmah8

mssql/functions.py

@@ -4,18 +4,16 @@
 import json
 
 from django import VERSION
-
+from django.core import validators
 from django.db import NotSupportedError, connections, transaction
-from django.db.models import BooleanField, Value
-from django.db.models.functions import Cast, NthValue
-from django.db.models.functions.math import ATan2, Log, Ln, Mod, Round
-from django.db.models.expressions import Case, Exists, OrderBy, When, Window, Expression
-from django.db.models.lookups import Lookup, In
-from django.db.models import lookups, CheckConstraint
+from django.db.models import BooleanField, CheckConstraint, Value
+from django.db.models.expressions import Case, Exists, Expression, OrderBy, When, Window
 from django.db.models.fields import BinaryField, Field
-from django.db.models.sql.query import Query
+from django.db.models.functions import Cast, NthValue
+from django.db.models.functions.math import ATan2, Ln, Log, Mod, Round
+from django.db.models.lookups import In, Lookup
 from django.db.models.query import QuerySet
-from django.core import validators
+from django.db.models.sql.query import Query
 
 if VERSION >= (3, 1):
     from django.db.models.fields.json import (
@@ -67,9 +65,11 @@ def sqlserver_nth_value(self, compiler, connection, **extra_content):
 def sqlserver_round(self, compiler, connection, **extra_context):
     return self.as_sql(compiler, connection, template='%(function)s(%(expressions)s, 0)', **extra_context)
 
+
 def sqlserver_random(self, compiler, connection, **extra_context):
     return self.as_sql(compiler, connection, function='RAND', **extra_context)
 
+
 def sqlserver_window(self, compiler, connection, template=None):
     # MSSQL window functions require an OVER clause with ORDER BY
     if self.order_by is None:
@@ -143,26 +143,29 @@ def split_parameter_list_as_sql(self, compiler, connection):
 
     return in_clause, ()
 
+
 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(lookups.Exact, self).process_rhs(compiler, connection)
-    rhs, rhs_params = super(KeyTransformExact, self).process_rhs(compiler, connection)
+    rhs, rhs_params = key_transform_exact_process_rhs(self, compiler, connection)
+    if connection.vendor == 'microsoft':
+        rhs_params = unquote_json_rhs(rhs_params)
+    return rhs, rhs_params
 
-    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))
 
+
 def json_HasKeyLookup(self, compiler, connection):
     # Process JSON path from the left-hand side.
     if isinstance(self.lhs, KeyTransform):
@@ -193,6 +196,7 @@ def json_HasKeyLookup(self, compiler, connection):
 
     return sql % tuple(rhs_params), []
 
+
 def BinaryField_init(self, *args, **kwargs):
     # Add max_length option for BinaryField, default to max
     kwargs.setdefault('editable', False)
@@ -202,6 +206,7 @@ def BinaryField_init(self, *args, **kwargs):
     else:
         self.max_length = 'max'
 
+
 def _get_check_sql(self, model, schema_editor):
     if VERSION >= (3, 1):
         query = Query(model=model, alias_cols=False)
@@ -217,6 +222,7 @@ def _get_check_sql(self, model, schema_editor):
 
     return sql % tuple(schema_editor.quote_value(p) for p in params)
 
+
 def bulk_update_with_default(self, objs, fields, batch_size=None, default=0):
     """
         Update the given fields in each of the given objects in the database.
@@ -255,7 +261,7 @@ def bulk_update_with_default(self, objs, fields, batch_size=None, default=0):
                 attr = getattr(obj, field.attname)
                 if not isinstance(attr, Expression):
                     if attr is None:
-                        value_none_counter+=1
+                        value_none_counter += 1
                     attr = Value(attr, output_field=field)
                 when_statements.append(When(pk=obj.pk, then=attr))
             if(value_none_counter == len(when_statements)):
@@ -272,10 +278,13 @@ def bulk_update_with_default(self, objs, fields, batch_size=None, default=0):
             rows_updated += self.filter(pk__in=pks).update(**update_kwargs)
     return rows_updated
 
+
 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
+    # Need copy of old KeyTransformExact.process_rhs to call later
+    key_transform_exact_process_rhs = KeyTransformExact.process_rhs
     KeyTransformExact.process_rhs = json_KeyTransformExact_process_rhs
     HasKeyLookup.as_microsoft = json_HasKeyLookup
 Ln.as_microsoft = sqlserver_ln

testapp/migrations/0016_jsonmodel.py

@@ -0,0 +1,28 @@
+# Generated by Django 4.0.1 on 2022-02-01 15:58
+
+from django import VERSION
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('testapp', '0015_test_rename_m2mfield_part2'),
+    ]
+
+    # JSONField added in Django 3.1
+    if VERSION >= (3, 1):
+        operations = [
+            migrations.CreateModel(
+                name='JSONModel',
+                fields=[
+                    ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                    ('value', models.JSONField()),
+                ],
+                options={
+                    'required_db_features': {'supports_json_field'},
+                },
+            ),
+        ]
+    else:
+        pass

testapp/models.py

@@ -3,6 +3,7 @@
 
 import uuid
 
+from django import VERSION
 from django.db import models
 from django.db.models import Q
 from django.utils import timezone
@@ -151,3 +152,11 @@ class Meta:
 
     _type = models.CharField(max_length=50)
     status = models.CharField(max_length=50)
+
+
+if VERSION >= (3, 1):
+    class JSONModel(models.Model):
+        value = models.JSONField()
+
+        class Meta:
+            required_db_features = {'supports_json_field'}

testapp/settings.py

@@ -1,6 +1,11 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the BSD license.
 import os
+from pathlib import Path
+
+from django import VERSION
+
+BASE_DIR = Path(__file__).resolve().parent.parent
 
 DATABASES = {
     "default": {
@@ -23,6 +28,14 @@
     },
 }
 
+# Django 3.0 and below unit test doesn't handle more than 2 databases in DATABASES correctly
+if VERSION >= (3, 1):
+    DATABASES['sqlite'] = {
+        "ENGINE": "django.db.backends.sqlite3",
+        "NAME": str(BASE_DIR / "db.sqlitetest"),
+    }
+
+
 # Set to `True` locally if you want SQL queries logged to django_sql.log
 DEBUG = False
 
@@ -267,6 +280,7 @@
     'backends.tests.BackendTestCase.test_queries_logger',
     'migrations.test_operations.OperationTests.test_alter_field_pk_mti_fk',
     'migrations.test_operations.OperationTests.test_run_sql_add_missing_semicolon_on_collect_sql',
+    'migrations.test_operations.OperationTests.test_alter_field_pk_mti_and_fk_to_base'
 ]
 
 REGEX_TESTS = [

testapp/tests/test_jsonfield.py

@@ -0,0 +1,71 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the BSD license.
+
+from unittest import skipUnless
+
+from django import VERSION
+from django.test import TestCase
+
+if VERSION >= (3, 1):
+    from ..models import JSONModel
+
+
+def _check_jsonfield_supported_sqlite():
+    # Info about JSONField support in SQLite: https://code.djangoproject.com/wiki/JSON1Extension
+    import sqlite3
+
+    supports_jsonfield = True
+    try:
+        conn = sqlite3.connect(':memory:')
+        cursor = conn.cursor()
+        cursor.execute('SELECT JSON(\'{"a": "b"}\')')
+    except sqlite3.OperationalError:
+        supports_jsonfield = False
+    finally:
+        return supports_jsonfield
+
+
+class TestJSONField(TestCase):
+    databases = ['default']
+    # Django 3.0 and below unit test doesn't handle more than 2 databases in DATABASES correctly
+    if VERSION >= (3, 1):
+        databases.append('sqlite')
+
+    json = {
+        'a': 'b',
+        'b': 1,
+        'c': '1',
+        'd': [],
+        'e': [1, 2],
+        'f': ['a', 'b'],
+        'g': [1, 'a'],
+        'h': {},
+        'i': {'j': 1},
+        'j': False,
+        'k': True,
+        'l': {
+            'foo': 'bar',
+            'baz': {'a': 'b', 'c': 'd'},
+            'bar': ['foo', 'bar'],
+            'bax': {'foo': 'bar'},
+        },
+    }
+
+    @skipUnless(VERSION >= (3, 1), "JSONField not support in Django versions < 3.1")
+    @skipUnless(
+        _check_jsonfield_supported_sqlite(),
+        "JSONField not support by SQLite on this platform and Python version",
+    )
+    def test_keytransformexact_not_overriding(self):
+        # Issue https://github.com/microsoft/mssql-django/issues/82
+        json_obj = JSONModel(value=self.json)
+        json_obj.save()
+        self.assertSequenceEqual(
+            JSONModel.objects.filter(value__a='b'),
+            [json_obj],
+        )
+        json_obj.save(using='sqlite')
+        self.assertSequenceEqual(
+            JSONModel.objects.using('sqlite').filter(value__a='b'),
+            [json_obj],
+        )