diff --git a/piccolo/query/functions/__init__.py b/piccolo/query/functions/__init__.py index 3163f6d1..7d51dca2 100644 --- a/piccolo/query/functions/__init__.py +++ b/piccolo/query/functions/__init__.py @@ -1,11 +1,13 @@ from .aggregate import Avg, Count, Max, Min, Sum from .datetime import Day, Extract, Hour, Month, Second, Strftime, Year +from .json import Arrow from .math import Abs, Ceil, Floor, Round from .string import Concat, Length, Lower, Ltrim, Reverse, Rtrim, Upper from .type_conversion import Cast __all__ = ( "Abs", + "Arrow", "Avg", "Cast", "Ceil", diff --git a/piccolo/query/functions/json.py b/piccolo/query/functions/json.py index b4c9f4d5..4805698d 100644 --- a/piccolo/query/functions/json.py +++ b/piccolo/query/functions/json.py @@ -17,7 +17,12 @@ class Arrow(QueryString): ``.output(load_json=True)``. """ - def __init__(self, column: JSONB, key: str, alias: t.Optional[str] = None): + def __init__( + self, + column: JSONB | QueryString, + key: str, + alias: t.Optional[str] = None, + ): super().__init__("{} -> {}", column, key, alias=alias) def clean_value(self, value: t.Any): diff --git a/tests/query/functions/test_json.py b/tests/query/functions/test_json.py new file mode 100644 index 00000000..3ec3c4c3 --- /dev/null +++ b/tests/query/functions/test_json.py @@ -0,0 +1,25 @@ +from piccolo.columns import JSONB +from piccolo.query.functions.json import Arrow +from piccolo.table import Table +from piccolo.testing.test_case import AsyncTableTest + + +class RecordingStudio(Table): + facilities = JSONB(null=True) + + +class TestArrow(AsyncTableTest): + + tables = [RecordingStudio] + + async def test_nested(self): + await RecordingStudio( + {RecordingStudio.facilities: {"a": {"b": {"c": 1}}}} + ).save() + + response = await RecordingStudio.select( + Arrow(Arrow(RecordingStudio.facilities, "a"), "b").as_alias( + "b_value" + ) + ) + self.assertListEqual(response, [{"b_value": '{"c": 1}'}])