From 63f269f1df3e6f3de09ce6da738febff0e022814 Mon Sep 17 00:00:00 2001
From: Daniel Townsend <dan@dantownsend.co.uk>
Date: Fri, 18 Oct 2024 22:34:34 +0100
Subject: [PATCH] improve arrow function

---
 piccolo/columns/column_types.py | 33 +++------------------------------
 piccolo/querystring.py          |  6 ++++++
 2 files changed, 9 insertions(+), 30 deletions(-)

diff --git a/piccolo/columns/column_types.py b/piccolo/columns/column_types.py
index d16329b49..229295c66 100644
--- a/piccolo/columns/column_types.py
+++ b/piccolo/columns/column_types.py
@@ -2352,40 +2352,13 @@ class JSONB(JSON):
     def column_type(self):
         return "JSONB"  # Must be defined, we override column_type() in JSON()
 
-    def arrow(self, key: str) -> JSONB:
+    def arrow(self, key: str) -> QueryString:
         """
         Allows part of the JSON structure to be returned - for example,
         for {"a": 1}, and a key value of "a", then 1 will be returned.
         """
-        instance = t.cast(JSONB, self.copy())
-        instance.json_operator = f"-> '{key}'"
-        return instance
-
-    def get_select_string(
-        self, engine_type: str, with_alias: bool = True
-    ) -> QueryString:
-        select_string = self._meta.get_full_name(with_alias=False)
-
-        if self.json_operator is not None:
-            select_string += f" {self.json_operator}"
-
-        if with_alias:
-            alias = self._alias or self._meta.get_default_alias()
-            select_string += f' AS "{alias}"'
-
-        return QueryString(select_string)
-
-    def eq(self, value) -> Where:
-        """
-        See ``Boolean.eq`` for more details.
-        """
-        return self.__eq__(value)
-
-    def ne(self, value) -> Where:
-        """
-        See ``Boolean.ne`` for more details.
-        """
-        return self.__ne__(value)
+        alias = self._alias or self._meta.get_default_alias()
+        return QueryString("{} -> {}", self, key, alias=alias)
 
     ###########################################################################
     # Descriptors
diff --git a/piccolo/querystring.py b/piccolo/querystring.py
index 22f5f215a..ae8748117 100644
--- a/piccolo/querystring.py
+++ b/piccolo/querystring.py
@@ -264,6 +264,12 @@ def __eq__(self, value) -> QueryString:  # type: ignore[override]
     def __ne__(self, value) -> QueryString:  # type: ignore[override]
         return QueryString("{} != {}", self, value)
 
+    def eq(self, value) -> QueryString:
+        return self.__eq__(value)
+
+    def ne(self, value) -> QueryString:
+        return self.__ne__(value)
+
     def __add__(self, value) -> QueryString:
         return QueryString("{} + {}", self, value)