add test for nested arrow functions

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",

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):

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}'}])