Add STOPPING state and update when shutdown is requested

Author: j4mie

django_dbq/management/commands/worker.py

@@ -30,6 +30,9 @@ def init_signals(self):
 
     def shutdown(self, signum, frame):
         self.alive = False
+        if self.current_job:
+            self.current_job.state = Job.STATES.STOPPING
+            self.current_job.save(update_fields=["state"])
 
     def run(self):
         while self.alive:

django_dbq/migrations/0005_alter_job_state.py

@@ -0,0 +1,30 @@
+# Generated by Django 3.2rc1 on 2021-11-29 04:43
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("django_dbq", "0004_auto_20210818_0247"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="job",
+            name="state",
+            field=models.CharField(
+                choices=[
+                    ("NEW", "New"),
+                    ("READY", "Ready"),
+                    ("PROCESSING", "Processing"),
+                    ("STOPPING", "Stopping"),
+                    ("FAILED", "Failed"),
+                    ("COMPLETE", "Complete"),
+                ],
+                db_index=True,
+                default="NEW",
+                max_length=20,
+            ),
+        ),
+    ]

django_dbq/models.py

@@ -77,6 +77,7 @@ class STATES(TextChoices):
         NEW = "NEW"
         READY = "READY"
         PROCESSING = "PROCESSING"
+        STOPPING = "STOPPING"
         FAILED = "FAILED"
         COMPLETE = "COMPLETE"
 

django_dbq/tests.py

@@ -152,6 +152,19 @@ def test_process_job_previous_job_long_time_ago(self, mock_sleep):
         self.assertEqual(self.mock_worker.last_job_finished, timezone.now())
 
 
+@override_settings(JOBS={"testjob": {"tasks": ["a"]}})
+class ShutdownTestCase(TestCase):
+    def test_shutdown_sets_state_to_stopping(self):
+        job = Job.objects.create(name="testjob")
+        worker = Worker("default", 1)
+        worker.current_job = job
+
+        worker.shutdown(None, None)
+
+        job.refresh_from_db()
+        self.assertEqual(job.state, Job.STATES.STOPPING)
+
+
 @override_settings(JOBS={"testjob": {"tasks": ["a"]}})
 class JobTestCase(TestCase):
     def test_create_job(self):