optimize GET /schedules internal API endpoint (#1169)

# What this PR does

Fixes slow internal`GET /schedules` endpoints. Using the fake-data
generation script in #1128, I generated 65 calendar schedules in my
local setup. This resulted in the following endpoint performance:
![Screenshot 2023-01-24 at 12 03
16](https://user-images.githubusercontent.com/9406895/214276618-1a9848ba-eb84-49ec-a099-fdd96beac93f.png)

The responses which show ~76 queries were run on the latest `dev`
branch. Responses w/ ~26 queries were run on this branch.

Additionally:
- add typing to a few methods in `apps/schedules/ical_utils.py`
- document `apps/api/permissions/__init__.py:user_is_authorized`
function

## Which issue(s) this PR fixes

grafana/oncall-private#1552

## Checklist

- [ ] Tests updated
- [ ] Documentation added
- [ ] `CHANGELOG.md` updated

Co-authored-by: Vadim Stepanov <[email protected]>

Author: joeyorlando

CHANGELOG.md

@@ -27,6 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Allow users with `viewer` role to fetch cloud connection status using the internal API ([#1181](https://github.com/grafana/oncall/pull/1181))
 - When removing the Slack ChatOps integration, make it more explicit to the user what the implications of doing so are
+- Improve performance of `GET /api/internal/v1/schedules` endpoint ([#1169](https://github.com/grafana/oncall/pull/1169))
 
 ### Fixed
 

engine/apps/api/permissions/__init__.py

@@ -73,6 +73,17 @@ def get_most_authorized_role(
 
 
 def user_is_authorized(user, required_permissions: typing.List[LegacyAccessControlCompatiblePermission]) -> bool:
+    """
+    This function checks whether `user` has all permissions in `required_permissions`. RBAC permissions are used
+    if RBAC is enabled for the organization, otherwise the fallback basic role is checked.
+
+    Parameters
+    ----------
+    user : apps.user_management.models.user.User
+        The user to check permissions for
+    required_permissions : typing.List[LegacyAccessControlCompatiblePermission]
+        A list of permissions that a user must have to be considered authorized
+    """
     if user.organization.is_rbac_permissions_enabled:
         user_permissions = [u["action"] for u in user.permissions]
         required_permissions = [p.value for p in required_permissions]

engine/apps/api/serializers/schedule_base.py

@@ -1,8 +1,6 @@
-from django.utils import timezone
 from rest_framework import serializers
 
 from apps.api.serializers.user_group import UserGroupSerializer
-from apps.schedules.ical_utils import list_users_to_notify_from_ical
 from apps.schedules.models import OnCallSchedule
 from apps.schedules.tasks import schedule_notify_about_empty_shifts_in_schedule, schedule_notify_about_gaps_in_schedule
 from common.api_helpers.custom_fields import TeamPrimaryKeyRelatedField
@@ -67,11 +65,9 @@ def get_warnings(self, obj):
         return warnings
 
     def get_on_call_now(self, obj):
-        users_on_call = list_users_to_notify_from_ical(obj, timezone.datetime.now(timezone.utc))
-        if users_on_call is not None:
-            return [user.short() for user in users_on_call]
-        else:
-            return []
+        # Serializer context is set here: apps.api.views.schedule.ScheduleView.get_serializer_context
+        users = self.context["oncall_users"].get(obj.pk, [])
+        return [user.short() for user in users]
 
     def get_number_of_escalation_chains(self, obj):
         # num_escalation_chains param added in queryset via annotate. Check ScheduleView.get_queryset

engine/apps/api/views/schedule.py

@@ -104,9 +104,19 @@ def can_update_user_groups(self):
 
         return user_group.can_be_updated
 
+    @cached_property
+    def oncall_users(self):
+        """
+        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.
+        """
+        queryset = self.get_queryset()
+        return queryset.get_oncall_users()
+
     def get_serializer_context(self):
         context = super().get_serializer_context()
         context.update({"can_update_user_groups": self.can_update_user_groups})
+        context.update({"oncall_users": self.oncall_users})
         return context
 
     def _annotate_queryset(self, queryset):

engine/apps/schedules/ical_utils.py

@@ -3,6 +3,7 @@
 import datetime
 import logging
 import re
+import typing
 from collections import namedtuple
 from typing import TYPE_CHECKING
 
@@ -35,60 +36,64 @@
 """
 if TYPE_CHECKING:
     from apps.schedules.models import OnCallSchedule
-    from apps.user_management.models import User
+    from apps.user_management.models import Organization, User
+    from apps.user_management.models.user import UserQuerySet
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
 
 
-def users_in_ical(usernames_from_ical, organization, include_viewers=False):
+def users_in_ical(
+    usernames_from_ical: typing.List[str],
+    organization: Organization,
+    include_viewers=False,
+    users_to_filter: typing.Optional[UserQuerySet] = None,
+) -> UserQuerySet:
     """
-    Parse ical file and return list of users found
-    NOTE: only grafana username will be used, consider adding grafana email and id
+    This method returns a `UserQuerySet`, filtered by users whose username, or case-insensitive e-mail,
+    is present in `usernames_from_ical`. If `include_viewers` is set to `True`, users are further filtered down
+    based on their granted permissions.
+
+    Parameters
+    ----------
+    usernames_from_ical : typing.List[str]
+        A list of usernames present in the ical feed
+    organization : apps.user_management.models.organization.Organization
+        The organization in question
+    include_viewers : bool
+        Whether or not the list should be further filtered to exclude users based on granted permissions
+    users_to_filter : typing.Optional[UserQuerySet]
+        Filter users without making SQL queries if users_to_filter arg is provided
+        users_to_filter is passed in `apps.schedules.ical_utils.get_oncall_users_for_multiple_schedules`
     """
     from apps.user_management.models import User
 
+    emails_from_ical = [username.lower() for username in usernames_from_ical]
+
+    if users_to_filter is not None:
+        return list(
+            {user for user in users_to_filter if user.username in usernames_from_ical or user.email in emails_from_ical}
+        )
+
     users_found_in_ical = organization.users
     if not include_viewers:
-        # TODO: this is a breaking change....
         users_found_in_ical = users_found_in_ical.filter(
             **User.build_permissions_query(RBACPermission.Permissions.SCHEDULES_WRITE, organization)
         )
 
-    user_emails = [v.lower() for v in usernames_from_ical]
     users_found_in_ical = users_found_in_ical.filter(
-        (Q(username__in=usernames_from_ical) | Q(email__lower__in=user_emails))
+        (Q(username__in=usernames_from_ical) | Q(email__lower__in=emails_from_ical))
     ).distinct()
 
-    # Here is the example how we extracted users previously, using slack fields too
-    # user_roles_found_in_ical = team.org_user_role.filter(role__in=[ROLE_ADMIN, ROLE_USER]).filter(
-    #     Q(
-    #         Q(amixr_user__slack_user_identities__slack_team_identity__amixr_team=team) &
-    #         Q(
-    #             Q(amixr_user__slack_user_identities__profile_display_name__in=usernames_from_ical) |
-    #             Q(amixr_user__slack_user_identities__cached_name__in=usernames_from_ical) |
-    #             Q(amixr_user__slack_user_identities__slack_id__in=[username.split(" ")[0] for username in
-    #                                                                usernames_from_ical]) |
-    #             Q(amixr_user__slack_user_identities__cached_slack_login__in=usernames_from_ical) |
-    #             Q(amixr_user__slack_user_identities__profile_real_name__in=usernames_from_ical)
-    #         )
-    #     )
-    #     |
-    #     Q(username__in=usernames_from_ical)
-    # ).annotate(is_deleted_sui=Subquery(slack_user_identity_subquery.values("deleted")[:1])).filter(
-    #     ~Q(is_deleted_sui=True) | Q(is_deleted_sui__isnull=True)).distinct()
-    # return user_roles_found_in_ical
-
     return users_found_in_ical
 
 
 @timed_lru_cache(timeout=100)
-def memoized_users_in_ical(usernames_from_ical, organization):
+def memoized_users_in_ical(usernames_from_ical: typing.List[str], organization: Organization) -> UserQuerySet:
     # using in-memory cache instead of redis to avoid pickling python objects
     return users_in_ical(usernames_from_ical, organization)
 
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-
-
 # used for display schedule events on web
 def list_of_oncall_shifts_from_ical(
     schedule,
@@ -288,17 +293,29 @@ def list_of_empty_shifts_in_schedule(schedule, start_date, end_date):
     return sorted(empty_shifts, key=lambda dt: dt.start)
 
 
-def list_users_to_notify_from_ical(schedule, events_datetime=None, include_viewers=False):
+def list_users_to_notify_from_ical(
+    schedule, events_datetime=None, include_viewers=False, users_to_filter=None
+) -> UserQuerySet:
     """
     Retrieve on-call users for the current time
     """
     events_datetime = events_datetime if events_datetime else timezone.datetime.now(timezone.utc)
     return list_users_to_notify_from_ical_for_period(
-        schedule, events_datetime, events_datetime, include_viewers=include_viewers
+        schedule,
+        events_datetime,
+        events_datetime,
+        include_viewers=include_viewers,
+        users_to_filter=users_to_filter,
     )
 
 
-def list_users_to_notify_from_ical_for_period(schedule, start_datetime, end_datetime, include_viewers=False):
+def list_users_to_notify_from_ical_for_period(
+    schedule,
+    start_datetime,
+    end_datetime,
+    include_viewers=False,
+    users_to_filter=None,
+) -> UserQuerySet:
     # get list of iCalendars from current iCal files. If there is more than one calendar, primary calendar will always
     # be the first
     calendars = schedule.get_icalendars()
@@ -316,7 +333,9 @@ def list_users_to_notify_from_ical_for_period(schedule, start_datetime, end_date
             parsed_ical_events.setdefault(current_priority, []).extend(current_usernames)
         # find users by usernames. if users are not found for shift, get users from lower priority
         for _, usernames in sorted(parsed_ical_events.items(), reverse=True):
-            users_found_in_ical = users_in_ical(usernames, schedule.organization, include_viewers=include_viewers)
+            users_found_in_ical = users_in_ical(
+                usernames, schedule.organization, include_viewers=include_viewers, users_to_filter=users_to_filter
+            )
             if users_found_in_ical:
                 break
         if users_found_in_ical:
@@ -325,6 +344,52 @@ def list_users_to_notify_from_ical_for_period(schedule, start_datetime, end_date
     return users_found_in_ical
 
 
+def get_oncall_users_for_multiple_schedules(
+    schedules, events_datetime=None
+) -> typing.Dict[OnCallSchedule, typing.List[User]]:
+    from apps.user_management.models import User
+
+    if events_datetime is None:
+        events_datetime = timezone.datetime.now(timezone.utc)
+
+    # Exit early if there are no schedules
+    if not schedules.exists():
+        return {}
+
+    # Assume all schedules from the queryset belong to the same organization
+    organization = schedules[0].organization
+
+    # Gather usernames from all events from all schedules
+    usernames = set()
+    for schedule in schedules.all():
+        calendars = schedule.get_icalendars()
+        for calendar in calendars:
+            if calendar is None:
+                continue
+            events = ical_events.get_events_from_ical_between(calendar, events_datetime, events_datetime)
+            for event in events:
+                current_usernames, _ = get_usernames_from_ical_event(event)
+                usernames.update(current_usernames)
+
+    # Fetch relevant users from the db
+    emails = [username.lower() for username in usernames]
+    users = organization.users.filter(
+        Q(**User.build_permissions_query(RBACPermission.Permissions.SCHEDULES_WRITE, organization))
+        & (Q(username__in=usernames) | Q(email__lower__in=emails))
+    )
+
+    # Get on-call users
+    oncall_users = {}
+    for schedule in schedules.all():
+        # pass user list to list_users_to_notify_from_ical
+        schedule_oncall_users = list_users_to_notify_from_ical(
+            schedule, events_datetime=events_datetime, users_to_filter=users
+        )
+        oncall_users.update({schedule.pk: schedule_oncall_users})
+
+    return oncall_users
+
+
 def parse_username_from_string(string):
     """
     Parse on-call shift user from the given string

engine/apps/schedules/models/on_call_schedule.py

@@ -18,10 +18,10 @@
 
 from apps.schedules.ical_utils import (
     fetch_ical_file_or_get_error,
+    get_oncall_users_for_multiple_schedules,
     list_of_empty_shifts_in_schedule,
     list_of_gaps_in_schedule,
     list_of_oncall_shifts_from_ical,
-    list_users_to_notify_from_ical,
 )
 from apps.schedules.models import CustomOnCallShift
 from common.public_primary_keys import generate_public_primary_key, increase_public_primary_key_length
@@ -43,19 +43,7 @@ def generate_public_primary_key_for_oncall_schedule_channel():
 
 class OnCallScheduleQuerySet(PolymorphicQuerySet):
     def get_oncall_users(self, events_datetime=None):
-        if events_datetime is None:
-            events_datetime = timezone.datetime.now(timezone.utc)
-
-        users = set()
-
-        for schedule in self.all():
-            schedule_oncall_users = list_users_to_notify_from_ical(schedule, events_datetime=events_datetime)
-            if schedule_oncall_users is None:
-                continue
-
-            users.update(schedule_oncall_users)
-
-        return list(users)
+        return get_oncall_users_for_multiple_schedules(self, events_datetime)
 
 
 class OnCallSchedule(PolymorphicModel):

engine/apps/schedules/tests/test_custom_on_call_shift.py

@@ -1295,7 +1295,7 @@ def test_get_oncall_users_for_empty_schedule(
     schedule = make_schedule(organization, schedule_class=OnCallScheduleCalendar)
     schedules = OnCallSchedule.objects.filter(pk=schedule.pk)
 
-    assert schedules.get_oncall_users() == []
+    assert schedules.get_oncall_users()[schedule.pk] == []
 
 
 @pytest.mark.django_db
@@ -1357,15 +1357,29 @@ def test_get_oncall_users_for_multiple_schedules(
 
     schedules = OnCallSchedule.objects.filter(pk__in=[schedule_1.pk, schedule_2.pk])
 
-    expected = set(schedules.get_oncall_users(events_datetime=now + timezone.timedelta(seconds=1)))
+    def _extract_oncall_users_from_schedules(schedules):
+        return set(user for schedule in schedules.values() for user in schedule)
+
+    expected = _extract_oncall_users_from_schedules(
+        schedules.get_oncall_users(events_datetime=now + timezone.timedelta(seconds=1))
+    )
     assert expected == {user_1, user_2}
 
-    expected = set(schedules.get_oncall_users(events_datetime=now + timezone.timedelta(minutes=10, seconds=1)))
+    expected = _extract_oncall_users_from_schedules(
+        schedules.get_oncall_users(events_datetime=now + timezone.timedelta(minutes=10, seconds=1))
+    )
     assert expected == {user_1, user_2, user_3}
 
-    assert schedules.get_oncall_users(events_datetime=now + timezone.timedelta(minutes=30, seconds=1)) == [user_3]
+    assert _extract_oncall_users_from_schedules(
+        schedules.get_oncall_users(events_datetime=now + timezone.timedelta(minutes=30, seconds=1))
+    ) == {user_3}
 
-    assert schedules.get_oncall_users(events_datetime=now + timezone.timedelta(minutes=40, seconds=1)) == []
+    assert (
+        _extract_oncall_users_from_schedules(
+            schedules.get_oncall_users(events_datetime=now + timezone.timedelta(minutes=40, seconds=1))
+        )
+        == set()
+    )
 
 
 @pytest.mark.django_db

engine/apps/slack/models/slack_usergroup.py

@@ -69,7 +69,7 @@ def can_be_updated(self) -> bool:
 
     @property
     def oncall_slack_user_identities(self):
-        users = self.oncall_schedules.get_oncall_users()
+        users = set(user for schedule in self.oncall_schedules.get_oncall_users().values() for user in schedule)
         slack_user_identities = [user.slack_user_identity for user in users if user.slack_user_identity is not None]
         return slack_user_identities
 

engine/apps/slack/tests/test_user_group.py

@@ -39,7 +39,7 @@ def test_oncall_slack_user_identities(
     )
     user_3 = make_user_for_organization(organization)
 
-    with patch.object(OnCallScheduleQuerySet, "get_oncall_users", return_value=[user_1, user_2, user_3]):
+    with patch.object(OnCallScheduleQuerySet, "get_oncall_users", return_value={"schedule1": [user_1, user_2, user_3]}):
         assert user_group.oncall_slack_user_identities == [slack_user_identity_1, slack_user_identity_2]