Skip to content

Commit d5abadd

Browse files
committed
fix _get_distinct_timezone
the previous query did not actually get distinct timezones
1 parent 673dbc5 commit d5abadd

File tree

2 files changed

+38
-7
lines changed

2 files changed

+38
-7
lines changed

django_celery_beat/schedulers.py

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
"""Beat Scheduler Implementation."""
2+
from __future__ import annotations
3+
24
import datetime
35
import logging
46
import math
@@ -15,7 +17,7 @@
1517
from celery.utils.time import maybe_make_aware
1618
from django.conf import settings
1719
from django.core.exceptions import ObjectDoesNotExist
18-
from django.db import close_old_connections, transaction
20+
from django.db import close_old_connections, connection, transaction
1921
from django.db.models import Case, F, IntegerField, Q, When
2022
from django.db.models.functions import Cast
2123
from django.db.utils import DatabaseError, InterfaceError
@@ -340,7 +342,7 @@ def _get_crontab_exclude_query(self):
340342
+ 24
341343
) % 24
342344
)
343-
for timezone_name in self._get_unique_timezone_names()
345+
for timezone_name in self._get_unique_timezones()
344346
],
345347
# Default case - use hour as is
346348
default=F('hour_int')
@@ -359,11 +361,25 @@ def _get_crontab_exclude_query(self):
359361

360362
return exclude_query
361363

362-
def _get_unique_timezone_names(self):
363-
"""Get a list of all unique timezone names used in CrontabSchedule"""
364-
return CrontabSchedule.objects.values_list(
365-
'timezone', flat=True
366-
).distinct()
364+
def _get_unique_timezones(self) -> set[ZoneInfo]:
365+
"""Get a set of all unique timezones used in CrontabSchedule"""
366+
# sqlite does not support distinct on a column
367+
if connection.vendor == 'sqlite':
368+
return set(
369+
CrontabSchedule.objects.values_list(
370+
'timezone', flat=True
371+
)
372+
)
373+
374+
return set(
375+
CrontabSchedule.objects.order_by(
376+
'timezone'
377+
).distinct(
378+
'timezone'
379+
).values_list(
380+
'timezone', flat=True
381+
)
382+
)
367383

368384
def _get_timezone_offset(self, timezone_name):
369385
"""

t/unit/test_schedulers.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1059,6 +1059,21 @@ def test_crontab_special_hour_four(self):
10591059
# The hour=4 task should never be excluded
10601060
assert task_hour_four.id not in excluded_tasks
10611061

1062+
@pytest.mark.django_db
1063+
def test_get_unique_timezones(self):
1064+
"""
1065+
Test that get unique timezones returns a list of unique timezones
1066+
"""
1067+
# Create 2 crontabs with same timezone, and 1 with different timezone
1068+
CrontabSchedule.objects.create(hour='4', timezone='UTC')
1069+
CrontabSchedule.objects.create(hour='4', timezone='UTC')
1070+
CrontabSchedule.objects.create(hour='4', timezone='America/New_York')
1071+
1072+
timezones = self.s._get_unique_timezones()
1073+
1074+
assert len(timezones) == 2
1075+
assert set(timezones) == {ZoneInfo('UTC'), ZoneInfo('America/New_York')}
1076+
10621077
@pytest.mark.django_db
10631078
@patch('django_celery_beat.schedulers.aware_now')
10641079
@patch('django.utils.timezone.get_current_timezone')

0 commit comments

Comments
 (0)