added support for aliases (SELECT foo as bar)

Author: dantownsend

piccolo/columns/base.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from abc import ABCMeta, abstractmethod
+import copy
 from dataclasses import dataclass, field
 import datetime
 import decimal
@@ -226,6 +227,8 @@ def __init__(
             required=required,
         )
 
+        self.alias: t.Optional[str] = None
+
     def _validate_default(
         self,
         default: t.Any,
@@ -306,6 +309,20 @@ def __ne__(self, value) -> Where:  # type: ignore
     def __hash__(self):
         return hash(self._meta.name)
 
+    def as_alias(self, name: str) -> Column:
+        """
+        Allows column names to be changed in the result of a select.
+
+        For example:
+
+        >>> await Band.select(Band.name.as_alias('title')).run()
+        {'title': 'Pythonistas'}
+
+        """
+        column = copy.deepcopy(self)
+        column.alias = name
+        return column
+
     def get_default_value(self) -> t.Any:
         """
         If the column has a default attribute, return it. If it's callable,
@@ -323,7 +340,14 @@ def get_select_string(self, engine_type: str, just_alias=False) -> str:
         """
         How to refer to this column in a SQL query.
         """
-        return self._meta.get_full_name(just_alias=just_alias)
+        if self.alias is None:
+            return self._meta.get_full_name(just_alias=just_alias)
+        else:
+            original_name = self._meta.get_full_name(just_alias=True)
+            return f"{original_name} AS {self.alias}"
+
+    def get_where_string(self, engine_type: str) -> str:
+        return self.get_select_string(engine_type=engine_type, just_alias=True)
 
     def get_sql_value(self, value: t.Any) -> t.Any:
         """

piccolo/columns/column_types.py

@@ -1094,10 +1094,11 @@ def arrow(self, key: str) -> JSON:
         return self
 
     def get_select_string(self, engine_type: str, just_alias=False) -> str:
-        select_string = super().get_select_string(
-            engine_type=engine_type, just_alias=just_alias
-        )
+        select_string = self._meta.get_full_name(just_alias=just_alias)
         if self.json_operator is None:
             return select_string
         else:
-            return f"{select_string} {self.json_operator} AS {self._meta.name}"
+            if self.alias is None:
+                return f"{select_string} {self.json_operator}"
+            else:
+                return f"{select_string} {self.json_operator} AS {self.alias}"

piccolo/columns/combination.py

@@ -60,7 +60,7 @@ class Where(CombinableMixin):
 
     def __init__(
         self,
-        column: "Column",
+        column: Column,
         value: t.Any = UNDEFINED,
         values: t.Union[Iterable, Undefined] = UNDEFINED,
         operator: t.Type[ComparisonOperator] = ComparisonOperator,
@@ -91,7 +91,9 @@ def querystring(self) -> QueryString:
             args.append(self.values_querystring)
 
         template = self.operator.template.format(
-            name=self.column._meta.get_full_name(just_alias=True),
+            name=self.column.get_where_string(
+                engine_type=self.column._meta.engine_type, just_alias=True
+            ),
             value="{}",
             values="{}",
         )

piccolo/query/methods/select.py

@@ -248,7 +248,9 @@ def querystrings(self) -> t.Sequence[QueryString]:
 
         #######################################################################
 
-        select = "SELECT DISTINCT" if self.distinct else "SELECT"
+        select = (
+            "SELECT DISTINCT" if self.distinct_delegate._distinct else "SELECT"
+        )
         query = f"{select} {columns_str} FROM {self.table._meta.tablename}"
 
         for join in joins:

tests/columns/test_jsonb.py

@@ -11,7 +11,7 @@ class MyTable(Table):
 
 
 @postgres_only
-class TestJSON(TestCase):
+class TestJSONB(TestCase):
     def setUp(self):
         MyTable.create_table().run_sync()
 
@@ -32,4 +32,29 @@ def test_arrow(self):
         """
         MyTable(json='{"a": 1}').save().run_sync()
         row = MyTable.select(MyTable.json.arrow("a")).first().run_sync()
-        self.assertEqual(row["json"], "1")
+        self.assertEqual(row["?column?"], "1")
+
+    def test_arrow_as_alias(self):
+        """
+        Test using the arrow function to retrieve a subset of the JSON.
+        """
+        MyTable(json='{"a": 1}').save().run_sync()
+        row = (
+            MyTable.select(MyTable.json.arrow("a").as_alias("a"))
+            .first()
+            .run_sync()
+        )
+        self.assertEqual(row["a"], "1")
+
+    def test_arrow_where(self):
+        """
+        Make sure the arrow function can be used within a WHERE clause.
+        """
+        MyTable(json='{"a": 1}').save().run_sync()
+        self.assertEqual(
+            MyTable.count().where(MyTable.json.arrow("a") == "1").run_sync(), 1
+        )
+
+        self.assertEqual(
+            MyTable.count().where(MyTable.json.arrow("a") == "2").run_sync(), 0
+        )

tests/table/test_select.py

@@ -285,6 +285,14 @@ def test_distinct(self):
         self.insert_rows()
         self.insert_rows()
 
+        response = (
+            Band.select(Band.name).where(Band.name == "Pythonistas").run_sync()
+        )
+
+        self.assertTrue(
+            response == [{"name": "Pythonistas"}, {"name": "Pythonistas"}]
+        )
+
         response = (
             Band.select(Band.name)
             .where(Band.name == "Pythonistas")
@@ -410,9 +418,43 @@ def test_call_chain(self):
         """
         Make sure the call chain lengths are the correct size.
         """
-        # self.assertEqual(len(Concert.band_1.name._meta.call_chain), 1)
+        self.assertEqual(len(Concert.band_1.name._meta.call_chain), 1)
         self.assertEqual(len(Concert.band_1.manager.name._meta.call_chain), 2)
 
+    def test_as_alias(self):
+        """
+        Make sure we can specify aliases for the columns.
+        """
+        self.insert_row()
+        response = Band.select(Band.name.as_alias("title")).run_sync()
+        self.assertEqual(response, [{"title": "Pythonistas"}])
+
+    def test_as_alias_with_join(self):
+        """
+        Make sure we can specify aliases for the column, when performing a
+        join.
+        """
+        self.insert_row()
+        response = Band.select(
+            Band.manager.name.as_alias("manager_name")
+        ).run_sync()
+        self.assertEqual(response, [{"manager_name": "Guido"}])
+
+    def test_as_alias_with_where_clause(self):
+        """
+        Make sure we can specify aliases for the column, when the column is
+        also being used in a where clause.
+        """
+        self.insert_row()
+        response = (
+            Band.select(Band.name, Band.manager.name.as_alias("manager_name"))
+            .where(Band.manager.name == "Guido")
+            .run_sync()
+        )
+        self.assertEqual(
+            response, [{"name": "Pythonistas", "manager_name": "Guido"}]
+        )
+
 
 class TestSelectSecret(TestCase):
     def setUp(self):