simplify alias handling for Function

piccolo-orm · May 29, 2024 · 48fabb0 · 48fabb0
1 parent 3000aed
commit 48fabb0
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 9 deletions.
diff --git a/piccolo/query/functions/base.py b/piccolo/query/functions/base.py
@@ -13,23 +13,29 @@ def __init__(
         identifier: t.Union[Column, QueryString, str],
         alias: t.Optional[str] = None,
     ):
-        self._alias = alias
-
         if isinstance(identifier, Column):
-            if not alias:
-                self._alias = identifier._meta.get_default_alias()
-
             # We track any columns just in case we need to perform joins
             self.columns = [identifier, *getattr(self, "columns", [])]
 
             column_full_name = identifier._meta.get_full_name(with_alias=False)
-            super().__init__(f"{self.function_name}({column_full_name})")
+            super().__init__(
+                f"{self.function_name}({column_full_name})",
+                alias=alias,
+            )
         elif isinstance(identifier, QueryString):
             if identifier._alias:
                 self._alias = identifier._alias
 
-            super().__init__(f"{self.function_name}({{}})", identifier)
+            super().__init__(
+                f"{self.function_name}({{}})",
+                identifier,
+                alias=alias,
+            )
         elif isinstance(identifier, str):
-            super().__init__(f"{self.function_name}({{}})", identifier)
+            super().__init__(
+                f"{self.function_name}({{}})",
+                identifier,
+                alias=alias,
+            )
         else:
             raise ValueError("Unrecognised identifier type")
diff --git a/piccolo/querystring.py b/piccolo/querystring.py
@@ -75,6 +75,7 @@ def __init__(
         *args: t.Any,
         query_type: str = "generic",
         table: t.Optional[t.Type[Table]] = None,
+        alias: t.Optional[str] = None,
     ) -> None:
         """
         :param template:
@@ -100,7 +101,11 @@ def __init__(
         self._frozen_compiled_strings: t.Optional[
             t.Tuple[str, t.List[t.Any]]
         ] = None
-        self._alias = None
+        self._alias = alias
+
+    def as_alias(self, alias: str) -> QueryString:
+        self._alias = alias
+        return self
 
     def __str__(self):
         """