add on_job_start and on_job_end hooks (#274)

Author: tobymao

arq/worker.py

@@ -138,6 +138,8 @@ class Worker:
     :param burst: whether to stop the worker once all jobs have been run
     :param on_startup: coroutine function to run at startup
     :param on_shutdown: coroutine function to run at shutdown
+    :param on_job_start: coroutine function to run on job start
+    :param on_job_end: coroutine function to run on job end
     :param handle_signals: default true, register signal handlers,
       set to false when running inside other async framework
     :param max_jobs: maximum number of jobs to run at a time
@@ -169,6 +171,8 @@ def __init__(
         burst: bool = False,
         on_startup: Optional['StartupShutdown'] = None,
         on_shutdown: Optional['StartupShutdown'] = None,
+        on_job_start: Optional['StartupShutdown'] = None,
+        on_job_end: Optional['StartupShutdown'] = None,
         handle_signals: bool = True,
         max_jobs: int = 10,
         job_timeout: 'SecondsTimedelta' = 300,
@@ -202,6 +206,8 @@ def __init__(
         self.burst = burst
         self.on_startup = on_startup
         self.on_shutdown = on_shutdown
+        self.on_job_start = on_job_start
+        self.on_job_end = on_job_end
         self.sem = asyncio.BoundedSemaphore(max_jobs)
         self.job_timeout_s = to_seconds(job_timeout)
         self.keep_result_s = to_seconds(keep_result)
@@ -509,6 +515,10 @@ async def job_failed(exc: BaseException) -> None:
             'score': score,
         }
         ctx = {**self.ctx, **job_ctx}
+
+        if self.on_job_start:
+            await self.on_job_start(ctx)
+
         start_ms = timestamp_ms()
         success = False
         try:
@@ -585,6 +595,9 @@ async def job_failed(exc: BaseException) -> None:
                 serializer=self.job_serializer,
             )
 
+        if self.on_job_end:
+            await self.on_job_end(ctx)
+
         await asyncio.shield(
             self.finish_job(
                 job_id, finish, result_data, result_timeout_s, keep_result_forever, incr_score, keep_in_progress,

tests/test_worker.py

@@ -835,3 +835,32 @@ async def longfunc(ctx):
     assert worker.jobs_retried == 0
     log = re.sub(r'\d+.\d\ds', 'X.XXs', '\n'.join(r.message for r in caplog.records))
     assert 'X.XXs ! testing:longfunc failed, TimeoutError:' in log
+
+
+async def test_on_job(arq_redis: ArqRedis, worker):
+    result = {'called': 0}
+
+    async def on_start(ctx):
+        assert ctx['job_id'] == 'testing'
+        result['called'] += 1
+
+    async def on_end(ctx):
+        assert ctx['job_id'] == 'testing'
+        result['called'] += 1
+
+    async def test(ctx):
+        return
+
+    await arq_redis.enqueue_job('func', _job_id='testing')
+    worker: Worker = worker(
+        functions=[func(test, name='func')], on_job_start=on_start, on_job_end=on_end, job_timeout=0.2, poll_delay=0.1,
+    )
+    assert worker.jobs_complete == 0
+    assert worker.jobs_failed == 0
+    assert worker.jobs_retried == 0
+    assert result['called'] == 0
+    await worker.main()
+    assert worker.jobs_complete == 1
+    assert worker.jobs_failed == 0
+    assert worker.jobs_retried == 0
+    assert result['called'] == 2