Repository: adamchainz/django-mysql

Issue:
JSONField extra lookups
Summary: `django.contrib.postgres`'s JSONField has several lookups that ours doesn't (at current). These are:
- `contains`
- `contained_by`
- `has_key`, `has_keys`
- `has_any_key`

Also as a bonus we can add:
- `length`

Any of them that can be implemented (easily) with the JSON functions, should be.



Discussion:

Resolving patch:
diff --git a/django_mysql/models/fields/json.py b/django_mysql/models/fields/json.py
--- a/django_mysql/models/fields/json.py
+++ b/django_mysql/models/fields/json.py
@@ -1,16 +1,20 @@
 # -*- coding:utf-8 -*-
 from __future__ import absolute_import, unicode_literals
 
+import collections
 import json
 
 import django
 from django.core import checks
 from django.db import connections
-from django.db.models import Field, Transform
+from django.db.models import Field, IntegerField, Transform
 from django.utils import six
 
 from django_mysql import forms
 from django_mysql.compat import field_class
+from django_mysql.models.lookups import (
+    JSONContainedBy, JSONContains, JSONHasAnyKeys, JSONHasKey, JSONHasKeys
+)
 
 __all__ = ('JSONField',)
 
@@ -96,17 +100,31 @@ def get_prep_value(self, value):
     def get_prep_lookup(self, lookup_type, value):
         if (
             not hasattr(value, '_prepare') and
-            lookup_type in ('exact', 'gt', 'gte', 'lt', 'lte') and
+            lookup_type in ('exact', 'gt', 'gte', 'lt', 'lte', 'contains') and
             value is not None
         ):
             return JSONValue(value)
+        elif lookup_type == 'has_key':
+            if not isinstance(value, six.text_type):
+                raise ValueError(
+                    "JSONField's 'has_key' lookup only works with {} values"
+                    .format(six.text_type.__name__)
+                )
+            return value
+        elif lookup_type in ('has_keys', 'has_any_keys'):
+            if not isinstance(value, collections.Sequence):
+                raise ValueError(
+                    "JSONField's '{}' lookup only works with Sequences"
+                    .format(lookup_type)
+                )
+            return value
 
         return super(JSONField, self).get_prep_lookup(lookup_type, value)
 
     def get_lookup(self, lookup_name):
         # Have to 'unregister' some incompatible lookups
         if lookup_name in {
-            'range', 'in', 'iexact', 'contains', 'icontains', 'startswith',
+            'range', 'in', 'iexact', 'icontains', 'startswith',
             'istartswith', 'endswith', 'iendswith', 'search', 'regex', 'iregex'
         }:
             raise NotImplementedError(
@@ -125,7 +143,27 @@ def __init__(self, value):
         self.json_string = json.dumps(value)
 
     def as_sql(self, *args, **kwargs):
-        return 'CAST(%s AS JSON)', (self.json_string,)
+        return 'CAST(%s AS JSON)', [self.json_string]
+
+
+class JSONLength(Transform):
+    lookup_name = 'length'
+
+    output_field = IntegerField()
+
+    if django.VERSION[:2] < (1, 9):
+        def as_sql(self, compiler, connection):
+            lhs, params = compiler.compile(self.lhs)
+            return 'JSON_LENGTH({})'.format(lhs), params
+    else:
+        function = 'JSON_LENGTH'
+
+JSONField.register_lookup(JSONContains)
+JSONField.register_lookup(JSONContainedBy)
+JSONField.register_lookup(JSONHasAnyKeys)
+JSONField.register_lookup(JSONHasKey)
+JSONField.register_lookup(JSONHasKeys)
+JSONField.register_lookup(JSONLength)
 
 
 class KeyTransform(Transform):
diff --git a/django_mysql/models/functions.py b/django_mysql/models/functions.py
--- a/django_mysql/models/functions.py
+++ b/django_mysql/models/functions.py
@@ -229,8 +229,8 @@ def __init__(self, expression, path=None):
 class JSONLength(Func):
     function = 'JSON_LENGTH'
 
-    def __init__(self, expression, path=None):
-        from django_mysql.models.fields import JSONField
+    def __init__(self, expression, path=None, **extra):
+        output_field = extra.pop('output_field', IntegerField())
 
         exprs = [expression]
         if path is not None:
@@ -238,7 +238,7 @@ def __init__(self, expression, path=None):
                 path = Value(path)
             exprs.append(path)
 
-        super(JSONLength, self).__init__(*exprs, output_field=JSONField())
+        super(JSONLength, self).__init__(*exprs, output_field=output_field)
 
 
 # MariaDB Regexp Functions
diff --git a/django_mysql/models/lookups.py b/django_mysql/models/lookups.py
--- a/django_mysql/models/lookups.py
+++ b/django_mysql/models/lookups.py
@@ -1,4 +1,6 @@
 # -*- coding:utf-8 -*-
+import json
+
 from django.db.models import CharField, Lookup, Transform
 from django.db.models.lookups import BuiltinLookup
 
@@ -29,7 +31,78 @@ def as_sql(self, compiler, connection):
         return "SOUNDEX(%s)" % lhs, params
 
 
-# These are mostly used for custom fields
+# Custom field class lookups
+
+
+# JSONField
+
+
+class JSONContainedBy(Lookup):
+    lookup_name = 'contained_by'
+
+    def as_sql(self, qn, connection):
+        lhs, lhs_params = self.process_lhs(qn, connection)
+        rhs, rhs_params = self.process_rhs(qn, connection)
+        params = rhs_params + lhs_params
+        return 'JSON_CONTAINS({}, {})'.format(rhs, lhs), params
+
+
+class JSONContains(Lookup):
+    lookup_name = 'contains'
+
+    def as_sql(self, qn, connection):
+        lhs, lhs_params = self.process_lhs(qn, connection)
+        rhs, rhs_params = self.process_rhs(qn, connection)
+        params = lhs_params + rhs_params
+        return 'JSON_CONTAINS({}, {})'.format(lhs, rhs), params
+
+
+class JSONHasKey(Lookup):
+    lookup_name = 'has_key'
+
+    def as_sql(self, qn, connection):
+        lhs, lhs_params = self.process_lhs(qn, connection)
+        key_name = self.rhs
+        path = '$.{}'.format(json.dumps(key_name))
+        params = lhs_params + [path]
+        return "JSON_CONTAINS_PATH({}, 'one', %s)".format(lhs), params
+
+
+class JSONHasKeys(Lookup):
+    lookup_name = 'has_keys'
+
+    def as_sql(self, qn, connection):
+        lhs, lhs_params = self.process_lhs(qn, connection)
+        paths = [
+            '$.{}'.format(json.dumps(key_name))
+            for key_name in self.rhs
+        ]
+        params = lhs_params + paths
+
+        sql = ['JSON_CONTAINS_PATH(', lhs, ", 'all', "]
+        sql.append(', '.join('%s' for _ in paths))
+        sql.append(')')
+        return ''.join(sql), params
+
+
+class JSONHasAnyKeys(Lookup):
+    lookup_name = 'has_any_keys'
+
+    def as_sql(self, qn, connection):
+        lhs, lhs_params = self.process_lhs(qn, connection)
+        paths = [
+            '$.{}'.format(json.dumps(key_name))
+            for key_name in self.rhs
+        ]
+        params = lhs_params + paths
+
+        sql = ['JSON_CONTAINS_PATH(', lhs, ", 'one', "]
+        sql.append(', '.join('%s' for _ in paths))
+        sql.append(')')
+        return ''.join(sql), params
+
+
+# Set{Char,Text}Field
 
 
 class SetContains(Lookup):
@@ -47,6 +120,9 @@ class SetIContains(SetContains):
     lookup_name = 'icontains'
 
 
+# DynamicField
+
+
 class DynColHasKey(Lookup):
     lookup_name = 'has_key'
 


Test patch:
diff --git a/tests/testapp/test_functions.py b/tests/testapp/test_functions.py
--- a/tests/testapp/test_functions.py
+++ b/tests/testapp/test_functions.py
@@ -374,6 +374,14 @@ def test_json_length_path(self):
         )
         assert results == [len(self.obj.attrs['sub'])]
 
+    def test_json_length_type(self):
+        results = list(
+            JSONModel.objects.annotate(
+                x=JSONLength('attrs')
+            ).filter(x__range=[3, 5])
+        )
+        assert results == [self.obj]
+
 
 @requiresDatabaseFunctions
 class RegexpFunctionTests(TestCase):
diff --git a/tests/testapp/test_jsonfield.py b/tests/testapp/test_jsonfield.py
--- a/tests/testapp/test_jsonfield.py
+++ b/tests/testapp/test_jsonfield.py
@@ -1,4 +1,6 @@
 # -*- coding:utf-8 -*-
+from __future__ import unicode_literals
+
 import mock
 from unittest import SkipTest
 
@@ -187,16 +189,18 @@ def test_gte_2(self):
         )
 
 
-class SubKeyQueryTests(JSONFieldTestCase):
+class ExtraLookupsQueryTests(JSONFieldTestCase):
 
     def setUp(self):
-        super(SubKeyQueryTests, self).setUp()
+        super(ExtraLookupsQueryTests, self).setUp()
 
         self.objs = [
             JSONModel.objects.create(attrs={}),
             JSONModel.objects.create(attrs={
                 'a': 'b',
                 'c': 1,
+                '"': 'awkward',
+                '\n': 'super awkward'
             }),
             JSONModel.objects.create(attrs={
                 'a': 'b',
@@ -211,9 +215,138 @@ def setUp(self):
             JSONModel.objects.create(attrs={
                 'k': True,
                 'l': False,
+                '\\': 'awkward'
             }),
         ]
 
+    def test_has_key_invalid_type(self):
+        with pytest.raises(ValueError) as excinfo:
+            JSONModel.objects.filter(attrs__has_key=1)
+        assert "'has_key' lookup only works with" in str(excinfo.value)
+
+    def test_has_key(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_key='a')) ==
+            [self.objs[1], self.objs[2]]
+        )
+
+    def test_has_key_2(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_key='l')) ==
+            [self.objs[4]]
+        )
+
+    def test_has_key_awkward(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_key='"')) ==
+            [self.objs[1]]
+        )
+
+    def test_has_key_awkward_2(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_key='\n')) ==
+            [self.objs[1]]
+        )
+
+    def test_has_key_awkward_3(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_key='\\')) ==
+            [self.objs[4]]
+        )
+
+    def test_has_keys_invalid_type(self):
+        with pytest.raises(ValueError) as excinfo:
+            JSONModel.objects.filter(attrs__has_keys={'a': 'b'})
+        assert "only works with Sequences" in str(excinfo.value)
+
+    def test_has_keys(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_keys=['a', 'c'])) ==
+            [self.objs[1], self.objs[2]]
+        )
+
+    def test_has_keys_2(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_keys=['l'])) ==
+            [self.objs[4]]
+        )
+
+    def test_has_keys_awkward(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_keys=['\n', '"'])) ==
+            [self.objs[1]]
+        )
+
+    def test_has_any_keys_invalid_type(self):
+        with pytest.raises(ValueError) as excinfo:
+            JSONModel.objects.filter(attrs__has_any_keys={'a': 'b'})
+        assert "only works with Sequences" in str(excinfo.value)
+
+    def test_has_any_keys(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_any_keys=['a', 'k'])) ==
+            [self.objs[1], self.objs[2], self.objs[4]]
+        )
+
+    def test_has_any_keys_awkward(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__has_any_keys=['\\'])) ==
+            [self.objs[4]]
+        )
+
+    def test_contains(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__contains={'a': 'b'})) ==
+            [self.objs[1], self.objs[2]]
+        )
+
+    def test_contains_2(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__contains={'l': False})) ==
+            [self.objs[4]]
+        )
+
+    def test_contains_array(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__contains=[[2]])) ==
+            [self.objs[3]]
+        )
+
+    def test_contained_by(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__contained_by={
+                'k': True,
+                'l': False,
+                '\\': 'awkward',
+                'spare_key': ['unused', 'array'],
+            })) ==
+            [self.objs[0], self.objs[4]]
+        )
+
+    def test_contained_by_array(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__contained_by=[1, [2], 3])) ==
+            [self.objs[3]]
+        )
+
+    def test_length_lookup(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__length=0)) ==
+            [self.objs[0]]
+        )
+
+    def test_length_lookup_2(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__length=2)) ==
+            [self.objs[3]]
+        )
+
+    def test_length_lookup_chains(self):
+        assert (
+            list(JSONModel.objects.filter(attrs__length__range=[0, 10])) ==
+            self.objs
+        )
+
     def test_shallow_list_lookup(self):
         assert (
             list(JSONModel.objects.filter(attrs__0=1)) ==
diff --git a/tests/testapp/utils.py b/tests/testapp/utils.py
--- a/tests/testapp/utils.py
+++ b/tests/testapp/utils.py
@@ -70,7 +70,7 @@ def query(self):
         return self.capturer.captured_queries[-1]['sql']
 
 
-class PrintAllQueries(object):
+class print_all_queries(object):
     def __init__(self, conn=None):
         if conn is None:
             self.conn = connection