implement order_by (#70)

Co-authored-by: Amin Alaee <[email protected]>

Author: AndTheDaysGoBy

README.md

@@ -88,6 +88,10 @@ notes = await Note.objects.exclude(completed=False).all()
 # exact, iexact, contains, icontains, lt, lte, gt, gte, in
 notes = await Note.objects.filter(text__icontains="mum").all()
 
+# .order_by()
+# order by ascending name and descending id
+notes = await Note.objects.order_by("name", "-id").all()
+
 # .get()
 note = await Note.objects.get(id=1)
 
@@ -100,6 +104,7 @@ await note.delete()
 # 'pk' always refers to the primary key
 note = await Note.objects.get(pk=2)
 note.pk  # 2
+
 ```
 
 ORM supports loading and filtering across foreign keys...

orm/models.py

@@ -56,12 +56,14 @@ def __init__(
         select_related=None,
         limit_count=None,
         offset=None,
+        order_by=None,
     ):
         self.model_cls = model_cls
         self.filter_clauses = [] if filter_clauses is None else filter_clauses
         self._select_related = [] if select_related is None else select_related
         self.limit_count = limit_count
         self.query_offset = offset
+        self._order_by = [] if order_by is None else order_by
 
     def __get__(self, instance, owner):
         return self.__class__(model_cls=owner)
@@ -96,6 +98,10 @@ def build_select_expression(self):
                 clause = sqlalchemy.sql.and_(*self.filter_clauses)
             expr = expr.where(clause)
 
+        if self._order_by:
+            order_by = list(map(self._prepare_order_by, self._order_by))
+            expr = expr.order_by(*order_by)
+
         if self.limit_count:
             expr = expr.limit(self.limit_count)
 
@@ -186,6 +192,7 @@ def _filter_query(self, _exclude: bool = False, **kwargs):
             select_related=select_related,
             limit_count=self.limit_count,
             offset=self.query_offset,
+            order_by=self._order_by,
         )
 
     def select_related(self, related):
@@ -199,6 +206,17 @@ def select_related(self, related):
             select_related=related,
             limit_count=self.limit_count,
             offset=self.query_offset,
+            order_by=self._order_by,
+        )
+
+    def order_by(self, *order_by):
+        return self.__class__(
+            model_cls=self.model_cls,
+            filter_clauses=self.filter_clauses,
+            select_related=self._select_related,
+            limit_count=self.limit_count,
+            offset=self.query_offset,
+            order_by=order_by,
         )
 
     async def exists(self) -> bool:
@@ -213,6 +231,7 @@ def limit(self, limit_count: int):
             select_related=self._select_related,
             limit_count=limit_count,
             offset=self.query_offset,
+            order_by=self._order_by,
         )
 
     def offset(self, offset: int):
@@ -222,6 +241,7 @@ def offset(self, offset: int):
             select_related=self._select_related,
             limit_count=self.limit_count,
             offset=offset,
+            order_by=self._order_by,
         )
 
     async def count(self) -> int:
@@ -285,6 +305,12 @@ async def create(self, **kwargs):
         instance.pk = await self.database.execute(expr)
         return instance
 
+    def _prepare_order_by(self, order_by: str):
+        reverse = order_by.startswith("-")
+        order_by = order_by.lstrip("-")
+        order_col = self.table.columns[order_by]
+        return order_col.desc() if reverse else order_col
+
 
 class Model(typesystem.Schema, metaclass=ModelMetaclass):
     __abstract__ = True

tests/test_models.py

@@ -156,6 +156,41 @@ async def test_model_filter():
         assert await products.count() == 3
 
 
+async def test_model_order_by():
+    async with database:
+        await User.objects.create(name="Bob")
+        await User.objects.create(name="Allen")
+        await User.objects.create(name="Bob")
+
+        users = await User.objects.order_by("name").all()
+        assert users[0].name == "Allen"
+        assert users[1].name == "Bob"
+
+        users = await User.objects.order_by("-name").all()
+        assert users[1].name == "Bob"
+        assert users[2].name == "Allen"
+
+        users = await User.objects.order_by("name", "-id").all()
+        assert users[0].name == "Allen"
+        assert users[0].id == 2
+        assert users[1].name == "Bob"
+        assert users[1].id == 3
+
+        users = await User.objects.filter(name="Bob").order_by("-id").all()
+        assert users[0].name == "Bob"
+        assert users[0].id == 3
+        assert users[1].name == "Bob"
+        assert users[1].id == 1
+
+        users = await User.objects.order_by("id").limit(1).all()
+        assert users[0].name == "Bob"
+        assert users[0].id == 1
+
+        users = await User.objects.order_by("id").limit(1).offset(1).all()
+        assert users[0].name == "Allen"
+        assert users[0].id == 2
+
+
 async def test_model_exists():
     async with database:
         await User.objects.create(name="Tom")