Merge pull request #33 from dabapps/queue-depth-utils

Queue Depth Utils

Author: jordaneb

README.md

@@ -182,13 +182,34 @@ Jobs have a `state` field which can have one of the following values:
 
 ### API
 
+#### Model methods
+
+##### Job.get_queue_depths
+If you need to programatically get the depth of any queue you can run the following:
+```python
+from django_dbq.models import Job
+
+...
+
+Job.objects.create(name='do_work', workspace={})
+Job.objects.create(name='do_other_work', queue_name='other_queue', workspace={})
+
+queue_depths = Job.get_queue_depths()
+print(queue_depths)  # {"default": 1, "other_queue": 1}
+```
+
+**Important:** When checking queue depths, do not assume that the key for your queue will always be available. Queue depths of zero won't be included
+in the dict returned by this method.
+
 #### Management commands
 
+##### manage.py delete_old_jobs
 There is a management command, `manage.py delete_old_jobs`, which deletes any
 jobs from the database which are in state `COMPLETE` or `FAILED` and were
 created more than 24 hours ago. This could be run, for example, as a cron task,
 to ensure the jobs table remains at a reasonable size.
 
+##### manage.py create_job
 For debugging/development purposes, a simple management command is supplied to create jobs:
 
     manage.py create_job <job_name> --queue_name 'my_queue_name' --workspace '{"key": "value"}'
@@ -197,13 +218,20 @@ The `workspace` flag is optional. If supplied, it must be a valid JSON string.
 
 `queue_name` is optional and defaults to `default`
 
+##### manage.py worker
 To start a worker:
 
     manage.py worker [queue_name] [--rate_limit]
 
 - `queue_name` is optional, and will default to `default`
 - The `--rate_limit` flag is optional, and will default to `1`. It is the minimum number of seconds that must have elapsed before a subsequent job can be run.
 
+##### manage.py queue_depth
+If you'd like to check your queue depth from the command line, you can run `manage.py queue_depth [queue_name [queue_name ...]]` and any
+jobs in the "NEW" or "READY" states will be returned.
+
+**Important:** If you misspell or provide a queue name which does not have any jobs, a depth of 0 will always be returned.
+
 ## Testing
 
 It may be necessary to supply a DATABASE_PORT environment variable.

django_dbq/management/commands/queue_depth.py

@@ -0,0 +1,27 @@
+from django.core.management.base import BaseCommand
+from django_dbq.models import Job
+
+
+class Command(BaseCommand):
+
+    help = "Print the current depth of the given queue"
+
+    def add_arguments(self, parser):
+        parser.add_argument("queue_name", nargs="*", default=["default"], type=str)
+
+    def handle(self, *args, **options):
+        queue_names = options["queue_name"]
+        queue_depths = Job.get_queue_depths()
+
+        queue_depths_string = " ".join(
+            [
+                "{queue_name}={queue_depth}".format(
+                    queue_name=queue_name, queue_depth=queue_depths.get(queue_name, 0),
+                )
+                for queue_name in queue_names
+            ]
+        )
+
+        self.stdout.write(
+            "event=queue_depths {queue_depths}".format(queue_depths=queue_depths_string)
+        )

django_dbq/models.py

@@ -6,7 +6,7 @@
     get_creation_hook_name,
 )
 from jsonfield import JSONField
-from django.db.models import UUIDField
+from django.db.models import UUIDField, Count
 import datetime
 import logging
 import uuid
@@ -137,3 +137,17 @@ def run_creation_hook(self):
             logger.info("Running creation hook %s for new job", creation_hook_name)
             creation_hook_function = import_string(creation_hook_name)
             creation_hook_function(self)
+
+    @staticmethod
+    def get_queue_depths():
+        annotation_dicts = (
+            Job.objects.filter(state__in=(Job.STATES.READY, Job.STATES.NEW))
+            .values("queue_name")
+            .order_by("queue_name")
+            .annotate(Count("queue_name"))
+        )
+
+        return {
+            annotation_dict["queue_name"]: annotation_dict["queue_name__count"]
+            for annotation_dict in annotation_dicts
+        }

django_dbq/tests.py

@@ -81,6 +81,70 @@ def test_worker_with_queue_name(self):
         self.assertTrue("test_queue" in output)
 
 
+@override_settings(JOBS={"testjob": {"tasks": ["a"]}})
+class JobModelMethodTestCase(TestCase):
+    def test_get_queue_depths(self):
+        Job.objects.create(name="testjob", queue_name="default")
+        Job.objects.create(name="testjob", queue_name="testworker")
+        Job.objects.create(name="testjob", queue_name="testworker")
+        Job.objects.create(
+            name="testjob", queue_name="testworker", state=Job.STATES.FAILED
+        )
+        Job.objects.create(
+            name="testjob", queue_name="testworker", state=Job.STATES.COMPLETE
+        )
+
+        queue_depths = Job.get_queue_depths()
+        self.assertDictEqual(queue_depths, {"default": 1, "testworker": 2})
+
+
+@override_settings(JOBS={"testjob": {"tasks": ["a"]}})
+class QueueDepthTestCase(TestCase):
+    def test_queue_depth(self):
+
+        Job.objects.create(name="testjob", state=Job.STATES.FAILED)
+        Job.objects.create(name="testjob", state=Job.STATES.NEW)
+        Job.objects.create(name="testjob", state=Job.STATES.FAILED)
+        Job.objects.create(name="testjob", state=Job.STATES.COMPLETE)
+        Job.objects.create(name="testjob", state=Job.STATES.READY)
+        Job.objects.create(
+            name="testjob", queue_name="testqueue", state=Job.STATES.READY
+        )
+        Job.objects.create(
+            name="testjob", queue_name="testqueue", state=Job.STATES.READY
+        )
+
+        stdout = StringIO()
+        call_command("queue_depth", stdout=stdout)
+        output = stdout.getvalue()
+        self.assertEqual(output.strip(), "event=queue_depths default=2")
+
+    def test_queue_depth_multiple_queues(self):
+
+        Job.objects.create(name="testjob", state=Job.STATES.FAILED)
+        Job.objects.create(name="testjob", state=Job.STATES.NEW)
+        Job.objects.create(name="testjob", state=Job.STATES.FAILED)
+        Job.objects.create(name="testjob", state=Job.STATES.COMPLETE)
+        Job.objects.create(name="testjob", state=Job.STATES.READY)
+        Job.objects.create(
+            name="testjob", queue_name="testqueue", state=Job.STATES.READY
+        )
+        Job.objects.create(
+            name="testjob", queue_name="testqueue", state=Job.STATES.READY
+        )
+
+        stdout = StringIO()
+        call_command("queue_depth", queue_name=("default", "testqueue",), stdout=stdout)
+        output = stdout.getvalue()
+        self.assertEqual(output.strip(), "event=queue_depths default=2 testqueue=2")
+
+    def test_queue_depth_for_queue_with_zero_jobs(self):
+        stdout = StringIO()
+        call_command("queue_depth", queue_name=("otherqueue",), stdout=stdout)
+        output = stdout.getvalue()
+        self.assertEqual(output.strip(), "event=queue_depths otherqueue=0")
+
+
 @freezegun.freeze_time()
 @mock.patch("django_dbq.management.commands.worker.sleep")
 @mock.patch("django_dbq.management.commands.worker.process_job")