Repository: adamchainz/django-mysql

Issue:
Allow passing output_field to JSONExtract
Summary: Currently some JSON field related database functions accept output_field as parameters, but `JSONExtract` always assumes that `JSONField` is needed, even though the value extracted is a flat type value.

Consider this example:
assume we have a model `Document` with `JSONField` `dynamic_data`. Assume the following content of a `Document` instance's `dynamic_data`:
`{
    'title': 'amazing title'
}`

Currently this query is not possible:
`Document.objects.annotate(title=JSONExtract("dynamic_data", "$.title")).filter(title__icontains="amazing")`
because
`NotImplementedError: Lookup 'icontains' doesn't work with JSONField`
since `JSONField` is hardcoded in `JSONExtract`.

I suggest adding `output_field` parameter to `JSONExtract` so that it would be possible to use `icontains` on annotated data.


Discussion:
I'm not sure if this works in general, as `JSON_EXTRACT` returns an expression of type `json` in MySQL. I think you'd need to `CAST` to `CHAR` as before you could use `LIKE` (what `icontains` maps to).
I have tried it like this:

    class EzJSONExtract(Func):
        function = 'JSON_EXTRACT'

        def __init__(self, expression, *paths, output_field=None):
            from django_mysql.models.fields import JSONField

            exprs = [expression]
            for path in paths:
                if not hasattr(path, 'resolve_expression'):
                    path = Value(path)
                exprs.append(path)
            output_field = output_field or JSONField()
            super().__init__(*exprs, output_field=output_field)

Then
`Document.objects.annotate(title=EzJSONExtract("dynamic_data", "$.title", output_field=CharField())).filter(title__icontains="amazing")` works just fine
Seems `ez` easy.

Will look into it a bit more before adding.

Resolving patch:
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
@@ -220,7 +220,7 @@ def get(cls, using=DEFAULT_DB_ALIAS):
 class JSONExtract(Func):
     function = 'JSON_EXTRACT'
 
-    def __init__(self, expression, *paths):
+    def __init__(self, expression, *paths, **kwargs):
         from django_mysql.models.fields import JSONField
 
         exprs = [expression]
@@ -229,7 +229,23 @@ def __init__(self, expression, *paths):
                 path = Value(path)
             exprs.append(path)
 
-        super(JSONExtract, self).__init__(*exprs, output_field=JSONField())
+        # kwarg validation - for Python 2 compat
+        if list(kwargs.keys()) not in ([], ['output_field']):
+            raise TypeError(
+                'Only supported keyword argument is "output_field"'
+            )
+
+        if 'output_field' in kwargs:
+            if len(paths) > 1:
+                raise TypeError(
+                    "output_field won't work with more than one path, as a "
+                    "JSON Array will be returned"
+                )
+            output_field = kwargs['output_field']
+        else:
+            output_field = JSONField()
+
+        super(JSONExtract, self).__init__(*exprs, output_field=output_field)
 
 
 class JSONKeys(Func):


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
@@ -8,7 +8,7 @@
 
 import pytest
 from django.db import connection
-from django.db.models import F, Q, Value
+from django.db.models import F, FloatField, Q, Value
 from django.db.models.functions import Length, Lower, Upper
 from django.test import TestCase
 from django.utils import six
@@ -357,6 +357,19 @@ def setUp(self):
             'arr': ['dee', 'arr', 'arr'],
         })
 
+    def test_json_extract_kwarg_bad(self):
+        with pytest.raises(TypeError) as excinfo:
+            JSONExtract('foo', 'bar', foo=1)
+        assert 'Only supported keyword' in six.text_type(excinfo.value)
+
+    def test_json_extract_output_field_too_many_paths(self):
+        with pytest.raises(TypeError) as excinfo:
+            JSONExtract('foo', 'bar', 'baz', output_field=1)
+        assert (
+            "output_field won't work with more than one path" in
+            six.text_type(excinfo.value)
+        )
+
     def test_json_extract_flote(self):
         results = list(
             JSONModel.objects.annotate(
@@ -366,6 +379,17 @@ def test_json_extract_flote(self):
         assert results == [1.5]
         assert isinstance(results[0], float)
 
+    def test_json_extract_flote_as_float(self):
+        results = list(
+            JSONModel.objects.annotate(
+                x=JSONExtract('attrs', '$.flote', output_field=FloatField())
+            ).filter(
+                x__gt=0.1
+            ).values_list('x', flat=True)
+        )
+        assert results == [1.5]
+        assert isinstance(results[0], float)
+
     def test_json_extract_multiple(self):
         results = list(
             JSONModel.objects.annotate(