Merge pull request #711 from grafana/matiasb/schedules-list-filter-by-type

Add optional type filter to internal schedules endpoint

engine/apps/alerts/tests/test_notify_ical_schedule_shift.py

@@ -6,6 +6,7 @@ import pytz
 from django.utils import timezone
 
 from apps.alerts.tasks.notify_ical_schedule_shift import notify_ical_schedule_shift
+from apps.schedules.ical_utils import memoized_users_in_ical
 from apps.schedules.models import OnCallScheduleICal
 
 ICAL_DATA = """
@@ -80,6 +81,8 @@ def test_next_shift_notification_long_shifts(
     organization, _, _, _ = make_organization_and_user_with_slack_identities()
     make_user(organization=organization, username="user1")
     make_user(organization=organization, username="user2")
+    # clear users pks <-> organization cache (persisting between tests)
+    memoized_users_in_ical.cache_clear()
 
     ical_schedule = make_schedule(
         organization,

engine/apps/api/tests/test_schedules.py

@@ -134,6 +134,84 @@ def test_get_list_schedules(
     assert response.json() == expected_payload
 
 
+@pytest.mark.django_db
+def test_get_list_schedules_by_type(
+    schedule_internal_api_setup, make_escalation_chain, make_escalation_policy, make_user_auth_headers
+):
+    user, token, calendar_schedule, ical_schedule, web_schedule, slack_channel = schedule_internal_api_setup
+    client = APIClient()
+
+    # setup escalation chain linked to web schedule
+    escalation_chain = make_escalation_chain(user.organization)
+    make_escalation_policy(
+        escalation_chain=escalation_chain,
+        escalation_policy_step=EscalationPolicy.STEP_NOTIFY_SCHEDULE,
+        notify_schedule=web_schedule,
+    )
+
+    expected_payload = [
+        {
+            "id": calendar_schedule.public_primary_key,
+            "type": 0,
+            "team": None,
+            "name": "test_calendar_schedule",
+            "time_zone": "UTC",
+            "slack_channel": None,
+            "user_group": None,
+            "warnings": [],
+            "ical_url_overrides": None,
+            "on_call_now": [],
+            "has_gaps": False,
+            "mention_oncall_next": False,
+            "mention_oncall_start": True,
+            "notify_empty_oncall": 0,
+            "notify_oncall_shift_freq": 1,
+            "number_of_escalation_chains": 0,
+        },
+        {
+            "id": ical_schedule.public_primary_key,
+            "type": 1,
+            "team": None,
+            "name": "test_ical_schedule",
+            "ical_url_primary": ICAL_URL,
+            "ical_url_overrides": None,
+            "slack_channel": None,
+            "user_group": None,
+            "warnings": [],
+            "on_call_now": [],
+            "has_gaps": False,
+            "mention_oncall_next": False,
+            "mention_oncall_start": True,
+            "notify_empty_oncall": 0,
+            "notify_oncall_shift_freq": 1,
+            "number_of_escalation_chains": 0,
+        },
+        {
+            "id": web_schedule.public_primary_key,
+            "type": 2,
+            "time_zone": "UTC",
+            "team": None,
+            "name": "test_web_schedule",
+            "slack_channel": None,
+            "user_group": None,
+            "warnings": [],
+            "on_call_now": [],
+            "has_gaps": False,
+            "mention_oncall_next": False,
+            "mention_oncall_start": True,
+            "notify_empty_oncall": 0,
+            "notify_oncall_shift_freq": 1,
+            "number_of_escalation_chains": 1,
+        },
+    ]
+
+    for expected, schedule_type in enumerate(("api", "ical", "web")):
+        url = reverse("api-internal:schedule-list") + "?type={}".format(schedule_type)
+        response = client.get(url, format="json", **make_user_auth_headers(user, token))
+        assert response.status_code == status.HTTP_200_OK
+        assert response.json() == [expected_payload[expected]]
+
+
 @pytest.mark.django_db
 def test_get_detail_calendar_schedule(schedule_internal_api_setup, make_user_auth_headers):
     user, token, calendar_schedule, _, _, _ = schedule_internal_api_setup

engine/apps/api/views/schedule.py

@@ -24,7 +24,7 @@ from apps.api.serializers.schedule_polymorphic import (
 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.schedules.models import OnCallSchedule
+from apps.schedules.models import OnCallSchedule, OnCallScheduleCalendar, OnCallScheduleICal, OnCallScheduleWeb
 from apps.slack.models import SlackChannel
 from apps.slack.tasks import update_slack_user_group_for_schedules
 from common.api_helpers.exceptions import BadRequest, Conflict
@@ -42,6 +42,8 @@ EVENTS_FILTER_BY_ROTATION = "rotation"
 EVENTS_FILTER_BY_OVERRIDE = "override"
 EVENTS_FILTER_BY_FINAL = "final"
 
+SCHEDULE_TYPE_TO_CLASS = {"api": OnCallScheduleCalendar, "ical": OnCallScheduleICal, "web": OnCallScheduleWeb}
+
 
 class ScheduleView(
     TeamFilteringMixin,
@@ -123,6 +125,7 @@ class ScheduleView(
 
     def get_queryset(self):
         is_short_request = self.request.query_params.get("short", "false") == "true"
+        filter_by_type = self.request.query_params.get("type")
         organization = self.request.auth.organization
         queryset = OnCallSchedule.objects.filter(organization=organization, team=self.request.user.current_team).defer(
             # avoid requesting large text fields which are not used when listing schedules
@@ -134,6 +137,8 @@ class ScheduleView(
         if not is_short_request:
             queryset = self._annotate_queryset(queryset)
             queryset = self.serializer_class.setup_eager_loading(queryset)
+        if filter_by_type is not None and filter_by_type in SCHEDULE_TYPE_TO_CLASS:
+            queryset = queryset.filter().instance_of(SCHEDULE_TYPE_TO_CLASS[filter_by_type])
         return queryset
 
     def perform_create(self, serializer):