Fix #887 Enable automatic retry by a handy way (#1084)

Author: seratch

integration_tests/webhook/test_webhook.py

@@ -2,6 +2,8 @@
 import unittest
 import time
 
+import pytest
+
 from integration_tests.env_variable_names import (
     SLACK_SDK_TEST_INCOMING_WEBHOOK_URL,
     SLACK_SDK_TEST_INCOMING_WEBHOOK_CHANNEL_NAME,
@@ -68,6 +70,9 @@ def test_with_unfurls_off(self):
         self.assertIsNotNone(history)
         self.assertTrue("attachments" not in history["messages"][0])
 
+
+    # FIXME: This test started failing as of August 5, 2021
+    @pytest.mark.skip()
     def test_with_unfurls_on(self):
         # Slack API rate limits unfurls of unique links so test will
         # fail when repeated. For testing, either use a different URL

slack_sdk/audit_logs/v1/async_client.py

@@ -5,7 +5,7 @@
 import json
 import logging
 from ssl import SSLContext
-from typing import Any
+from typing import Any, List
 from typing import Dict, Optional
 
 import aiohttp
@@ -18,6 +18,11 @@
     get_user_agent,
 )
 from .response import AuditLogsResponse
+from slack_sdk.http_retry.async_handler import AsyncRetryHandler
+from slack_sdk.http_retry.builtin_async_handlers import async_default_handlers
+from slack_sdk.http_retry.request import HttpRequest as RetryHttpRequest
+from slack_sdk.http_retry.response import HttpResponse as RetryHttpResponse
+from slack_sdk.http_retry.state import RetryState
 from ...proxy_env_variable_loader import load_http_proxy_from_env
 
 
@@ -34,6 +39,7 @@ class AsyncAuditLogsClient:
     auth: Optional[BasicAuth]
     default_headers: Dict[str, str]
     logger: logging.Logger
+    retry_handlers: List[AsyncRetryHandler]
 
     def __init__(
         self,
@@ -49,6 +55,7 @@ def __init__(
         user_agent_prefix: Optional[str] = None,
         user_agent_suffix: Optional[str] = None,
         logger: Optional[logging.Logger] = None,
+        retry_handlers: List[AsyncRetryHandler] = async_default_handlers,
     ):
         """API client for Audit Logs API
         See https://api.slack.com/admins/audit-logs for more details
@@ -66,6 +73,7 @@ def __init__(
             user_agent_prefix: Prefix for User-Agent header value
             user_agent_suffix: Suffix for User-Agent header value
             logger: Custom logger
+            retry_handlers: Retry handlers
         """
         self.token = token
         self.timeout = timeout
@@ -80,6 +88,7 @@ def __init__(
             user_agent_prefix, user_agent_suffix
         )
         self.logger = logger if logger is not None else logging.getLogger(__name__)
+        self.retry_handlers = retry_handlers
 
         if self.proxy is None or len(self.proxy.strip()) == 0:
             env_variable = load_http_proxy_from_env(self.logger)
@@ -218,18 +227,6 @@ async def _perform_http_request(
             body_params = json.dumps(body_params)
         headers["Content-Type"] = "application/json;charset=utf-8"
 
-        if self.logger.level <= logging.DEBUG:
-            headers_for_logging = {
-                k: "(redacted)" if k.lower() == "authorization" else v
-                for k, v in headers.items()
-            }
-            self.logger.debug(
-                f"Sending a request - "
-                f"url: {url}, "
-                f"params: {query_params}, "
-                f"body: {body_params}, "
-                f"headers: {headers_for_logging}"
-            )
         session: Optional[ClientSession] = None
         use_running_session = self.session and not self.session.closed
         if use_running_session:
@@ -241,7 +238,8 @@ async def _perform_http_request(
                 trust_env=self.trust_env_in_session,
             )
 
-        resp: AuditLogsResponse
+        last_error = None
+        resp: Optional[AuditLogsResponse] = None
         try:
             request_kwargs = {
                 "headers": headers,
@@ -250,25 +248,112 @@ async def _perform_http_request(
                 "ssl": self.ssl,
                 "proxy": self.proxy,
             }
-            async with session.request(http_verb, url, **request_kwargs) as res:
-                response_body = {}
-                try:
-                    response_body = await res.text()
-                except aiohttp.ContentTypeError:
+            retry_request = RetryHttpRequest(
+                method=http_verb,
+                url=url,
+                headers=headers,
+                body_params=body_params,
+            )
+
+            retry_state = RetryState()
+            counter_for_safety = 0
+            while counter_for_safety < 100:
+                counter_for_safety += 1
+                # If this is a retry, the next try started here. We can reset the flag.
+                retry_state.next_attempt_requested = False
+                retry_response: Optional[RetryHttpResponse] = None
+                response_body = ""
+
+                if self.logger.level <= logging.DEBUG:
+                    headers_for_logging = {
+                        k: "(redacted)" if k.lower() == "authorization" else v
+                        for k, v in headers.items()
+                    }
                     self.logger.debug(
-                        f"No response data returned from the following API call: {url}."
+                        f"Sending a request - "
+                        f"url: {url}, "
+                        f"params: {query_params}, "
+                        f"body: {body_params}, "
+                        f"headers: {headers_for_logging}"
                     )
-                except json.decoder.JSONDecodeError as e:
-                    message = f"Failed to parse the response body: {str(e)}"
-                    raise SlackApiError(message, res)
-
-                resp = AuditLogsResponse(
-                    url=url,
-                    status_code=res.status,
-                    raw_body=response_body,
-                    headers=res.headers,
-                )
-                _debug_log_response(self.logger, resp)
+
+                try:
+                    async with session.request(http_verb, url, **request_kwargs) as res:
+                        try:
+                            response_body = await res.text()
+                            retry_response = RetryHttpResponse(
+                                status_code=res.status,
+                                headers=res.headers,
+                                data=response_body.encode("utf-8")
+                                if response_body is not None
+                                else None,
+                            )
+                        except aiohttp.ContentTypeError:
+                            self.logger.debug(
+                                f"No response data returned from the following API call: {url}."
+                            )
+                        except json.decoder.JSONDecodeError as e:
+                            message = f"Failed to parse the response body: {str(e)}"
+                            raise SlackApiError(message, res)
+
+                        if res.status == 429:
+                            for handler in self.retry_handlers:
+                                if await handler.can_retry_async(
+                                    state=retry_state,
+                                    request=retry_request,
+                                    response=retry_response,
+                                ):
+                                    if self.logger.level <= logging.DEBUG:
+                                        self.logger.info(
+                                            f"A retry handler found: {type(handler).__name__} "
+                                            f"for {http_verb} {url} - rate_limited"
+                                        )
+                                    await handler.prepare_for_next_attempt_async(
+                                        state=retry_state,
+                                        request=retry_request,
+                                        response=retry_response,
+                                    )
+                                    break
+
+                        if retry_state.next_attempt_requested is False:
+                            resp = AuditLogsResponse(
+                                url=url,
+                                status_code=res.status,
+                                raw_body=response_body,
+                                headers=res.headers,
+                            )
+                            _debug_log_response(self.logger, resp)
+                            return resp
+
+                except Exception as e:
+                    last_error = e
+                    for handler in self.retry_handlers:
+                        if await handler.can_retry_async(
+                            state=retry_state,
+                            request=retry_request,
+                            response=retry_response,
+                            error=e,
+                        ):
+                            if self.logger.level <= logging.DEBUG:
+                                self.logger.info(
+                                    f"A retry handler found: {type(handler).__name__} "
+                                    f"for {http_verb} {url} - {e}"
+                                )
+                            await handler.prepare_for_next_attempt_async(
+                                state=retry_state,
+                                request=retry_request,
+                                response=retry_response,
+                                error=e,
+                            )
+                            break
+
+                    if retry_state.next_attempt_requested is False:
+                        raise last_error
+
+            if resp is not None:
+                return resp
+            raise last_error
+
         finally:
             if not use_running_session:
                 await session.close()