Home > Software engineering >  Sort Django queryset with custom order
Sort Django queryset with custom order

Time:06-18

Let's say I have the following Django model:

class Todo(models.Model):
    class Priority(models.IntegerChoices):
        HIGH = 1, "High"
        LOW = 2, "Low"
        MEDIUM = 3, "Medium"

    title = models.CharField(max_length=255)
    priority = models.PositiveSmallIntegerField(choices=Priority.choices, db_index=True)

And the following test data via factory-boy:

import factory
from factory.django import DjangoModelFactory

class TodoFactory(DjangoModelFactory):
    class Meta:
        model = Todo

    title = factory.Sequence(lambda n: f"Todo {n}")

def test_custom_db_sort():
    todo1 = TodoFactory(priority=Todo.Priority.MEDIUM)
    todo2 = TodoFactory(priority=Todo.Priority.HIGH)
    todo3 = TodoFactory(priority=Todo.Priority.LOW)
    todo4 = TodoFactory(priority=Todo.Priority.MEDIUM)
    todo5 = TodoFactory(priority=Todo.Priority.HIGH)

I want to fetch the Todo objects from the database and sort from HIGH -> LOW priority. Unfortunately Todo.Priority isn't ordered correctly so I can't sort on that. I've come up with the following solution:

PREFERENCE = {
    Priority.HIGH: 1,
    Priority.MEDIUM: 2,
    Priority.LOW: 3,
}
result = sorted(
    Todo.objects.all(),
    key=lambda x: [PREFERENCE[x.priority], x.id]
)

This works but I'd rather sort in the database for several reasons (e.g. faster, better to do work in the DB rather than Python). Is there anyway to perform this query in the database?

CodePudding user response:

Use Conditional expression (Case(...), When(...)) expressions

from django.db.models import Case, When, Value

new_priority = Case(
    When(priority=Todo.Priority.HIGH, then=Value(1)),
    When(priority=Todo.Priority.MEDIUM, then=Value(2)),
    When(priority=Todo.Priority.LOW, then=Value(3)),
)

qs = Todo.objects.annotate(new_priority=new_priority).order_by("new_priority")

If I have a choice, I would rather update the Priority choice values instead of annotating the expression.

  • Related