1399 load_json and array values with prefetch (#1408)

* add test

* wip

* break up join tests

* fix bug

* add test assertions

* simplify

* add comment

Author: dantownsend

piccolo/query/base.py

@@ -105,28 +105,27 @@ async def _process_results(self, results) -> QueryResponseType:
 
         raw = await self.response_handler(raw)
 
-        if output:
-            if output._output.as_objects:
-                if output._output.nested:
-                    return cast(
-                        QueryResponseType,
-                        [
-                            make_nested_object(
-                                row,
-                                self.table,
-                                load_json=output._output.load_json,
-                            )
-                            for row in raw
-                        ],
-                    )
-                else:
-                    return cast(
-                        QueryResponseType,
-                        [
-                            self.table(**columns, _exists_in_db=True)
-                            for columns in raw
-                        ],
-                    )
+        if output and output._output.as_objects:
+            if output._output.nested:
+                return cast(
+                    QueryResponseType,
+                    [
+                        make_nested_object(
+                            row,
+                            self.table,
+                            load_json=output._output.load_json,
+                        )
+                        for row in raw
+                    ],
+                )
+            else:
+                return cast(
+                    QueryResponseType,
+                    [
+                        self.table(**columns, _exists_in_db=True)
+                        for columns in raw
+                    ],
+                )
 
         return cast(QueryResponseType, raw)
 

piccolo/utils/encoding.py

@@ -63,7 +63,7 @@ class JSONDict(dict):
     ...
 
 
-def load_json(data: str) -> Any:
+def load_json(data: str | bytes | bytearray) -> Any:
     response = (
         orjson.loads(data) if ORJSON else json.loads(data)  # type: ignore
     )

piccolo/utils/objects.py

@@ -69,7 +69,15 @@ def make_nested_object(
             else:
                 # The value doesn't belong to a foreign key, so just append it.
                 table_params[key] = value
-        elif load_json and key in json_column_names and value is not None:
+        elif (
+            load_json
+            and key in json_column_names
+            # The top level might already have the JSON loaded, so check for
+            # a string, rather than it being non-null.
+            # https://github.com/piccolo-orm/piccolo/issues/1399
+            # Ideally we want a cleaner solution than this.
+            and isinstance(value, (str, bytes, bytearray))
+        ):
             table_params[key] = encoding.load_json(value)
         else:
             table_params[key] = value

tests/table/test_join.py

@@ -1,6 +1,9 @@
 import decimal
+from unittest import TestCase
 
-from piccolo.testing.test_case import TableTest
+from piccolo.columns.column_types import JSON, ForeignKey, Serial, Varchar
+from piccolo.table import Table
+from piccolo.testing.test_case import AsyncTableTest, TableTest
 from tests.base import engine_is
 from tests.example_apps.music.tables import (
     Band,
@@ -18,6 +21,9 @@
 
 class TestCreateJoin:
     def test_create_join(self):
+        """
+        A simple test to make sure tables with foreign keys can be created.
+        """
         for table in TABLES:
             table.create_table().run_sync()
 
@@ -74,9 +80,10 @@ def setUp(self):
         )
         instrument.save().run_sync()
 
-    ###########################################################################
 
-    def test_join(self):
+class TestSelectJoin(TestJoin):
+
+    def test_select_join(self):
         select_query = Concert.select(
             Concert.band_1.name,
             Concert.band_2.name,
@@ -211,25 +218,6 @@ def test_select_all_columns_deep(self):
                 },
             )
 
-    def test_proxy_columns(self):
-        """
-        Make sure that ``proxy_columns`` are set correctly.
-
-        There used to be a bug which meant queries got slower over time:
-
-        https://github.com/piccolo-orm/piccolo/issues/691
-
-        """
-        # We call it multiple times to make sure it doesn't change with time.
-        for _ in range(2):
-            self.assertEqual(
-                len(Concert.band_1._.manager._foreign_key_meta.proxy_columns),
-                2,
-            )
-            self.assertEqual(
-                len(Concert.band_1._foreign_key_meta.proxy_columns), 4
-            )
-
     def test_select_all_columns_root(self):
         """
         Make sure that using ``all_columns`` at the root doesn't interfere
@@ -343,7 +331,29 @@ def test_select_all_columns_exclude(self):
                 },
             )
 
-    ###########################################################################
+
+class TestProxyColumns(TestCase):
+    def test_proxy_columns(self):
+        """
+        Make sure that ``proxy_columns`` are set correctly.
+
+        There used to be a bug which meant queries got slower over time:
+
+        https://github.com/piccolo-orm/piccolo/issues/691
+
+        """
+        # We call it multiple times to make sure it doesn't change with time.
+        for _ in range(2):
+            self.assertEqual(
+                len(Concert.band_1._.manager._foreign_key_meta.proxy_columns),
+                2,
+            )
+            self.assertEqual(
+                len(Concert.band_1._foreign_key_meta.proxy_columns), 4
+            )
+
+
+class TestObjectsJoin(TestJoin):
 
     def test_objects_nested(self):
         """
@@ -413,52 +423,6 @@ def test_objects__all_related__deep(self):
         self.assertIsInstance(ticket.concert.band_1.manager, Manager)
         self.assertIsInstance(ticket.concert.band_2.manager, Manager)
 
-    def test_objects_nested_with_load_json(self):
-        """
-        Make sure that nested objects works alongside ``load_json`` (i.e. the
-        JSON on nested objects gets loaded).
-
-        https://github.com/piccolo-orm/piccolo/issues/1383
-
-        """
-        instrument = (
-            Instrument.objects(Instrument.recording_studio)
-            .output(load_json=True)
-            .first()
-            .run_sync()
-        )
-        assert instrument is not None
-        self.assertDictEqual(
-            instrument.recording_studio.facilities,
-            {"restaurant": True},
-        )
-
-    def test_objects_nested_with_load_json_null(self):
-        """
-        Make sure that nested objects works alongside ``load_json``, when
-        the nested object has a null value for a JSON column.
-
-        https://github.com/piccolo-orm/piccolo/issues/1391
-
-        """
-        RecordingStudio.update(
-            {
-                RecordingStudio.facilities: None,
-            },
-            force=True,
-        ).run_sync()
-
-        instrument = (
-            Instrument.objects(Instrument.recording_studio)
-            .output(load_json=True)
-            .first()
-            .run_sync()
-        )
-        assert instrument is not None
-        self.assertIsNone(
-            instrument.recording_studio.facilities,
-        )
-
     def test_objects_prefetch_clause(self):
         """
         Make sure that ``prefetch`` clause works correctly.
@@ -562,3 +526,96 @@ def test_objects_prefetch_db_column_name(self):
         signing = Signing.objects().prefetch(Signing.with_).first().run_sync()
         assert signing is not None
         self.assertIsInstance(signing.with_, Band)
+
+
+class TestObjectsNestedLoadJSON(TestJoin):
+    def test_objects_nested_with_load_json(self):
+        """
+        Make sure that nested objects works alongside ``load_json`` (i.e. the
+        JSON on nested objects gets loaded).
+
+        https://github.com/piccolo-orm/piccolo/issues/1383
+
+        """
+        instrument = (
+            Instrument.objects(Instrument.recording_studio)
+            .output(load_json=True)
+            .first()
+            .run_sync()
+        )
+        assert instrument is not None
+        self.assertDictEqual(
+            instrument.recording_studio.facilities,
+            {"restaurant": True},
+        )
+
+    def test_objects_nested_with_load_json_null(self):
+        """
+        Make sure that nested objects works alongside ``load_json``, when
+        the nested object has a null value for a JSON column.
+
+        https://github.com/piccolo-orm/piccolo/issues/1391
+
+        """
+        RecordingStudio.update(
+            {
+                RecordingStudio.facilities: None,
+            },
+            force=True,
+        ).run_sync()
+
+        instrument = (
+            Instrument.objects(Instrument.recording_studio)
+            .output(load_json=True)
+            .first()
+            .run_sync()
+        )
+        assert instrument is not None
+        self.assertIsNone(
+            instrument.recording_studio.facilities,
+        )
+
+
+class TableA(Table):
+    id = Serial(primary_key=True)
+    name = Varchar()
+
+
+class TableB(Table):
+    data = JSON()
+    table_a = ForeignKey(TableA, null=False)
+
+
+class TestObjectsNestedLoadJSONArray(AsyncTableTest):
+    tables = [TableA, TableB]
+
+    async def test_objects_nested_with_load_json_array(self):
+        """
+        Make sure that nested objects works alongside ``load_json``, when
+        the parent object has a JSON column, and it contains an array value.
+
+        https://github.com/piccolo-orm/piccolo/issues/1399#issuecomment-4864444701
+
+        The problem was because we already loaded the JSON string on the top
+        level row, and then were trying to load it again in
+        `make_nested_object`.
+
+        """  # noqa: E501
+        table_a = TableA()
+        await table_a.save()
+
+        data = ["a", "b", "c"]
+        table_b = TableB(data=data, table_a=table_a.id)
+        await table_b.save()
+
+        # Without prefetch:
+        row = await TableB.objects().output(load_json=True).first()
+        assert row
+        assert row.data == data
+
+        # With prefetch:
+        row = (
+            await TableB.objects(TableB.table_a).output(load_json=True).first()
+        )
+        assert row
+        assert row.data == data