order_by using Column object instead of string

Author: dantownsend

piccolo/query/base.py

@@ -88,9 +88,3 @@ def response_handler(self, response):
         the database driver.
         """
         return response
-
-    def _is_valid_column_name(self, column_name: str):
-        if column_name.startswith('-'):
-            column_name = column_name[1:]
-        if column_name not in [i._name for i in self.table.Meta.columns]:
-            raise ValueError(f"{column_name} isn't a valid column name")

piccolo/query/mixins.py

@@ -21,12 +21,13 @@ def __str__(self):
 
 @dataclasses.dataclass
 class OrderBy():
-    column_name: str
+    columns: t.Tuple[Column]
     ascending: bool
 
     def __str__(self):
         order = 'ASC' if self.ascending else 'DESC'
-        return f' ORDER BY {self.column_name} {order}'
+        columns_names = ', '.join([i._name for i in self.columns])
+        return f' ORDER BY {columns_names} {order}'
 
 
 @dataclasses.dataclass
@@ -56,15 +57,8 @@ def __init__(self):
         super().__init__()
         self._order_by: t.Optional[OrderBy] = None
 
-    def order_by(self, column_name: str):
-        self._is_valid_column_name(column_name)
-
-        ascending = True
-        if column_name.startswith('-'):
-            ascending = False
-            column_name = column_name[1:]
-
-        self._order_by = OrderBy(column_name, ascending)
+    def order_by(self, *columns: Column, ascending=True):
+        self._order_by = OrderBy(columns, ascending)
         return self
 
 

tests/table/test_select.py

@@ -71,7 +71,7 @@ def test_where_greater_equal_than(self):
         ).where(
             Band.popularity >= 1000
         ).order_by(
-            'name'
+            Band.name
         ).run_sync()
 
         print(f'response = {response}')
@@ -136,6 +136,8 @@ def test_where_or(self):
             Band.name
         ).where(
             (Band.name == 'Rustaceans') | (Band.name == 'CSharps')
+        ).order_by(
+            Band.name
         ).run_sync()
 
         print(f'response = {response}')
@@ -182,6 +184,8 @@ def test_complex_where(self):
         ).where(
             ((Band.popularity == 2000) & (Band.manager == 'Graydon')) |
             ((Band.popularity == 10) & (Band.manager == 'Mads'))
+        ).order_by(
+            Band.name
         ).run_sync()
 
         print(f'response = {response}')
@@ -197,7 +201,7 @@ def test_limit(self):
         response = Band.select.columns(
             Band.name
         ).order_by(
-            'name'
+            Band.name
         ).limit(
             1
         ).run_sync()
@@ -215,7 +219,7 @@ def test_first(self):
         response = Band.select.columns(
             Band.name
         ).order_by(
-            'name'
+            Band.name
         ).first().run_sync()
 
         print(f'response = {response}')
@@ -231,7 +235,7 @@ def test_order_by_ascending(self):
         response = Band.select.columns(
             Band.name
         ).order_by(
-            'name'
+            Band.name
         ).limit(1).run_sync()
 
         print(f'response = {response}')
@@ -247,7 +251,8 @@ def test_order_by_decending(self):
         response = Band.select.columns(
             Band.name
         ).order_by(
-            '-name'
+            Band.name,
+            ascending=False
         ).limit(1).run_sync()
 
         print(f'response = {response}')