diff --git a/django_celery_beat/schedulers.py b/django_celery_beat/schedulers.py index 99d98f8c..9db2f062 100644 --- a/django_celery_beat/schedulers.py +++ b/django_celery_beat/schedulers.py @@ -1,4 +1,6 @@ """Beat Scheduler Implementation.""" +from __future__ import annotations + import datetime import logging import math @@ -15,7 +17,7 @@ from celery.utils.time import maybe_make_aware from django.conf import settings from django.core.exceptions import ObjectDoesNotExist -from django.db import close_old_connections, transaction +from django.db import close_old_connections, connection, transaction from django.db.models import Case, F, IntegerField, Q, When from django.db.models.functions import Cast from django.db.utils import DatabaseError, InterfaceError @@ -330,14 +332,14 @@ def _get_crontab_exclude_query(self): # Handle each timezone specifically *[ When( - timezone=timezone_name, + timezone=timezone, then=( F('hour_int') - + self._get_timezone_offset(timezone_name) + + self._get_timezone_offset(timezone) + 24 ) % 24 ) - for timezone_name in self._get_unique_timezone_names() + for timezone in self._get_unique_timezones() ], # Default case - use hour as is default=F('hour_int') @@ -365,13 +367,27 @@ def _get_valid_hour_formats(self): f"{hour:02d}" for hour in range(10) ] - def _get_unique_timezone_names(self): - """Get a list of all unique timezone names used in CrontabSchedule""" - return CrontabSchedule.objects.values_list( - 'timezone', flat=True - ).distinct() + def _get_unique_timezones(self) -> set[ZoneInfo]: + """Get a set of all unique timezones used in CrontabSchedule""" + # sqlite does not support distinct on a column + if connection.vendor == 'sqlite': + return set( + CrontabSchedule.objects.values_list( + 'timezone', flat=True + ) + ) + + return set( + CrontabSchedule.objects.order_by( + 'timezone' + ).distinct( + 'timezone' + ).values_list( + 'timezone', flat=True + ) + ) - def _get_timezone_offset(self, timezone_name): + def _get_timezone_offset(self, timezone: ZoneInfo) -> int: """ Args: timezone_name: The name of the timezone or a ZoneInfo object @@ -387,17 +403,12 @@ def _get_timezone_offset(self, timezone_name): else: server_tz = ZoneInfo(str(server_time.tzinfo)) - if isinstance(timezone_name, ZoneInfo): - timezone_name = timezone_name.key - - target_tz = ZoneInfo(timezone_name) - # Use a fixed point in time for the calculation to avoid DST issues fixed_dt = datetime.datetime(2023, 1, 1, 12, 0, 0) # Calculate the offset dt1 = fixed_dt.replace(tzinfo=server_tz) - dt2 = fixed_dt.replace(tzinfo=target_tz) + dt2 = fixed_dt.replace(tzinfo=timezone) # Calculate hour difference offset_seconds = ( diff --git a/t/unit/test_schedulers.py b/t/unit/test_schedulers.py index f7dfac8b..45c4a609 100644 --- a/t/unit/test_schedulers.py +++ b/t/unit/test_schedulers.py @@ -333,7 +333,7 @@ def test_entry_is_due__no_use_tz(self): assert self.app.timezone.key == 'Europe/Berlin' # simulate last_run_at from DB - not TZ aware but localtime - right_now = datetime.utcnow() + right_now = datetime.now(dt_timezone.utc) m = self.create_model_crontab( crontab(minute='*/10'), @@ -364,7 +364,7 @@ def test_entry_and_model_last_run_at_with_utc_no_use_tz(self, monkeypatch): time.tzset() assert self.app.timezone.key == 'Europe/Berlin' # simulate last_run_at from DB - not TZ aware but localtime - right_now = datetime.utcnow() + right_now = datetime.now(dt_timezone.utc) # make sure to use fixed date time monkeypatch.setattr(self.Entry, '_default_now', lambda o: right_now) m = self.create_model_crontab( @@ -398,7 +398,7 @@ def test_entry_and_model_last_run_at_when_model_changed(self, monkeypatch): time.tzset() assert self.app.timezone.key == 'Europe/Berlin' # simulate last_run_at from DB - not TZ aware but localtime - right_now = datetime.utcnow() + right_now = datetime.now(dt_timezone.utc) # make sure to use fixed date time monkeypatch.setattr(self.Entry, '_default_now', lambda o: right_now) m = self.create_model_crontab( @@ -451,7 +451,7 @@ def test_entry_is_due__celery_timezone_doesnt_match_time_zone(self): # simulate last_run_at all none, doing the same thing that # _default_now() would do - right_now = datetime.utcnow() + right_now = datetime.now(dt_timezone.utc) m = self.create_model_crontab( crontab(minute='*/10'), @@ -1224,6 +1224,21 @@ def test_crontab_special_hour_four(self): # The hour=4 task should never be excluded assert task_hour_four.id not in excluded_tasks + @pytest.mark.django_db + def test_get_unique_timezones(self): + """ + Test that get unique timezones returns a list of unique timezones + """ + # Create 2 crontabs with same timezone, and 1 with different timezone + CrontabSchedule.objects.create(hour='4', timezone='UTC') + CrontabSchedule.objects.create(hour='4', timezone='UTC') + CrontabSchedule.objects.create(hour='4', timezone='America/New_York') + + timezones = self.s._get_unique_timezones() + + assert len(timezones) == 2 + assert set(timezones) == {ZoneInfo('UTC'), ZoneInfo('America/New_York')} + @pytest.mark.django_db @patch('django_celery_beat.schedulers.aware_now') @patch('django.utils.timezone.get_current_timezone') @@ -1663,31 +1678,6 @@ def setup_method(self): def teardown_method(self): patch.stopall() - @patch("django_celery_beat.schedulers.aware_now") - def test_server_timezone_handling_with_zoneinfo(self, mock_aware_now): - """Test handling when server timezone - is already a ZoneInfo instance.""" - - # Create a mock scheduler with only the methods we need to test - class MockScheduler: - _get_timezone_offset = ( - schedulers.DatabaseScheduler._get_timezone_offset - ) - - s = MockScheduler() - - tokyo_tz = ZoneInfo("Asia/Tokyo") - mock_now = datetime(2023, 1, 1, 12, 0, 0, tzinfo=tokyo_tz) - mock_aware_now.return_value = mock_now - - # Test with a different timezone - new_york_tz = "America/New_York" - offset = s._get_timezone_offset(new_york_tz) # Pass self explicitly - - # Tokyo is UTC+9, New York is UTC-5, so difference should be 14 hours - assert offset == 14 - assert mock_aware_now.called - @patch("django_celery_beat.schedulers.aware_now") def test_timezone_offset_with_zoneinfo_object_param(self, mock_aware_now): """Test handling when timezone_name parameter is a ZoneInfo object.""" @@ -1709,3 +1699,4 @@ class MockScheduler: # Tokyo is UTC+9, New York is UTC-5, so difference should be 14 hours assert offset == 14 + assert mock_aware_now.called