diff --git a/engine/apps/alerts/tasks/notify_ical_schedule_shift.py b/engine/apps/alerts/tasks/notify_ical_schedule_shift.py index bf76d53d..9294353b 100644 --- a/engine/apps/alerts/tasks/notify_ical_schedule_shift.py +++ b/engine/apps/alerts/tasks/notify_ical_schedule_shift.py @@ -155,7 +155,10 @@ def notify_ical_schedule_shift(schedule_pk): if len(new_shifts) > 0 or schedule.empty_oncall: task_logger.info(f"new_shifts: {new_shifts}") if schedule.notify_oncall_shift_freq != OnCallSchedule.NotifyOnCallShiftFreq.NEVER: - slack_client = SlackClient(schedule.organization.slack_team_identity) + slack_client = SlackClient( + schedule.organization.slack_team_identity, + enable_ratelimit_retry=True, + ) step = scenario_step.ScenarioStep.get_step("schedules", "EditScheduleShiftNotifyStep") report_blocks = step.get_report_blocks_ical(new_shifts, upcoming_shifts, schedule, schedule.empty_oncall) diff --git a/engine/apps/integrations/tasks.py b/engine/apps/integrations/tasks.py index 83d7a0ff..46101b93 100644 --- a/engine/apps/integrations/tasks.py +++ b/engine/apps/integrations/tasks.py @@ -162,7 +162,7 @@ def notify_about_integration_ratelimit_in_slack(organization_id, text, **kwargs) slack_team_identity = organization.slack_team_identity if slack_team_identity is not None: try: - sc = SlackClient(slack_team_identity) + sc = SlackClient(slack_team_identity, enable_ratelimit_retry=True) sc.chat_postMessage(channel=organization.general_log_channel_id, text=text, team=slack_team_identity) except SlackAPIError as e: logger.warning(f"Slack exception {e} while sending message for organization {organization_id}") diff --git a/engine/apps/slack/client.py b/engine/apps/slack/client.py index 9947de95..ca05d498 100644 --- a/engine/apps/slack/client.py +++ b/engine/apps/slack/client.py @@ -44,16 +44,22 @@ class SlackServerErrorRetryHandler(RetryHandler): # retries when HTTP status 429 is returned using the Retry-After header information -rate_limit_handler = RateLimitErrorRetryHandler(max_retry_count=2) +rate_limit_handler = RateLimitErrorRetryHandler(max_retry_count=1) server_error_retry_handler = SlackServerErrorRetryHandler(max_retry_count=2) class SlackClient(WebClient): - def __init__(self, slack_team_identity: "SlackTeamIdentity", timeout: int = 30) -> None: + def __init__( + self, slack_team_identity: "SlackTeamIdentity", enable_ratelimit_retry=False, timeout: int = 30 + ) -> None: + retry_handlers = default_retry_handlers() + [server_error_retry_handler] + if enable_ratelimit_retry: + retry_handlers += [rate_limit_handler] + super().__init__( token=slack_team_identity.bot_access_token, timeout=timeout, - retry_handlers=default_retry_handlers() + [server_error_retry_handler, rate_limit_handler], + retry_handlers=retry_handlers, ) self.slack_team_identity = slack_team_identity diff --git a/engine/apps/slack/models/slack_usergroup.py b/engine/apps/slack/models/slack_usergroup.py index a99c9ee8..3bdda3de 100644 --- a/engine/apps/slack/models/slack_usergroup.py +++ b/engine/apps/slack/models/slack_usergroup.py @@ -116,7 +116,7 @@ class SlackUserGroup(models.Model): pass def update_members(self, slack_ids): - sc = SlackClient(self.slack_team_identity) + sc = SlackClient(self.slack_team_identity, enable_ratelimit_retry=True) try: sc.usergroups_users_update(usergroup=self.slack_id, users=slack_ids) diff --git a/engine/apps/slack/tasks.py b/engine/apps/slack/tasks.py index f4a04a57..0a81c889 100644 --- a/engine/apps/slack/tasks.py +++ b/engine/apps/slack/tasks.py @@ -154,7 +154,7 @@ def send_message_to_thread_if_bot_not_in_channel(alert_group_pk, slack_team_iden slack_team_identity = SlackTeamIdentity.objects.get(pk=slack_team_identity_pk) alert_group = AlertGroup.objects.get(pk=alert_group_pk) - sc = SlackClient(slack_team_identity) + sc = SlackClient(slack_team_identity, enable_ratelimit_retry=True) bot_user_id = slack_team_identity.bot_user_id members = slack_team_identity.get_conversation_members(sc, channel_id) @@ -341,7 +341,7 @@ def populate_slack_usergroups_for_team(slack_team_identity_id): from apps.slack.models import SlackTeamIdentity, SlackUserGroup slack_team_identity = SlackTeamIdentity.objects.get(pk=slack_team_identity_id) - sc = SlackClient(slack_team_identity) + sc = SlackClient(slack_team_identity, enable_ratelimit_retry=True) try: usergroups = sc.usergroups_list()["usergroups"] @@ -454,7 +454,7 @@ def populate_slack_channels_for_team(slack_team_identity_id: int, cursor: Option from apps.slack.models import SlackChannel, SlackTeamIdentity slack_team_identity = SlackTeamIdentity.objects.get(pk=slack_team_identity_id) - sc = SlackClient(slack_team_identity) + sc = SlackClient(slack_team_identity, enable_ratelimit_retry=True) active_task_id_key = get_populate_slack_channel_task_id_key(slack_team_identity_id) active_task_id = cache.get(active_task_id_key) diff --git a/engine/apps/slack/tests/test_slack_client.py b/engine/apps/slack/tests/test_slack_client.py index 9ad935d7..7214aa96 100644 --- a/engine/apps/slack/tests/test_slack_client.py +++ b/engine/apps/slack/tests/test_slack_client.py @@ -167,6 +167,25 @@ def test_slack_client_ratelimit(monkeypatch, error, make_organization_with_slack _, slack_team_identity = make_organization_with_slack_team_identity() client = SlackClient(slack_team_identity) + return_value = {"status": 429, "body": json.dumps({"ok": False, "error": error}), "headers": {"Retry-After": "1"}} + with patch( + "slack_sdk.web.base_client.BaseClient._perform_urllib_http_request_internal", return_value=return_value + ) as mock_request: + with pytest.raises(SlackAPIRatelimitError): + client.api_call("auth.test") + + # no slack built-in retries by default + assert len(mock_request.mock_calls) == 1 + + +@pytest.mark.parametrize("error", ["ratelimited", "rate_limited", "message_limit_exceeded"]) +@pytest.mark.django_db +def test_slack_client_ratelimit_enable_retry(monkeypatch, error, make_organization_with_slack_team_identity): + monkeypatch.undo() # undo engine.conftest.mock_slack_api_call + + _, slack_team_identity = make_organization_with_slack_team_identity() + client = SlackClient(slack_team_identity, enable_ratelimit_retry=True) + return_value = {"status": 429, "body": json.dumps({"ok": False, "error": error}), "headers": {"Retry-After": "1"}} with patch( "slack_sdk.web.base_client.BaseClient._perform_urllib_http_request_internal", return_value=return_value @@ -174,7 +193,7 @@ def test_slack_client_ratelimit(monkeypatch, error, make_organization_with_slack with pytest.raises(SlackAPIRatelimitError) as exc_info: client.api_call("auth.test") - assert len(mock_request.mock_calls) == 3 + assert len(mock_request.mock_calls) == 2 assert exc_info.value.retry_after == 1