diff --git a/engine/apps/api/views/schedule.py b/engine/apps/api/views/schedule.py index 9a4c95e6..cde5b89a 100644 --- a/engine/apps/api/views/schedule.py +++ b/engine/apps/api/views/schedule.py @@ -4,7 +4,7 @@ import operator import pytz from django.core.exceptions import ObjectDoesNotExist -from django.db.models import Count, OuterRef, Subquery +from django.db.models import Count, OuterRef, Prefetch, Q, Subquery from django.db.utils import IntegrityError from django.urls import reverse from django.utils import dateparse, timezone @@ -34,8 +34,9 @@ from apps.auth_token.auth import PluginAuthentication from apps.auth_token.constants import SCHEDULE_EXPORT_TOKEN_NAME from apps.auth_token.models import ScheduleExportAuthToken from apps.mobile_app.auth import MobileAppAuthTokenAuthentication +from apps.schedules.constants import PREFETCHED_SHIFT_SWAPS from apps.schedules.ical_utils import get_oncall_users_for_multiple_schedules -from apps.schedules.models import OnCallSchedule +from apps.schedules.models import OnCallSchedule, ShiftSwapRequest from apps.slack.models import SlackChannel from apps.slack.tasks import update_slack_user_group_for_schedules from common.api_helpers.exceptions import BadRequest, Conflict @@ -138,13 +139,29 @@ class ScheduleView( since self.get_serializer_context() is called multiple times for every instance in the queryset. """ current_schedules = self.get_queryset(annotate=False).none() + events_datetime = datetime.datetime.now(datetime.timezone.utc) if self.action == "list": - # listing page, only get oncall users for current page schedules - current_schedules = self.paginate_queryset(self.filter_queryset(self.get_queryset(annotate=False))) + # listing page, only get oncall users for current page schedules, prefetch shift swap requests + current_schedules = self.filter_queryset(self.get_queryset(annotate=False)).prefetch_related( + self.prefetch_shift_swaps( + queryset=ShiftSwapRequest.objects.filter( + swap_start__lte=events_datetime, swap_end__gte=events_datetime + ) + ) + ) + current_schedules = self.paginate_queryset(current_schedules) elif self.kwargs.get("pk"): # if this is a particular schedule detail, only consider it as current current_schedules = [self.get_object(annotate=False)] - return get_oncall_users_for_multiple_schedules(current_schedules) + return get_oncall_users_for_multiple_schedules(current_schedules, events_datetime) + + @staticmethod + def prefetch_shift_swaps(queryset): + return Prefetch( + "shift_swap_requests", + queryset=queryset.select_related("benefactor", "beneficiary").order_by("created_at"), + to_attr=PREFETCHED_SHIFT_SWAPS, + ) def get_serializer_context(self): context = super().get_serializer_context() @@ -186,8 +203,9 @@ class ScheduleView( ) if not ignore_filtering_by_available_teams: queryset = queryset.filter(*self.available_teams_lookup_args).distinct() - if not is_short_request or annotate: - queryset = self._annotate_queryset(queryset) + if not is_short_request: + if annotate: + queryset = self._annotate_queryset(queryset) queryset = self.serializer_class.setup_eager_loading(queryset) if filter_by_type: valid_types = [i for i in filter_by_type if i in SCHEDULE_TYPE_TO_CLASS] @@ -425,8 +443,19 @@ class ScheduleView( user_tz, starting_date, days = get_date_range_from_request(self.request) pytz_tz = pytz.timezone(user_tz) datetime_start = datetime.datetime.combine(starting_date, datetime.time.min, tzinfo=pytz_tz) - - schedules = OnCallSchedule.objects.related_to_user(self.request.user) + datetime_end = datetime_start + datetime.timedelta(days=days) + schedules = ( + OnCallSchedule.objects.related_to_user(self.request.user) + .select_related("organization") + .prefetch_related( + self.prefetch_shift_swaps( + queryset=ShiftSwapRequest.objects.filter( + Q(swap_start__lt=datetime_start, swap_end__gte=datetime_start) + | Q(swap_start__gte=datetime_start, swap_start__lte=datetime_end) + ) + ) + ) + ) schedules_events = [] is_oncall = False for schedule in schedules: diff --git a/engine/apps/schedules/constants.py b/engine/apps/schedules/constants.py index 880cf50d..3f8f556e 100644 --- a/engine/apps/schedules/constants.py +++ b/engine/apps/schedules/constants.py @@ -29,3 +29,5 @@ EXPORT_WINDOW_DAYS_BEFORE = 15 SCHEDULE_ONCALL_CACHE_KEY_PREFIX = "schedule_oncall_users_" SCHEDULE_ONCALL_CACHE_TTL = 15 * 60 # 15 minutes in seconds + +PREFETCHED_SHIFT_SWAPS = "prefetched_shift_swaps" diff --git a/engine/apps/schedules/ical_utils.py b/engine/apps/schedules/ical_utils.py index 488faf36..8eeacddf 100644 --- a/engine/apps/schedules/ical_utils.py +++ b/engine/apps/schedules/ical_utils.py @@ -101,7 +101,7 @@ def users_in_ical( @timed_lru_cache(timeout=100) -def memoized_users_in_ical(usernames_from_ical: typing.List[str], organization: "Organization") -> typing.List["User"]: +def memoized_users_in_ical(usernames_from_ical: typing.Tuple[str], organization: "Organization") -> typing.List["User"]: # using in-memory cache instead of redis to avoid pickling python objects return users_in_ical(usernames_from_ical, organization) @@ -358,15 +358,13 @@ def list_users_to_notify_from_ical_for_period( schedule: "OnCallSchedule", start_datetime: datetime.datetime, end_datetime: datetime.datetime, -) -> typing.List["User"]: - users_found_in_ical: typing.Sequence["User"] = [] +) -> typing.Sequence["User"]: events = schedule.final_events(start_datetime, end_datetime) - usernames = [] + usernames: typing.List[str] = [] for event in events: usernames += [u["email"] for u in event.get("users", [])] - users_found_in_ical = users_in_ical(usernames, schedule.organization) - return users_found_in_ical + return memoized_users_in_ical(tuple(usernames), schedule.organization) def get_oncall_users_for_multiple_schedules( diff --git a/engine/apps/schedules/models/on_call_schedule.py b/engine/apps/schedules/models/on_call_schedule.py index afa99664..5dcd83e4 100644 --- a/engine/apps/schedules/models/on_call_schedule.py +++ b/engine/apps/schedules/models/on_call_schedule.py @@ -32,6 +32,7 @@ from apps.schedules.constants import ( ICAL_STATUS_CANCELLED, ICAL_SUMMARY, ICAL_UID, + PREFETCHED_SHIFT_SWAPS, ) from apps.schedules.ical_utils import ( EmptyShifts, @@ -465,7 +466,7 @@ class OnCallSchedule(PolymorphicModel): return events def filter_swap_requests( - self, datetime_start: datetime.datetime, datetime_end: datetime.time + self, datetime_start: datetime.datetime, datetime_end: datetime.datetime ) -> "RelatedManager['ShiftSwapRequest']": swap_requests = self.shift_swap_requests.filter( # starting before but ongoing swap_start__lt=datetime_start, swap_end__gte=datetime_start @@ -716,7 +717,12 @@ class OnCallSchedule(PolymorphicModel): ) -> ScheduleEvents: """Apply swap requests details to schedule events.""" # get swaps requests affecting this schedule / time range - swaps = self.filter_swap_requests(datetime_start, datetime_end) + prefetched_swaps = getattr(self, PREFETCHED_SHIFT_SWAPS, None) + swaps = ( + prefetched_swaps + if prefetched_swaps is not None + else self.filter_swap_requests(datetime_start, datetime_end) + ) def _insert_event(index: int, event: ScheduleEvent) -> int: # add event, if any, to events list in the specified index