Repository: piccolo-orm/piccolo
Issue: Fix `.as_alias()` bug with `JSONB` and joins

### Discussed in https://github.com/piccolo-orm/piccolo/discussions/551

<div type='discussions-op-text'>

<sup>Originally posted by **Anton-Karpenko** June 21, 2022</sup>
I struggle to apply `as_alias()` to a field from a joined table.

Here is an example:

```
class User(Table, table_name="user"):
    name = Varchar(length=120)
    config = JSONB(default={})

    
class Subscriber(Table, table_name="subscriber"):
    name = Varchar(length=120)
    user = ForeignKey(references=User)
```

Here is my query:

```
query = Subscriber.select(Subscriber.name.as_alias("subscriber_name"), Subscriber.user.config.as_alias("config"))
print(query)
```
The output is:
```
SELECT "subscriber"."name" AS subscriber_name,  "subscriber$user"."config" AS "user$config" AS config
FROM subscriber LEFT JOIN user subscriber$user ON (subscriber.user = subscriber$user.id)
```

Here is AS coupled, and I have an obvious error:
```
asyncpg.exceptions.PostgresSyntaxError: syntax error at or near "AS"
```

</div>

Changed files:
piccolo/columns/column_types.py, piccolo/query/methods/select.py, piccolo/columns/base.py, tests/columns/test_jsonb.py, tests/conftest.py, piccolo/columns/readable.py, piccolo/query/base.py, piccolo/apps/playground/commands/run.py, piccolo/table.py, piccolo/query/mixins.py, piccolo/columns/m2m.py

diff --git a/piccolo/apps/playground/commands/run.py b/piccolo/apps/playground/commands/run.py
index 598fb7e1..582a8763 100644
--- a/piccolo/apps/playground/commands/run.py
+++ b/piccolo/apps/playground/commands/run.py
@@ -6,6 +6,7 @@ import datetime
 import sys
 import uuid
 from decimal import Decimal
+from enum import Enum
 
 from piccolo.columns import (
     JSON,
@@ -47,8 +48,14 @@ class Concert(Table):
 
 
 class Ticket(Table):
+    class TicketType(Enum):
+        sitting = "sitting"
+        standing = "standing"
+        premium = "premium"
+
     concert = ForeignKey(Concert)
     price = Numeric(digits=(5, 2))
+    ticket_type = Varchar(choices=TicketType, default=TicketType.standing)
 
 
 class DiscountCode(Table):
diff --git a/piccolo/columns/base.py b/piccolo/columns/base.py
index 5a92322f..0cd435d3 100644
--- a/piccolo/columns/base.py
+++ b/piccolo/columns/base.py
@@ -203,6 +203,8 @@ class ColumnMeta:
     call_chain: t.List["ForeignKey"] = field(default_factory=list)
     table_alias: t.Optional[str] = None
 
+    ###########################################################################
+
     @property
     def engine_type(self) -> str:
         engine = self.table._meta.db
@@ -234,21 +236,52 @@ class ColumnMeta:
 
         return output
 
+    ###########################################################################
+
+    def get_default_alias(self):
+        column_name = self.db_column_name
+
+        if self.call_chain:
+            column_name = (
+                "$".join(
+                    t.cast(str, i._meta.db_column_name)
+                    for i in self.call_chain
+                )
+                + f"${column_name}"
+            )
+
+        return column_name
+
+    def _get_path(self, include_quotes: bool = False):
+        column_name = self.db_column_name
+
+        if self.call_chain:
+            table_alias = self.call_chain[-1]._meta.table_alias
+            if include_quotes:
+                return f'"{table_alias}"."{column_name}"'
+            else:
+                return f"{table_alias}.{column_name}"
+        else:
+            if include_quotes:
+                return f'"{self.table._meta.tablename}"."{column_name}"'
+            else:
+                return f"{self.table._meta.tablename}.{column_name}"
+
     def get_full_name(
-        self, just_alias: bool = False, include_quotes: bool = False
+        self, with_alias: bool = True, include_quotes: bool = True
     ) -> str:
         """
         Returns the full column name, taking into account joins.
 
-        :param just_alias:
+        :param with_alias:
             Examples:
 
             .. code-block python::
 
-                >>> Band.manager.name._meta.get_full_name(just_alias=True)
+                >>> Band.manager.name._meta.get_full_name(with_alias=False)
                 'band$manager.name'
 
-                >>> Band.manager.name._meta.get_full_name(just_alias=False)
+                >>> Band.manager.name._meta.get_full_name(with_alias=True)
                 'band$manager.name AS "manager$name"'
 
         :param include_quotes:
@@ -267,27 +300,16 @@ class ColumnMeta:
 
 
         """
-        column_name = self.db_column_name
+        full_name = self._get_path(include_quotes=include_quotes)
 
-        if not self.call_chain:
+        if with_alias and self.call_chain:
+            alias = self.get_default_alias()
             if include_quotes:
-                return f'"{self.table._meta.tablename}"."{column_name}"'
+                full_name += f' AS "{alias}"'
             else:
-                return f"{self.table._meta.tablename}.{column_name}"
-
-        column_name = (
-            "$".join(
-                t.cast(str, i._meta.db_column_name) for i in self.call_chain
-            )
-            + f"${column_name}"
-        )
-
-        if include_quotes:
-            alias = f'"{self.call_chain[-1]._meta.table_alias}"."{self.name}"'
-        else:
-            alias = f"{self.call_chain[-1]._meta.table_alias}.{self.name}"
+                full_name += f" AS {alias}"
 
-        return alias if just_alias else f'{alias} AS "{column_name}"'
+        return full_name
 
     ###########################################################################
 
@@ -321,22 +343,24 @@ class ColumnMeta:
 
 
 class Selectable(metaclass=ABCMeta):
-    alias: t.Optional[str]
+    _alias: t.Optional[str]
 
     @abstractmethod
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
         """
         In a query, what to output after the select statement - could be a
         column name, a sub query, a function etc. For a column it will be the
         column name.
         """
-        pass
+        raise NotImplementedError()
 
     def as_alias(self, alias: str) -> Selectable:
         """
         Allows column names to be changed in the result of a select.
         """
-        self.alias = alias
+        self._alias = alias
         return self
 
 
@@ -472,7 +496,7 @@ class Column(Selectable):
             secret=secret,
         )
 
-        self.alias: t.Optional[str] = None
+        self._alias: t.Optional[str] = None
 
     def _validate_default(
         self,
@@ -694,7 +718,7 @@ class Column(Selectable):
 
         """
         column = copy.deepcopy(self)
-        column.alias = name
+        column._alias = name
         return column
 
     def get_default_value(self) -> t.Any:
@@ -709,21 +733,30 @@ class Column(Selectable):
             return default() if is_callable else default  # type: ignore
         return None
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
         """
-        How to refer to this column in a SQL query.
+        How to refer to this column in a SQL query, taking account of any joins
+        and aliases.
         """
-        if self.alias is None:
-            return self._meta.get_full_name(
-                just_alias=just_alias, include_quotes=True
-            )
-        original_name = self._meta.get_full_name(
-            just_alias=True, include_quotes=True
-        )
-        return f"{original_name} AS {self.alias}"
+        if with_alias:
+            if self._alias:
+                original_name = self._meta.get_full_name(
+                    with_alias=False,
+                )
+                return f'{original_name} AS "{self._alias}"'
+            else:
+                return self._meta.get_full_name(
+                    with_alias=True,
+                )
+
+        return self._meta.get_full_name(with_alias=False)
 
     def get_where_string(self, engine_type: str) -> str:
-        return self.get_select_string(engine_type=engine_type, just_alias=True)
+        return self.get_select_string(
+            engine_type=engine_type, with_alias=False
+        )
 
     def get_sql_value(self, value: t.Any) -> t.Any:
         """
diff --git a/piccolo/columns/column_types.py b/piccolo/columns/column_types.py
index 84ee1db8..2fc76a72 100644
--- a/piccolo/columns/column_types.py
+++ b/piccolo/columns/column_types.py
@@ -2227,21 +2227,19 @@ class JSONB(JSON):
         instance.json_operator = f"-> '{key}'"
         return instance
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
-        select_string = self._meta.get_full_name(
-            just_alias=just_alias, include_quotes=True
-        )
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
+        select_string = self._meta.get_full_name(with_alias=False)
+
         if self.json_operator is not None:
-            return (
-                f"{select_string} {self.json_operator}"
-                if self.alias is None
-                else f"{select_string} {self.json_operator} AS {self.alias}"
-            )
+            select_string += f" {self.json_operator}"
 
-        if self.alias is None:
-            return select_string
-        else:
-            return f"{select_string} AS {self.alias}"
+        if with_alias:
+            alias = self._alias or self._meta.get_default_alias()
+            select_string += f' AS "{alias}"'
+
+        return select_string
 
     def eq(self, value) -> Where:
         """
@@ -2489,15 +2487,17 @@ class Array(Column):
         else:
             raise ValueError("Only integers can be used for indexing.")
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
-        select_string = self._meta.get_full_name(
-            just_alias=just_alias, include_quotes=True
-        )
+    def get_select_string(self, engine_type: str, with_alias=True) -> str:
+        select_string = self._meta.get_full_name(with_alias=False)
 
         if isinstance(self.index, int):
-            return f"{select_string}[{self.index}]"
-        else:
-            return select_string
+            select_string += f"[{self.index}]"
+
+        if with_alias:
+            alias = self._alias or self._meta.get_default_alias()
+            select_string += f' AS "{alias}"'
+
+        return select_string
 
     def any(self, value: t.Any) -> Where:
         """
diff --git a/piccolo/columns/m2m.py b/piccolo/columns/m2m.py
index 57f40307..db0d61d7 100644
--- a/piccolo/columns/m2m.py
+++ b/piccolo/columns/m2m.py
@@ -55,7 +55,7 @@ class M2MSelect(Selectable):
             for column in columns
         )
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
+    def get_select_string(self, engine_type: str, with_alias=True) -> str:
         m2m_table_name = self.m2m._meta.resolved_joining_table._meta.tablename
         m2m_relationship_name = self.m2m._meta.name
 
diff --git a/piccolo/columns/readable.py b/piccolo/columns/readable.py
index 651f5b1a..630299f9 100644
--- a/piccolo/columns/readable.py
+++ b/piccolo/columns/readable.py
@@ -24,7 +24,7 @@ class Readable(Selectable):
     @property
     def _columns_string(self) -> str:
         return ", ".join(
-            i._meta.get_full_name(just_alias=True) for i in self.columns
+            i._meta.get_full_name(with_alias=False) for i in self.columns
         )
 
     def _get_string(self, operator: str) -> str:
@@ -41,7 +41,7 @@ class Readable(Selectable):
     def postgres_string(self) -> str:
         return self._get_string(operator="FORMAT")
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
+    def get_select_string(self, engine_type: str, with_alias=True) -> str:
         try:
             return getattr(self, f"{engine_type}_string")
         except AttributeError as e:
diff --git a/piccolo/query/base.py b/piccolo/query/base.py
index b1857df8..8e6bc145 100644
--- a/piccolo/query/base.py
+++ b/piccolo/query/base.py
@@ -86,13 +86,10 @@ class Query:
 
             json_column_names = []
             for column in json_columns:
-                if column.alias is not None:
-                    json_column_names.append(column.alias)
+                if column._alias is not None:
+                    json_column_names.append(column._alias)
                 elif column.json_operator is not None:
-                    # If no alias is specified, then the default column name
-                    # that Postgres gives when using the `->` operator is
-                    # `?column?`.
-                    json_column_names.append("?column?")
+                    json_column_names.append(column._meta.name)
                 elif len(column._meta.call_chain) > 0:
                     json_column_names.append(
                         column.get_select_string(
diff --git a/piccolo/query/methods/select.py b/piccolo/query/methods/select.py
index 0c31e837..45fa675f 100644
--- a/piccolo/query/methods/select.py
+++ b/piccolo/query/methods/select.py
@@ -62,13 +62,13 @@ class Avg(Selectable):
             self.column = column
         else:
             raise ValueError("Column type must be numeric to run the query.")
-        self.alias = alias
+        self._alias = alias
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
-        column_name = self.column._meta.get_full_name(
-            just_alias=just_alias, include_quotes=True
-        )
-        return f"AVG({column_name}) AS {self.alias}"
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
+        column_name = self.column._meta.get_full_name(with_alias=False)
+        return f'AVG({column_name}) AS "{self._alias}"'
 
 
 class Count(Selectable):
@@ -100,16 +100,16 @@ class Count(Selectable):
         self, column: t.Optional[Column] = None, alias: str = "count"
     ):
         self.column = column
-        self.alias = alias
+        self._alias = alias
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
         if self.column is None:
             column_name = "*"
         else:
-            column_name = self.column._meta.get_full_name(
-                just_alias=just_alias, include_quotes=True
-            )
-        return f"COUNT({column_name}) AS {self.alias}"
+            column_name = self.column._meta.get_full_name(with_alias=False)
+        return f'COUNT({column_name}) AS "{self._alias}"'
 
 
 class Max(Selectable):
@@ -136,13 +136,13 @@ class Max(Selectable):
 
     def __init__(self, column: Column, alias: str = "max"):
         self.column = column
-        self.alias = alias
+        self._alias = alias
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
-        column_name = self.column._meta.get_full_name(
-            just_alias=just_alias, include_quotes=True
-        )
-        return f"MAX({column_name}) AS {self.alias}"
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
+        column_name = self.column._meta.get_full_name(with_alias=False)
+        return f'MAX({column_name}) AS "{self._alias}"'
 
 
 class Min(Selectable):
@@ -167,13 +167,13 @@ class Min(Selectable):
 
     def __init__(self, column: Column, alias: str = "min"):
         self.column = column
-        self.alias = alias
+        self._alias = alias
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
-        column_name = self.column._meta.get_full_name(
-            just_alias=just_alias, include_quotes=True
-        )
-        return f"MIN({column_name}) AS {self.alias}"
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
+        column_name = self.column._meta.get_full_name(with_alias=False)
+        return f'MIN({column_name}) AS "{self._alias}"'
 
 
 class Sum(Selectable):
@@ -203,13 +203,13 @@ class Sum(Selectable):
             self.column = column
         else:
             raise ValueError("Column type must be numeric to run the query.")
-        self.alias = alias
+        self._alias = alias
 
-    def get_select_string(self, engine_type: str, just_alias=False) -> str:
-        column_name = self.column._meta.get_full_name(
-            just_alias=just_alias, include_quotes=True
-        )
-        return f"SUM({column_name}) AS {self.alias}"
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
+        column_name = self.column._meta.get_full_name(with_alias=False)
+        return f'SUM({column_name}) AS "{self._alias}"'
 
 
 class Select(Query):
diff --git a/piccolo/query/mixins.py b/piccolo/query/mixins.py
index 75318aee..fe33092f 100644
--- a/piccolo/query/mixins.py
+++ b/piccolo/query/mixins.py
@@ -66,7 +66,7 @@ class OrderBy:
     def querystring(self) -> QueryString:
         order = "ASC" if self.ascending else "DESC"
         columns_names = ", ".join(
-            i._meta.get_full_name(just_alias=True) for i in self.columns
+            i._meta.get_full_name(with_alias=False) for i in self.columns
         )
 
         return QueryString(f" ORDER BY {columns_names} {order}")
@@ -438,7 +438,7 @@ class GroupBy:
     @property
     def querystring(self) -> QueryString:
         columns_names = ", ".join(
-            i._meta.get_full_name(just_alias=True) for i in self.columns
+            i._meta.get_full_name(with_alias=False) for i in self.columns
         )
 
         return QueryString(f" GROUP BY {columns_names}")
diff --git a/piccolo/table.py b/piccolo/table.py
index 3c5a5dd0..20870838 100644
--- a/piccolo/table.py
+++ b/piccolo/table.py
@@ -618,8 +618,7 @@ class Table(metaclass=TableMetaclass):
                         break
 
         alias_names = {
-            column._meta.name: getattr(column, "alias", None)
-            for column in filtered_columns
+            column._meta.name: column._alias for column in filtered_columns
         }
 
         output = {}
diff --git a/tests/columns/test_jsonb.py b/tests/columns/test_jsonb.py
index fab2b7d1..fd419179 100644
--- a/tests/columns/test_jsonb.py
+++ b/tests/columns/test_jsonb.py
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from piccolo.columns.column_types import JSONB, Varchar
+from piccolo.columns.column_types import JSONB, ForeignKey, Varchar
 from piccolo.table import Table
 from tests.base import postgres_only
 
@@ -10,12 +10,19 @@ class RecordingStudio(Table):
     facilities = JSONB(null=True)
 
 
+class Instrument(Table):
+    name = Varchar()
+    studio = ForeignKey(RecordingStudio)
+
+
 @postgres_only
 class TestJSONB(TestCase):
     def setUp(self):
         RecordingStudio.create_table().run_sync()
+        Instrument.create_table().run_sync()
 
     def tearDown(self):
+        Instrument.alter().drop_table().run_sync()
         RecordingStudio.alter().drop_table().run_sync()
 
     def test_json(self):
@@ -88,6 +95,32 @@ class TestJSONB(TestCase):
             [{"name": "Abbey Road"}],
         )
 
+    def test_as_alias_join(self):
+        """
+        Make sure that ``as_alias`` performs correctly when used via a join.
+        """
+        studio = (
+            RecordingStudio.objects()
+            .create(name="Abbey Road", facilities={"mixing_desk": True})
+            .run_sync()
+        )
+
+        Instrument.objects().create(name="Guitar", studio=studio).run_sync()
+
+        response = (
+            Instrument.select(
+                Instrument.name,
+                Instrument.studio.facilities.as_alias("studio_facilities"),
+            )
+            .output(load_json=True)
+            .run_sync()
+        )
+
+        self.assertListEqual(
+            response,
+            [{"name": "Guitar", "studio_facilities": {"mixing_desk": True}}],
+        )
+
     def test_arrow(self):
         """
         Test using the arrow function to retrieve a subset of the JSON.
@@ -103,7 +136,7 @@ class TestJSONB(TestCase):
             .first()
             .run_sync()
         )
-        self.assertEqual(row["?column?"], "true")
+        self.assertEqual(row["facilities"], "true")
 
         row = (
             RecordingStudio.select(
@@ -113,7 +146,7 @@ class TestJSONB(TestCase):
             .first()
             .run_sync()
         )
-        self.assertEqual(row["?column?"], True)
+        self.assertEqual(row["facilities"], True)
 
     def test_arrow_as_alias(self):
         """
diff --git a/tests/conftest.py b/tests/conftest.py
index 9ee3ff0c..72c20d27 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -18,6 +18,7 @@ async def drop_tables():
         "my_table",
         "recording_studio",
         "shirt",
+        "instrument",
         "mega_table",
         "small_table",
     ]: