From 3cf2fcf66031fb5acc1e0f22fa34e5432a284d26 Mon Sep 17 00:00:00 2001 From: Joey Orlando Date: Wed, 25 Jan 2023 11:08:09 +0100 Subject: [PATCH] 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 --- CHANGELOG.md | 1 + engine/apps/api/permissions/__init__.py | 11 ++ engine/apps/api/serializers/schedule_base.py | 10 +- engine/apps/api/views/schedule.py | 10 ++ engine/apps/schedules/ical_utils.py | 135 +++++++++++++----- .../apps/schedules/models/on_call_schedule.py | 16 +-- .../tests/test_custom_on_call_shift.py | 24 +++- engine/apps/slack/models/slack_usergroup.py | 2 +- engine/apps/slack/tests/test_user_group.py | 2 +- 9 files changed, 148 insertions(+), 63 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d3a4e049..3fef8af4 100644 --- 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 diff --git a/engine/apps/api/permissions/__init__.py b/engine/apps/api/permissions/__init__.py index c2fdf1f0..60caf3dc 100644 --- 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] diff --git a/engine/apps/api/serializers/schedule_base.py b/engine/apps/api/serializers/schedule_base.py index c2eb6ea6..ffe6f6ac 100644 --- 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 diff --git a/engine/apps/api/views/schedule.py b/engine/apps/api/views/schedule.py index fa61c5dd..7d475f7f 100644 --- 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): diff --git a/engine/apps/schedules/ical_utils.py b/engine/apps/schedules/ical_utils.py index 962766d5..fd6cd341 100644 --- 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 diff --git a/engine/apps/schedules/models/on_call_schedule.py b/engine/apps/schedules/models/on_call_schedule.py index d746e81c..df012d59 100644 --- 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): diff --git a/engine/apps/schedules/tests/test_custom_on_call_shift.py b/engine/apps/schedules/tests/test_custom_on_call_shift.py index 0a40a879..7974ad91 100644 --- 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 diff --git a/engine/apps/slack/models/slack_usergroup.py b/engine/apps/slack/models/slack_usergroup.py index a319f133..6cd98b66 100644 --- 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 diff --git a/engine/apps/slack/tests/test_user_group.py b/engine/apps/slack/tests/test_user_group.py index e03d1a72..d1a53c14 100644 --- 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]