added .output(load_json=True) option

Author: dantownsend

piccolo/query/base.py

@@ -3,12 +3,15 @@
 from time import time
 import typing as t
 
+from piccolo.columns.column_types import JSON, JSONB
+from piccolo.query.mixins import ColumnsDelegate
 from piccolo.querystring import QueryString
 from piccolo.utils.sync import run_sync
-from piccolo.utils.encoding import dump_json
+from piccolo.utils.encoding import dump_json, load_json
 
 if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.table import Table  # noqa
+    from piccolo.query.mixins import OutputDelegate
 
 
 class Timer:
@@ -51,9 +54,54 @@ async def _process_results(self, results):
         if hasattr(self, "run_callback"):
             self.run_callback(raw)
 
-        raw = await self.response_handler(raw)
+        output: t.Optional[OutputDelegate] = getattr(
+            self, "output_delegate", None
+        )
+
+        #######################################################################
+
+        if output and output._output.load_json:
+            columns_delegate: t.Optional[ColumnsDelegate] = getattr(
+                self, "columns_delegate", None
+            )
+
+            if columns_delegate is not None:
+                json_columns = [
+                    i
+                    for i in columns_delegate.selected_columns
+                    if isinstance(i, (JSON, JSONB))
+                ]
+            else:
+                json_columns = self.table._meta.json_columns
+
+            json_column_names = []
+            for column in json_columns:
+                if column.alias is not None:
+                    json_column_names.append(column.alias)
+                elif len(column._meta.call_chain) > 0:
+                    json_column_names.append(
+                        column.get_select_string(
+                            engine_type=column._meta.engine_type
+                        )
+                    )
+                else:
+                    json_column_names.append(column._meta.name)
+
+            processed_raw = []
+
+            for row in raw:
+                new_row = {**row}
+                for json_column_name in json_column_names:
+                    value = new_row.get(json_column_name)
+                    if value is not None:
+                        new_row[json_column_name] = load_json(value)
+                processed_raw.append(new_row)
 
-        output = getattr(self, "output_delegate", None)
+            raw = processed_raw
+
+        #######################################################################
+
+        raw = await self.response_handler(raw)
 
         if output:
             if output._output.as_objects:

piccolo/query/methods/objects.py

@@ -44,6 +44,12 @@ def __init__(self, table: t.Type[Table], **kwargs):
         self.output_delegate._output.as_objects = True
         self.where_delegate = WhereDelegate()
 
+    def output(self, load_json: bool = False) -> Objects:
+        self.output_delegate.output(
+            as_list=False, as_json=False, load_json=load_json
+        )
+        return self
+
     def limit(self, number: int) -> Objects:
         self.limit_delegate.limit(number)
         return self

piccolo/query/methods/select.py

@@ -129,8 +129,15 @@ def order_by(self, *columns: Column, ascending=True) -> Select:
         self.order_by_delegate.order_by(*_columns, ascending=ascending)
         return self
 
-    def output(self, as_list: bool = False, as_json: bool = False) -> Select:
-        self.output_delegate.output(as_list=as_list, as_json=as_json)
+    def output(
+        self,
+        as_list: bool = False,
+        as_json: bool = False,
+        load_json: bool = False,
+    ) -> Select:
+        self.output_delegate.output(
+            as_list=as_list, as_json=as_json, load_json=load_json
+        )
         return self
 
     def where(self, where: Combinable) -> Select:

piccolo/query/mixins.py

@@ -76,12 +76,14 @@ class Output:
     as_json: bool = False
     as_list: bool = False
     as_objects: bool = False
+    load_json: bool = False
 
     def copy(self) -> Output:
         return self.__class__(
             as_json=self.as_json,
             as_list=self.as_list,
             as_objects=self.as_objects,
+            load_json=self.load_json,
         )
 
 
@@ -191,6 +193,7 @@ def output(
         self,
         as_list: t.Optional[bool] = None,
         as_json: t.Optional[bool] = None,
+        load_json: t.Optional[bool] = None,
     ):
         """
         :param as_list:
@@ -199,13 +202,19 @@ def output(
         :param as_json:
             The results are serialised into JSON. It's equivalent to running
             `json.dumps` on the result.
+        :param load_json:
+            If True, any JSON fields will have the JSON values returned from
+            the database loaded as Python objects.
         """
         if as_list is not None:
             self._output.as_list = bool(as_list)
 
-        if type(as_json) is bool:
+        if as_json is not None:
             self._output.as_json = bool(as_json)
 
+        if load_json is not None:
+            self._output.load_json = bool(load_json)
+
     def copy(self) -> OutputDelegate:
         _output = self._output.copy() if self._output is not None else None
         return self.__class__(_output=_output)

piccolo/table.py

@@ -217,6 +217,7 @@ def __init_subclass__(
             default_columns=default_columns,
             non_default_columns=non_default_columns,
             foreign_key_columns=foreign_key_columns,
+            json_columns=json_columns,
             secret_columns=secret_columns,
             tags=tags,
             help_text=help_text,

tests/table/test_metaclass.py

@@ -1,10 +1,10 @@
-from piccolo.columns.column_types import ForeignKey
+from piccolo.columns.column_types import ForeignKey, JSON, JSONB
 from unittest import TestCase
 
 from piccolo.columns import Secret
 from piccolo.table import Table
 
-from ..example_app.tables import Band
+from tests.example_app.tables import Band
 
 
 class TestMetaClass(TestCase):
@@ -63,3 +63,16 @@ class Classified(Table):
         self.assertEqual(
             Classified._meta.secret_columns, [Classified.top_secret]
         )
+
+    def test_json_columns(self):
+        """
+        Make sure TableMeta.json_columns are setup correctly.
+        """
+
+        class MyTable(Table):
+            column_a = JSON()
+            column_b = JSONB()
+
+        self.assertEqual(
+            MyTable._meta.json_columns, [MyTable.column_a, MyTable.column_b]
+        )

tests/table/test_output.py

@@ -1,7 +1,8 @@
+from unittest import TestCase
 import json
 
-from ..base import DBTestCase
-from ..example_app.tables import Band
+from tests.base import DBTestCase
+from tests.example_app.tables import Band, RecordingStudio
 
 
 class TestOutputList(DBTestCase):
@@ -28,3 +29,33 @@ def test_output_as_json(self):
         response = Band.select(Band.name).output(as_json=True).run_sync()
 
         self.assertTrue(json.loads(response) == [{"name": "Pythonistas"}])
+
+
+class TestOutputLoadJSON(TestCase):
+    def setUp(self):
+        RecordingStudio.create_table().run_sync()
+
+    def tearDown(self):
+        RecordingStudio.alter().drop_table().run_sync()
+
+    def test_select(self):
+        json = {"a": 123}
+
+        RecordingStudio(facilities=json, facilities_b=json).save().run_sync()
+
+        results = RecordingStudio.select().output(load_json=True).run_sync()
+
+        self.assertEqual(
+            results,
+            [{"id": 1, "facilities": {"a": 123}, "facilities_b": {"a": 123}}],
+        )
+
+    def test_objects(self):
+        json = {"a": 123}
+
+        RecordingStudio(facilities=json, facilities_b=json).save().run_sync()
+
+        results = RecordingStudio.objects().output(load_json=True).run_sync()
+
+        self.assertEqual(results[0].facilities, json)
+        self.assertEqual(results[0].facilities_b, json)