Add timezone support for cron jobs (#354)

* Test that next_cron preserves tzinfo

* Add Worker.timezone attribute

* Pass timezone aware now to run_cron

* Make tests work in 3.7

Co-authored-by: Samuel Colvin <[email protected]>

Author: iamlikeme

arq/worker.py

@@ -3,7 +3,7 @@
 import logging
 import signal
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from functools import partial
 from signal import Signals
 from time import time
@@ -170,6 +170,8 @@ class Worker:
     :param job_deserializer: a function that deserializes bytes into Python objects, defaults to pickle.loads
     :param expires_extra_ms: the default length of time from when a job is expected to start
      after which the job expires, defaults to 1 day in ms.
+    :param timezone: timezone used for evaluation of cron schedules,
+        defaults to system timezone
     :param log_results: when set to true (default) results for successful jobs
       will be logged
     """
@@ -204,6 +206,7 @@ def __init__(
         job_serializer: Optional[Serializer] = None,
         job_deserializer: Optional[Deserializer] = None,
         expires_extra_ms: int = expires_extra_ms,
+        timezone: Optional[timezone] = None,
         log_results: bool = True,
     ):
         self.functions: Dict[str, Union[Function, CronJob]] = {f.name: f for f in map(func, functions)}
@@ -271,6 +274,9 @@ def __init__(
         self.expires_extra_ms = expires_extra_ms
         self.log_results = log_results
 
+        # default to system timezone
+        self.timezone = datetime.now().astimezone().tzinfo if timezone is None else timezone
+
     def run(self) -> None:
         """
         Sync function to run the worker, finally closes worker connections.
@@ -677,7 +683,7 @@ async def finish_failed_job(self, job_id: str, result_data: Optional[bytes]) ->
             await tr.execute()
 
     async def heart_beat(self) -> None:
-        now = datetime.now()
+        now = datetime.now(tz=self.timezone)
         await self.record_health()
 
         cron_window_size = max(self.poll_delay_s, 0.5)  # Clamp the cron delay to 0.5

tests/test_cron.py

@@ -5,6 +5,7 @@
 from random import random
 
 import pytest
+from pytest_mock import MockerFixture
 
 import arq
 from arq import Worker
@@ -51,6 +52,12 @@ def test_next_cron(previous, expected, kwargs):
     print(f'{diff.total_seconds() * 1000:0.3f}ms')
 
 
+def test_next_cron_preserves_tzinfo():
+    previous = datetime.fromisoformat('2016-06-01T12:10:10+02:00')
+    assert previous.tzinfo is not None
+    assert next_cron(previous, second=20).tzinfo is previous.tzinfo
+
+
 def test_next_cron_invalid():
     with pytest.raises(ValueError):
         next_cron(datetime(2001, 1, 1, 0, 0, 0), weekday='monday')
@@ -110,6 +117,16 @@ async def test_job_successful(worker, caplog, arq_redis, poll_delay):
     assert await arq_redis.pttl(keys[0]) > 0.0
 
 
+async def test_calculate_next_is_called_with_aware_datetime(worker, mocker: MockerFixture):
+    worker: Worker = worker(cron_jobs=[cron(foobar, hour=1)])
+    spy_run_cron = mocker.spy(worker, worker.run_cron.__name__)
+
+    await worker.main()
+
+    assert spy_run_cron.call_args[0][0].tzinfo is not None
+    assert worker.cron_jobs[0].next_run.tzinfo is not None
+
+
 async def test_job_successful_on_specific_queue(worker, caplog):
     caplog.set_level(logging.INFO)
     worker: Worker = worker(

tests/test_worker.py

@@ -940,3 +940,9 @@ async def test(ctx):
     assert worker.jobs_failed == 0
     assert worker.jobs_retried == 0
     assert result['called'] == 2
+
+
+async def test_worker_timezone_defaults_to_system_timezone(worker):
+    worker = worker(functions=[func(foobar)])
+    assert worker.timezone is not None
+    assert worker.timezone == datetime.now().astimezone().tzinfo