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 https://github.com/grafana/oncall-private/issues/1552 ## Checklist - [ ] Tests updated - [ ] Documentation added - [ ] `CHANGELOG.md` updated Co-authored-by: Vadim Stepanov <vadimkerr@gmail.com>
2023-01-25 11:08:09 +01:00 · 2023-01-25 11:08:09 +01:00 · 3cf2fcf660
commit 3cf2fcf660
parent de5d876d27
9 changed files with 148 additions and 63 deletions
--- a/CHANGELOG.md
+++ b/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

--- a/engine/apps/api/permissions/init.py
+++ b/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]
--- a/engine/apps/api/serializers/schedule_base.py
+++ b/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 @@ class ScheduleBaseSerializer(EagerLoadingMixin, serializers.ModelSerializer):
        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
--- a/engine/apps/api/views/schedule.py
+++ b/engine/apps/api/views/schedule.py
@ -104,9 +104,19 @@ class ScheduleView(

        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):
--- a/engine/apps/schedules/ical_utils.py
+++ b/engine/apps/schedules/ical_utils.py
@ -3,6 +3,7 @@ from __future__ import annotations
 import datetime
 import logging
 import re
+import typing
 from collections import namedtuple
 from typing import TYPE_CHECKING

@ -35,60 +36,64 @@ This module likely needs to refactored to be part of the OnCallSchedule module.
 """
 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
--- a/engine/apps/schedules/models/on_call_schedule.py
+++ b/engine/apps/schedules/models/on_call_schedule.py
@ -18,10 +18,10 @@ from polymorphic.query import PolymorphicQuerySet

 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):
--- a/engine/apps/schedules/tests/test_custom_on_call_shift.py
+++ b/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
--- a/engine/apps/slack/models/slack_usergroup.py
+++ b/engine/apps/slack/models/slack_usergroup.py
@ -69,7 +69,7 @@ class SlackUserGroup(models.Model):

    @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

--- a/engine/apps/slack/tests/test_user_group.py
+++ b/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]