Refactoring/optimizing some bits in schedule views (#3039)

Related to https://github.com/grafana/oncall-private/issues/2163

engine/apps/api/tests/test_schedules.py

@@ -2081,7 +2081,7 @@ def test_get_schedule_on_call_now(
     client = APIClient()
     url = reverse("api-internal:schedule-list")
     with patch(
-        "apps.schedules.models.on_call_schedule.OnCallScheduleQuerySet.get_oncall_users",
+        "apps.api.views.schedule.get_oncall_users_for_multiple_schedules",
         return_value={schedule.pk: [user]},
     ):
         response = client.get(url, format="json", **make_user_auth_headers(user, token))

engine/apps/api/views/schedule.py

@@ -34,6 +34,7 @@ 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.ical_utils import get_oncall_users_for_multiple_schedules
 from apps.schedules.models import OnCallSchedule
 from apps.slack.models import SlackChannel
 from apps.slack.tasks import update_slack_user_group_for_schedules
@@ -136,10 +137,8 @@ class ScheduleView(
         The result of this method is cached and is reused for the whole lifetime of a request,
         since self.get_serializer_context() is called multiple times for every instance in the queryset.
         """
-        current_page_schedules = self.paginate_queryset(self.filter_queryset(self.get_queryset()))
-        pks = [schedule.pk for schedule in current_page_schedules]
-        queryset = OnCallSchedule.objects.filter(pk__in=pks)
-        return queryset.get_oncall_users()
+        current_page_schedules = self.paginate_queryset(self.filter_queryset(self.get_queryset(annotate=False)))
+        return get_oncall_users_for_multiple_schedules(current_page_schedules)
 
     def get_serializer_context(self):
         context = super().get_serializer_context()
@@ -167,7 +166,7 @@ class ScheduleView(
         )
         return queryset
 
-    def get_queryset(self, ignore_filtering_by_available_teams=False):
+    def get_queryset(self, ignore_filtering_by_available_teams=False, annotate=True):
         is_short_request = self.request.query_params.get("short", "false") == "true"
         filter_by_type = self.request.query_params.getlist("type")
         mine = BooleanField(allow_null=True).to_internal_value(data=self.request.query_params.get("mine"))
@@ -181,7 +180,7 @@ class ScheduleView(
         )
         if not ignore_filtering_by_available_teams:
             queryset = queryset.filter(*self.available_teams_lookup_args).distinct()
-        if not is_short_request:
+        if not is_short_request or annotate:
             queryset = self._annotate_queryset(queryset)
             queryset = self.serializer_class.setup_eager_loading(queryset)
         if filter_by_type:
@@ -231,15 +230,16 @@ class ScheduleView(
         if instance.user_group is not None:
             update_slack_user_group_for_schedules.apply_async((instance.user_group.pk,))
 
-    def get_object(self) -> OnCallSchedule:
+    def get_object(self, annotate=True) -> OnCallSchedule:
         # get the object from the whole organization if there is a flag `get_from_organization=true`
         # otherwise get the object from the current team
         get_from_organization: bool = self.request.query_params.get("from_organization", "false") == "true"
         if get_from_organization:
-            return self.get_object_from_organization()
-        return super().get_object()
+            return self.get_object_from_organization(annotate=annotate)
+        queryset_kwargs = {"annotate": annotate}
+        return super().get_object(queryset_kwargs)
 
-    def get_object_from_organization(self, ignore_filtering_by_available_teams=False):
+    def get_object_from_organization(self, ignore_filtering_by_available_teams=False, annotate=True):
         # use this method to get the object from the whole organization instead of the current team
         pk = self.kwargs["pk"]
         organization = self.request.auth.organization
@@ -248,7 +248,10 @@ class ScheduleView(
         )
         if not ignore_filtering_by_available_teams:
             queryset = queryset.filter(*self.available_teams_lookup_args).distinct()
-        queryset = self._annotate_queryset(queryset)
+
+        if annotate:
+            queryset = self._annotate_queryset(queryset)
+            queryset = self.serializer_class.setup_eager_loading(queryset)
 
         try:
             obj = queryset.get()
@@ -283,7 +286,7 @@ class ScheduleView(
         with_empty = self.request.query_params.get("with_empty", False) == "true"
         with_gap = self.request.query_params.get("with_gap", False) == "true"
 
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=False)
 
         pytz_tz = pytz.timezone(user_tz)
         datetime_start = datetime.datetime.combine(starting_date, datetime.time.min, tzinfo=pytz_tz)
@@ -319,7 +322,7 @@ class ScheduleView(
             raise BadRequest(detail="Invalid type value")
         resolve_schedule = filter_by is None or filter_by == EVENTS_FILTER_BY_FINAL
 
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=False)
 
         pytz_tz = pytz.timezone(user_tz)
         datetime_start = datetime.datetime.combine(starting_date, datetime.time.min, tzinfo=pytz_tz)
@@ -349,7 +352,7 @@ class ScheduleView(
     @action(detail=True, methods=["get"])
     def filter_shift_swaps(self, request: Request, pk: str) -> Response:
         user_tz, starting_date, days = get_date_range_from_request(self.request)
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=False)
 
         pytz_tz = pytz.timezone(user_tz)
         datetime_start = datetime.datetime.combine(starting_date, datetime.time.min, tzinfo=pytz_tz)
@@ -367,7 +370,7 @@ class ScheduleView(
         """Return next shift for users in schedule."""
         now = timezone.now()
         datetime_end = now + datetime.timedelta(days=30)
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=False)
 
         events = schedule.final_events(now, datetime_end)
 
@@ -382,7 +385,7 @@ class ScheduleView(
 
     @action(detail=True, methods=["get"])
     def related_users(self, request, pk):
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=False)
         serializer = ScheduleUserSerializer(schedule.related_users(), many=True)
         result = {"users": serializer.data}
         return Response(result, status=status.HTTP_200_OK)
@@ -390,7 +393,7 @@ class ScheduleView(
     @action(detail=True, methods=["get"])
     def related_escalation_chains(self, request, pk):
         """Return escalation chains associated to schedule."""
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=True)
         escalation_chains = EscalationChain.objects.filter(escalation_policies__notify_schedule=schedule).distinct()
 
         result = [{"name": e.name, "pk": e.public_primary_key} for e in escalation_chains]
@@ -398,7 +401,7 @@ class ScheduleView(
 
     @action(detail=True, methods=["get"])
     def quality(self, request, pk):
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=False)
 
         _, date = self.get_request_timezone()
         datetime_start = datetime.datetime.combine(date, datetime.time.min, tzinfo=pytz.UTC)
@@ -440,7 +443,7 @@ class ScheduleView(
 
     @action(detail=True, methods=["post"])
     def reload_ical(self, request, pk):
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=False)
         schedule.drop_cached_ical()
         schedule.check_empty_shifts_for_next_week()
         schedule.check_gaps_for_next_week()
@@ -452,7 +455,7 @@ class ScheduleView(
 
     @action(detail=True, methods=["get", "post", "delete"])
     def export_token(self, request, pk):
-        schedule = self.get_object()
+        schedule = self.get_object(annotate=False)
 
         if self.request.method == "GET":
             try:

engine/apps/schedules/ical_utils.py

@@ -366,18 +366,18 @@ def list_users_to_notify_from_ical_for_period(
 
 
 def get_oncall_users_for_multiple_schedules(
-    schedules: "OnCallScheduleQuerySet", events_datetime=None
+    schedules: typing.List["OnCallSchedule"], events_datetime=None
 ) -> typing.Dict["OnCallSchedule", UserQuerySet]:
     if events_datetime is None:
         events_datetime = datetime.datetime.now(timezone.utc)
 
     # Exit early if there are no schedules
-    if not schedules.exists():
+    if not schedules:
         return {}
 
     # Get on-call users
     oncall_users = {}
-    for schedule in schedules.all():
+    for schedule in schedules:
         # pass user list to list_users_to_notify_from_ical
         schedule_oncall_users = list_users_to_notify_from_ical(schedule, events_datetime=events_datetime)
         oncall_users.update({schedule.pk: schedule_oncall_users})

engine/apps/schedules/models/on_call_schedule.py

@@ -151,7 +151,7 @@ def generate_public_primary_key_for_oncall_schedule_channel():
 
 class OnCallScheduleQuerySet(PolymorphicQuerySet):
     def get_oncall_users(self, events_datetime=None):
-        return get_oncall_users_for_multiple_schedules(self, events_datetime)
+        return get_oncall_users_for_multiple_schedules(self.all(), events_datetime)
 
     def related_to_user(self, user):
         username_regex = RE_ICAL_SEARCH_USERNAME.format(user.username)

engine/common/api_helpers/mixins.py

@@ -150,9 +150,11 @@ _MT = typing.TypeVar("_MT", bound=models.Model)
 
 
 class PublicPrimaryKeyMixin(typing.Generic[_MT]):
-    def get_object(self) -> _MT:
+    def get_object(self, queryset_kwargs=None) -> _MT:
         pk = self.kwargs["pk"]
-        queryset = self.filter_queryset(self.get_queryset())
+        if queryset_kwargs is None:
+            queryset_kwargs = {}
+        queryset = self.filter_queryset(self.get_queryset(**queryset_kwargs))
 
         try:
             obj = queryset.get(public_primary_key=pk)