Rework schedule related users, add endpoint (#1828)

Related to https://github.com/grafana/oncall/issues/1820.

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Add 2, 3 and 6 hours silence options
+- Add schedule related users endpoint to plugin API
 
 ## Fixed
 

engine/apps/api/serializers/user.py

@@ -186,6 +186,27 @@ class UserHiddenFieldsSerializer(UserSerializer):
         return ret
 
 
+class ScheduleUserSerializer(UserSerializer):
+    fields_to_keep = [
+        "pk",
+        "organization",
+        "email",
+        "username",
+        "name",
+        "avatar",
+        "avatar_full",
+        "timezone",
+        "working_hours",
+        "slack_user_identity",
+        "telegram_configuration",
+    ]
+
+    def to_representation(self, instance):
+        serialized = super(UserSerializer, self).to_representation(instance)
+        ret = {field: value for field, value in serialized.items() if field in self.fields_to_keep}
+        return ret
+
+
 class FastUserSerializer(serializers.ModelSerializer):
     pk = serializers.CharField(source="public_primary_key")
 

engine/apps/api/tests/test_schedules.py

@@ -11,6 +11,7 @@ from rest_framework.test import APIClient
 
 from apps.alerts.models import EscalationPolicy
 from apps.api.permissions import LegacyAccessControlRole
+from apps.api.serializers.user import ScheduleUserSerializer
 from apps.schedules.ical_utils import memoized_users_in_ical
 from apps.schedules.models import (
     CustomOnCallShift,
@@ -1274,6 +1275,68 @@ def test_next_shifts_per_user(
     assert returned_data == expected
 
 
+@pytest.mark.django_db
+def test_related_users(
+    make_organization_and_user_with_plugin_token,
+    make_user_for_organization,
+    make_user_auth_headers,
+    make_schedule,
+    make_on_call_shift,
+):
+    organization, admin, token = make_organization_and_user_with_plugin_token()
+    client = APIClient()
+
+    schedule = make_schedule(
+        organization,
+        schedule_class=OnCallScheduleWeb,
+        name="test_web_schedule",
+    )
+
+    tomorrow = timezone.now().replace(hour=0, minute=0, second=0, microsecond=0) + timezone.timedelta(days=1)
+    user_a, user_b, user_c, _ = (make_user_for_organization(organization, username=i) for i in "ABCD")
+    # clear users pks <-> organization cache (persisting between tests)
+    memoized_users_in_ical.cache_clear()
+
+    shifts = (
+        # user, priority, start time (h), duration (hs)
+        (user_a, 1, 8, 2),  # r1-1: 8-10 / A
+        (user_b, 2, 16, 2),  # r2-2: 16-18 / B
+    )
+    for user, priority, start_h, duration in shifts:
+        data = {
+            "start": tomorrow + timezone.timedelta(hours=start_h),
+            "rotation_start": tomorrow + timezone.timedelta(hours=start_h),
+            "duration": timezone.timedelta(hours=duration),
+            "priority_level": priority,
+            "frequency": CustomOnCallShift.FREQUENCY_DAILY,
+            "schedule": schedule,
+        }
+        on_call_shift = make_on_call_shift(
+            organization=organization, shift_type=CustomOnCallShift.TYPE_ROLLING_USERS_EVENT, **data
+        )
+        on_call_shift.add_rolling_users([[user]])
+
+    # override: 17-18 / C
+    override_data = {
+        "start": tomorrow + timezone.timedelta(hours=17),
+        "rotation_start": tomorrow + timezone.timedelta(hours=17),
+        "duration": timezone.timedelta(hours=1),
+        "schedule": schedule,
+    }
+    override = make_on_call_shift(
+        organization=organization, shift_type=CustomOnCallShift.TYPE_OVERRIDE, **override_data
+    )
+    override.add_rolling_users([[user_c]])
+    schedule.refresh_ical_file()
+
+    url = reverse("api-internal:schedule-related-users", kwargs={"pk": schedule.public_primary_key})
+    response = client.get(url, format="json", **make_user_auth_headers(admin, token))
+    assert response.status_code == status.HTTP_200_OK
+
+    expected = [ScheduleUserSerializer(u).data for u in (user_a, user_b, user_c)]
+    assert sorted(response.data["users"], key=lambda u: u["username"]) == sorted(expected, key=lambda u: u["username"])
+
+
 @pytest.mark.django_db
 def test_related_escalation_chains(
     make_organization_and_user_with_plugin_token,

engine/apps/api/views/schedule.py

@@ -23,6 +23,7 @@ from apps.api.serializers.schedule_polymorphic import (
     PolymorphicScheduleSerializer,
     PolymorphicScheduleUpdateSerializer,
 )
+from apps.api.serializers.user import ScheduleUserSerializer
 from apps.auth_token.auth import PluginAuthentication
 from apps.auth_token.constants import SCHEDULE_EXPORT_TOKEN_NAME
 from apps.auth_token.models import ScheduleExportAuthToken
@@ -80,6 +81,7 @@ class ScheduleView(
         "events": [RBACPermission.Permissions.SCHEDULES_READ],
         "filter_events": [RBACPermission.Permissions.SCHEDULES_READ],
         "next_shifts_per_user": [RBACPermission.Permissions.SCHEDULES_READ],
+        "related_users": [RBACPermission.Permissions.SCHEDULES_READ],
         "quality": [RBACPermission.Permissions.SCHEDULES_READ],
         "notify_empty_oncall_options": [RBACPermission.Permissions.SCHEDULES_READ],
         "notify_oncall_shift_freq_options": [RBACPermission.Permissions.SCHEDULES_READ],
@@ -336,7 +338,7 @@ class ScheduleView(
         schedule = self.get_object()
         events = schedule.final_events(user_tz, starting_date, days=30)
 
-        users = {u: None for u in schedule.related_users()}
+        users = {u.public_primary_key: None for u in schedule.related_users()}
         for e in events:
             user = e["users"][0]["pk"] if e["users"] else None
             if user is not None and users.get(user) is None and e["end"] > now:
@@ -345,6 +347,13 @@ class ScheduleView(
         result = {"users": users}
         return Response(result, status=status.HTTP_200_OK)
 
+    @action(detail=True, methods=["get"])
+    def related_users(self, request, pk):
+        schedule = self.get_object()
+        serializer = ScheduleUserSerializer(schedule.related_users(), many=True)
+        result = {"users": serializer.data}
+        return Response(result, status=status.HTTP_200_OK)
+
     @action(detail=True, methods=["get"])
     def related_escalation_chains(self, request, pk):
         """Return escalation chains associated to schedule."""

engine/apps/schedules/models/on_call_schedule.py

@@ -1,6 +1,6 @@
 import datetime
-import functools
 import itertools
+import re
 from collections import defaultdict
 from enum import Enum
 from typing import Iterable, Optional, TypedDict
@@ -45,6 +45,9 @@ from apps.user_management.models import User
 from common.database import NON_POLYMORPHIC_CASCADE, NON_POLYMORPHIC_SET_NULL
 from common.public_primary_keys import generate_public_primary_key, increase_public_primary_key_length
 
+RE_ICAL_SEARCH_USERNAME = r"SUMMARY:(\[L[0-9]+\] )?{}"
+RE_ICAL_FETCH_USERNAME = re.compile(r"SUMMARY:(?:\[L[0-9]+\] )?(\w+)")
+
 
 # Utility classes for schedule quality report
 class QualityReportCommentType(str, Enum):
@@ -88,7 +91,7 @@ class OnCallScheduleQuerySet(PolymorphicQuerySet):
         return get_oncall_users_for_multiple_schedules(self, events_datetime)
 
     def related_to_user(self, user):
-        username_regex = r"SUMMARY:(\[L[0-9]+\] )?{}".format(user.username)
+        username_regex = RE_ICAL_SEARCH_USERNAME.format(user.username)
         return self.filter(
             Q(cached_ical_file_primary__regex=username_regex)
             | Q(cached_ical_file_primary__contains=user.email)
@@ -244,8 +247,13 @@ class OnCallSchedule(PolymorphicModel):
         self.save(update_fields=["cached_ical_file_overrides", "prev_ical_file_overrides"])
 
     def related_users(self):
-        """Return public primary keys for all users referenced in the schedule."""
+        """Return users referenced in the schedule."""
-        return set()
+        usernames = []
+        if self.cached_ical_file_primary:
+            usernames += RE_ICAL_FETCH_USERNAME.findall(self.cached_ical_file_primary)
+        if self.cached_ical_file_overrides:
+            usernames += RE_ICAL_FETCH_USERNAME.findall(self.cached_ical_file_overrides)
+        return self.organization.users.filter(username__in=usernames)
 
     def filter_events(
         self,
@@ -981,22 +989,6 @@ class OnCallScheduleWeb(OnCallSchedule):
         self.cached_ical_file_overrides = self._generate_ical_file_overrides()
         self.save(update_fields=["cached_ical_file_overrides", "prev_ical_file_overrides"])
 
-    def related_users(self):
-        """Return public primary keys for all users referenced in the schedule."""
-        rolling_users = self.custom_shifts.values_list("rolling_users", flat=True)
-        users = functools.reduce(
-            set.union,
-            (
-                set(g.values())
-                for rolling_groups in rolling_users
-                if rolling_groups is not None
-                for g in rolling_groups
-                if g is not None
-            ),
-            set(),
-        )
-        return users
-
     # Insight logs
     @property
     def insight_logs_type_verbal(self):

engine/apps/schedules/tests/test_on_call_schedule.py

@@ -894,9 +894,10 @@ def test_schedule_related_users_empty_schedule(make_organization, make_schedule)
         schedule_class=OnCallScheduleWeb,
         name="test_web_schedule",
     )
+    schedule.refresh_ical_file()
 
     users = schedule.related_users()
-    assert users == set()
+    assert set(users) == set()
 
 
 @pytest.mark.django_db
@@ -930,7 +931,7 @@ def test_schedule_related_users(make_organization, make_user_for_organization, m
             "schedule": schedule,
         }
         on_call_shift = make_on_call_shift(
-            organization=organization, shift_type=CustomOnCallShift.TYPE_RECURRENT_EVENT, **data
+            organization=organization, shift_type=CustomOnCallShift.TYPE_ROLLING_USERS_EVENT, **data
         )
         on_call_shift.add_rolling_users([[user]])
 
@@ -946,9 +947,11 @@ def test_schedule_related_users(make_organization, make_user_for_organization, m
     )
     override.add_rolling_users([[user_e]])
 
+    schedule.refresh_ical_file()
     schedule.refresh_from_db()
+
     users = schedule.related_users()
-    assert users == set(u.public_primary_key for u in [user_a, user_d, user_e])
+    assert set(users) == set([user_a, user_d, user_e])
 
 
 @pytest.mark.django_db(transaction=True)