Fixed #35718 -- Add JSONArray to django.db.models.functions.

django · Sep 4, 2024 · d8bba27 · d8bba27
1 parent aa52930
commit d8bba27
Show file tree

Hide file tree

Showing 4 changed files with 257 additions and 1 deletion.
diff --git a/django/db/models/functions/__init__.py b/django/db/models/functions/__init__.py
@@ -1,4 +1,13 @@
-from .comparison import Cast, Coalesce, Collate, Greatest, JSONObject, Least, NullIf
+from .comparison import (
+    Cast,
+    Coalesce,
+    Collate,
+    Greatest,
+    JSONArray,
+    JSONObject,
+    Least,
+    NullIf,
+)
 from .datetime import (
     Extract,
     ExtractDay,
@@ -97,6 +106,7 @@
     "Coalesce",
     "Collate",
     "Greatest",
+    "JSONArray",
     "JSONObject",
     "Least",
     "NullIf",

diff --git a/django/db/models/functions/comparison.py b/django/db/models/functions/comparison.py
@@ -143,6 +143,46 @@ def as_sqlite(self, compiler, connection, **extra_context):
         return super().as_sqlite(compiler, connection, function="MAX", **extra_context)
 
 
+class JSONArray(Func):
+    function = "JSON_ARRAY"
+    output_field = JSONField()
+
+    def as_sql(self, compiler, connection, **extra_context):
+        if not connection.features.has_json_object_function:
+            raise NotSupportedError(
+                "JSONObject() is not supported on this database backend."
+            )
+        return super().as_sql(compiler, connection, **extra_context)
+
+    def as_native(self, compiler, connection, *, returning, **extra_context):
+        return self.as_sql(
+            compiler,
+            connection,
+            template=f"%(function)s(%(expressions)s RETURNING {returning})",
+            **extra_context,
+        )
+
+    def as_postgresql(self, compiler, connection, **extra_context):
+        if not connection.features.is_postgresql_16:
+            copy = self.copy()
+            copy.set_source_expressions(
+                [
+                    Cast(expression, TextField())
+                    for expression in enumerate(copy.get_source_expressions())
+                ]
+            )
+            return super(JSONArray, copy).as_sql(
+                compiler,
+                connection,
+                function="JSONB_BUILD_OBJECT",
+                **extra_context,
+            )
+        return self.as_native(compiler, connection, returning="JSONB", **extra_context)
+
+    def as_oracle(self, compiler, connection, **extra_context):
+        return self.as_native(compiler, connection, returning="CLOB", **extra_context)
+
+
 class JSONObject(Func):
     function = "JSON_OBJECT"
     output_field = JSONField()

diff --git a/tests/db_functions/comparison/test_json_array.py b/tests/db_functions/comparison/test_json_array.py
@@ -0,0 +1,108 @@
+from django.db import NotSupportedError
+from django.db.models import F, Value
+from django.db.models.functions import JSONArray, Lower
+from django.test import TestCase
+from django.test.testcases import skipIfDBFeature, skipUnlessDBFeature
+from django.utils import timezone
+
+from ..models import Article, Author
+
+
+@skipUnlessDBFeature("has_json_object_function")
+class JSONArrayTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        Author.objects.bulk_create(
+            [
+                Author(name="Ivan Ivanov", alias="iivanov"),
+                Author(name="Bertha Berthy", alias="bberthy"),
+            ]
+        )
+
+    def test_empty(self):
+        obj = Author.objects.annotate(json_array=JSONArray()).first()
+        self.assertEqual(obj.json_array, [])
+
+    def test_basic(self):
+        obj = Author.objects.annotate(
+            json_array=JSONArray(Value("name"), F("name"))
+        ).first()
+        self.assertEqual(obj.json_array, ["name", "Ivan Ivanov"])
+
+    def test_expressions(self):
+        obj = Author.objects.annotate(
+            json_array=JSONArray(
+                Lower("name"),
+                F("alias"),
+                F("goes_by"),
+                Value(30000.15),
+                F("age") * 2,
+            )
+        ).first()
+        self.assertEqual(
+            obj.json_array,
+            [
+                "ivan ivanov",
+                "iivanov",
+                None,
+                30000.15,
+                60,
+            ],
+        )
+
+    def test_nested_json_array(self):
+        obj = Author.objects.annotate(
+            json_array=JSONArray(
+                F("name"),
+                JSONArray(F("alias"), F("age")),
+            )
+        ).first()
+        self.assertEqual(
+            obj.json_array,
+            [
+                "Ivan Ivanov",
+                ["iivanov", 30],
+            ],
+        )
+
+    def test_nested_empty_json_array(self):
+        obj = Author.objects.annotate(
+            json_array=JSONArray(
+                F("name"),
+                JSONArray(),
+            )
+        ).first()
+        self.assertEqual(
+            obj.json_array,
+            [
+                "Ivan Ivanov",
+                [],
+            ],
+        )
+
+    def test_textfield(self):
+        Article.objects.create(
+            title="The Title",
+            text="x" * 4000,
+            written=timezone.now(),
+        )
+        obj = Article.objects.annotate(json_array=JSONArray(F("text"))).first()
+        self.assertEqual(obj.json_array, ["x" * 4000])
+
+    def test_order_by_key(self):
+        qs = Author.objects.annotate(attrs=JSONArray(F("arr"))).order_by("arr__0")
+        self.assertQuerySetEqual(qs, Author.objects.order_by("alias"))
+
+    def test_order_by_nested_key(self):
+        qs = Author.objects.annotate(arr=JSONArray(JSONArray(F("alias")))).order_by(
+            "-arr__0__0"
+        )
+        self.assertQuerySetEqual(qs, Author.objects.order_by("-alias"))
+
+
+@skipIfDBFeature("has_json_object_function")
+class JSONObjectNotSupportedTests(TestCase):
+    def test_not_supported(self):
+        msg = "JSONArray() is not supported on this database backend."
+        with self.assertRaisesMessage(NotSupportedError, msg):
+            Author.objects.annotate(json_array=JSONArray()).get()
diff --git a/tests/db_functions/comparison/test_json_array_object.py b/tests/db_functions/comparison/test_json_array_object.py
@@ -0,0 +1,98 @@
+from django.db import NotSupportedError
+from django.db.models import F, Value
+from django.db.models.functions import JSONArray, JSONObject, Lower
+from django.test import TestCase
+from django.test.testcases import skipIfDBFeature, skipUnlessDBFeature
+from django.utils import timezone
+
+from ..models import Article, Author
+
+
+@skipUnlessDBFeature("has_json_object_function")
+class JSONArrayObjectTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        Author.objects.bulk_create(
+            [
+                Author(name="Ivan Ivanov", alias="iivanov"),
+                Author(name="Bertha Berthy", alias="bberthy"),
+            ]
+        )
+
+    def test_nested_json_array_object(self):
+        obj = Author.objects.annotate(
+            json_array=JSONArray(
+                JSONObject(
+                    name1="name",
+                    nested_json_object1=JSONObject(
+                        alias1="alias",
+                        age1="age",
+                    ),
+                ),
+                JSONObject(
+                    name2="name",
+                    nested_json_object2=JSONObject(
+                        alias2="alias",
+                        age2="age",
+                    ),
+                ),
+            )
+        ).first()
+        self.assertEqual(
+            obj.json_array,
+            [
+                {
+                    "name1": "Ivan Ivanov",
+                    "nested_json_object1": {
+                        "alias1": "iivanov",
+                        "age1": 30,
+                    },
+                },
+                {
+                    "name2": "Ivan Ivanov",
+                    "nested_json_object2": {
+                        "alias2": "iivanov",
+                        "age2": 30,
+                    },
+                },
+            ],
+        )
+
+    def test_nested_json_object_array(self):
+        obj = Author.objects.annotate(
+            json_object=JSONObject(
+                name="name",
+                nested_json_array=JSONArray(
+                    JSONObject(
+                        alias1="alias",
+                        age1="age",
+                    ),
+                    JSONObject(
+                        alias2="alias",
+                        age2="age",
+                    ),
+                ),
+            )
+        ).first()
+        self.assertEqual(
+            obj.json_object,
+            {
+                "name": "Ivan Ivanov",
+                "nested_json_array": [
+                    {
+                        "alias1": "iivanov",
+                        "age1": 30,
+                    },
+                    {
+                        "alias2": "iivanov",
+                        "age2": 30,
+                    },
+                ],
+            },
+        )
+
+    def test_order_by_nested_key(self):
+        qs = Author.objects.annotate(
+            arr=JSONArray(JSONObject(alias=F("alias")))
+        ).order_by("-arr__0__alias")
+        self.assertQuerySetEqual(qs, Author.objects.order_by("-alias"))