Allow setting job_id on cron jobs (#293)

* Allow setting job_id on cron jobs

* Add tests for cron job singleton

Author: JonasKs

arq/cron.py

@@ -102,6 +102,7 @@ class CronJob:
     microsecond: int
     run_at_startup: bool
     unique: bool
+    job_id: Optional[str]
     timeout_s: Optional[float]
     keep_result_s: Optional[float]
     keep_result_forever: Optional[bool]
@@ -137,6 +138,7 @@ def cron(
     microsecond: int = 123_456,
     run_at_startup: bool = False,
     unique: bool = True,
+    job_id: Optional[str] = None,
     timeout: Optional[SecondsTimedelta] = None,
     keep_result: Optional[float] = 0,
     keep_result_forever: Optional[bool] = False,
@@ -159,7 +161,8 @@ def cron(
     :param microsecond: microsecond(s) to run the job on,
         defaults to 123456 as the world is busier at the top of a second, 0 - 1e6
     :param run_at_startup: whether to run as worker starts
-    :param unique: whether the job should be only be executed once at each time
+    :param unique: whether the job should only be executed once at each time (useful if you have multiple workers)
+    :param job_id: ID of the job, can be used to enforce job uniqueness, spanning multiple cron schedules
     :param timeout: job timeout
     :param keep_result: how long to keep the result for
     :param keep_result_forever: whether to keep results forever
@@ -188,6 +191,7 @@ def cron(
         microsecond,
         run_at_startup,
         unique,
+        job_id,
         timeout,
         keep_result,
         keep_result_forever,

arq/worker.py

@@ -687,7 +687,10 @@ async def run_cron(self, n: datetime, delay: float, num_windows: int = 2) -> Non
             # We queue up the cron if the next execution time is in the next
             # delay * num_windows (by default 0.5 * 2 = 1 second).
             if cron_job.next_run < this_hb_cutoff:
-                job_id = f'{cron_job.name}:{to_unix_ms(cron_job.next_run)}' if cron_job.unique else None
+                if cron_job.job_id:
+                    job_id: Optional[str] = cron_job.job_id
+                else:
+                    job_id = f'{cron_job.name}:{to_unix_ms(cron_job.next_run)}' if cron_job.unique else None
                 job_futures.add(
                     self.pool.enqueue_job(
                         cron_job.name, _job_id=job_id, _queue_name=self.queue_name, _defer_until=cron_job.next_run

tests/test_cron.py

@@ -161,3 +161,64 @@ async def try_sleep(ctx):
     assert worker.jobs_complete == 1
     assert worker.jobs_retried == 1
     assert worker.jobs_failed == 0
+
+
+async def barfoo(ctx):
+    """In order to test cron job singleton, we must have two different functions when bursting"""
+    return 24
+
+
+async def test_job_custom_id(worker):
+    """
+    Test that two different functions with the same job_id, will only be executed once.
+    """
+    worker: Worker = worker(
+        cron_jobs=[
+            cron(barfoo, minute=10, run_at_startup=True, job_id='singleton_job'),
+            cron(foobar, minute=20, run_at_startup=True, job_id='singleton_job'),
+        ],
+        poll_delay=0.01,
+    )
+    await worker.main()
+
+    assert worker.jobs_complete == 1
+    assert worker.jobs_failed == 0
+    assert worker.jobs_retried == 0
+
+
+async def test_job_same_function_different_id(worker):
+    """
+    Set a different ID on the same function job, which runs at startup and ensure both functions are run.
+    See next test for behaviour without setting job_id.
+    """
+    worker: Worker = worker(
+        cron_jobs=[
+            cron(foobar, minute=10, run_at_startup=True, job_id='custom_id'),
+            cron(foobar, minute=20, run_at_startup=True, job_id='custom_id2'),
+        ],
+        poll_delay=0.01,
+    )
+    await worker.main()
+
+    assert worker.jobs_complete == 2
+    assert worker.jobs_failed == 0
+    assert worker.jobs_retried == 0
+
+
+async def test_run_at_startup_no_id_only_runs_once(worker):
+    """
+    Without a custom job ID, and `run_at_startup=True` on two jobs, it will only execute once, since the job_id
+    will be equal, and should only run once.
+    """
+    worker: Worker = worker(
+        cron_jobs=[
+            cron(foobar, minute=10, run_at_startup=True),
+            cron(foobar, minute=20, run_at_startup=True),
+        ],
+        poll_delay=0.01,
+    )
+    await worker.main()
+
+    assert worker.jobs_complete == 1
+    assert worker.jobs_failed == 0
+    assert worker.jobs_retried == 0