Add DB indexes and make run_after not null

Author: jonatron

django_tasks/backends/database/migrations/0016_alter_dbtaskresult_options_and_more.py

@@ -0,0 +1,45 @@
+# Generated by Django 5.1.4 on 2025-04-28 11:28
+
+import datetime
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('django_tasks_database', '0015_correctly_order_run_after'),
+    ]
+
+    operations = [
+        migrations.AlterModelOptions(
+            name='dbtaskresult',
+            options={'ordering': [models.OrderBy(models.F('priority'), descending=True), models.OrderBy(models.F('run_after'))], 'verbose_name': 'Task Result', 'verbose_name_plural': 'Task Results'},
+        ),
+        migrations.AlterField(
+            model_name='dbtaskresult',
+            name='backend_name',
+            field=models.CharField(max_length=32, verbose_name='backend name'),
+        ),
+        migrations.AlterField(
+            model_name='dbtaskresult',
+            name='queue_name',
+            field=models.CharField(default='default', max_length=32, verbose_name='queue name'),
+        ),
+        migrations.AlterField(
+            model_name='dbtaskresult',
+            name='run_after',
+            field=models.DateTimeField(default=datetime.datetime(9999, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), verbose_name='run after'),
+        ),
+        migrations.AddIndex(
+            model_name='dbtaskresult',
+            index=models.Index(models.F('status'), models.OrderBy(models.F('priority'), descending=True), models.OrderBy(models.F('run_after')), condition=models.Q(('status', 'NEW')), name='django_task_new_ordering_idx'),
+        ),
+        migrations.AddIndex(
+            model_name='dbtaskresult',
+            index=models.Index(fields=['queue_name'], name='django_task_queue_n_99c321_idx'),
+        ),
+        migrations.AddIndex(
+            model_name='dbtaskresult',
+            index=models.Index(fields=['backend_name'], name='django_task_backend_071754_idx'),
+        ),
+    ]

django_tasks/backends/database/models.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 import uuid
 from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
@@ -46,14 +47,17 @@ def __class_getitem__(cls, _):
             return cls
 
 
+DATE_MAX = datetime.datetime(9999, 1, 1, tzinfo=datetime.timezone.utc)
+
+
 class DBTaskResultQuerySet(models.QuerySet):
     def ready(self) -> "DBTaskResultQuerySet":
         """
         Return tasks which are ready to be processed.
         """
         return self.filter(
             status=ResultStatus.NEW,
-        ).filter(models.Q(run_after=None) | models.Q(run_after__lte=timezone.now()))
+        ).filter(models.Q(run_after=DATE_MAX) | models.Q(run_after__lte=timezone.now()))
 
     def succeeded(self) -> "DBTaskResultQuerySet":
         return self.filter(status=ResultStatus.SUCCEEDED)
@@ -95,10 +99,12 @@ class DBTaskResult(GenericBase[P, T], models.Model):
 
     task_path = models.TextField(_("task path"))
 
-    queue_name = models.TextField(_("queue name"), default=DEFAULT_QUEUE_NAME)
-    backend_name = models.TextField(_("backend name"))
+    queue_name = models.CharField(
+        _("queue name"), default=DEFAULT_QUEUE_NAME, max_length=32
+    )
+    backend_name = models.CharField(_("backend name"), max_length=32)
 
-    run_after = models.DateTimeField(_("run after"), null=True)
+    run_after = models.DateTimeField(_("run after"), default=DATE_MAX)
 
     return_value = models.JSONField(_("return value"), default=None, null=True)
 
@@ -108,9 +114,18 @@ class DBTaskResult(GenericBase[P, T], models.Model):
     objects = DBTaskResultQuerySet.as_manager()
 
     class Meta:
-        ordering = [F("priority").desc(), F("run_after").asc(nulls_last=True)]
+        ordering = [F("priority").desc(), F("run_after").asc()]
         verbose_name = _("Task Result")
         verbose_name_plural = _("Task Results")
+        indexes = [
+            models.Index(
+                "status", *ordering,
+                name="django_task_new_ordering_idx",
+                condition=Q(status=ResultStatus.NEW),
+            ),
+            models.Index(fields=["queue_name"]),
+            models.Index(fields=["backend_name"]),
+        ]
 
         if django.VERSION >= (5, 1):
             constraints = [
@@ -127,6 +142,11 @@ class Meta:
                 )
             ]
 
+    def save(self, **kwargs):
+        if self.run_after is None:
+            self.run_after = DATE_MAX
+        super().save(**kwargs)
+
     @property
     def task(self) -> Task[P, T]:
         task = import_string(self.task_path)
@@ -136,6 +156,9 @@ def task(self) -> Task[P, T]:
                 f"Task {self.id} does not point to a Task ({self.task_path})"
             )
 
+        if self.run_after == DATE_MAX:
+            self.run_after = None
+
         return task.using(
             priority=self.priority,
             queue_name=self.queue_name,

tests/tests/test_database_backend.py

@@ -387,6 +387,38 @@ def test_enqueue_logs(self) -> None:
         self.assertIn("enqueued", captured_logs.output[0])
         self.assertIn(result.id, captured_logs.output[0])
 
+    def test_index_scan_for_ready(self) -> None:
+        test_tasks.noop_task.enqueue()
+        db_task = DBTaskResult.objects.first()
+        new_db_tasks = []
+        for _ in range(5000):
+            new_db_task = DBTaskResult()
+            for f in DBTaskResult._meta.fields:
+                if f.name != "id":
+                    setattr(new_db_task, f.attname, getattr(db_task, f.attname))
+            new_db_tasks.append(new_db_task)
+        DBTaskResult.objects.bulk_create(new_db_tasks)
+
+        # Update query plan for certain databases
+        if connection.vendor == "postgresql":
+            with connection.cursor() as c:
+                c.execute(f"ANALYZE {DBTaskResult._meta.db_table};")
+        elif connection.vendor == "mysql":
+            with connection.cursor() as c:
+                c.execute(f"ANALYZE TABLE {DBTaskResult._meta.db_table};")
+
+        plan = DBTaskResult.objects.ready().explain()
+
+        if connection.vendor == "postgresql":
+            self.assertIn("django_task_new_ordering_idx", plan)
+        elif connection.vendor == "sqlite":
+            self.assertIn("USING INDEX django_task_new_ordering_idx", plan)
+        elif connection.vendor == "mysql":
+            self.assertIn("Index lookup", plan)
+            self.assertIn("using django_task_new_ordering_idx", plan)
+        else:
+            self.fail("Unknown database engine")
+
 
 @override_settings(
     TASKS={