diff --git a/slack_sdk/web/base_client.py b/slack_sdk/web/base_client.py index 1f5ad58c7..0021be0a7 100644 --- a/slack_sdk/web/base_client.py +++ b/slack_sdk/web/base_client.py @@ -588,18 +588,117 @@ def _upload_file( proxy: Optional[str], ssl: Optional[SSLContext], ) -> FileUploadV2Result: - """Upload a file using the issued upload URL""" - result = _upload_file_via_v2_url( - url=url, - data=data, - logger=logger, - timeout=timeout, - proxy=proxy, - ssl=ssl, - ) - return FileUploadV2Result( - status=result.get("status"), # type: ignore[arg-type] - body=result.get("body"), # type: ignore[arg-type] + """Upload a file using the issued upload URL. + + Unlike files.getUploadURLExternal / files.completeUploadExternal, this POST + does not go through api_call(). Apply the same retry handlers so transient + HTTP errors (e.g. 504 on files.slack.com) can be retried. + """ + retry_request = RetryHttpRequest(method="POST", url=url, headers={}, data=data) + retry_state = RetryState() + last_error: Optional[Exception] = None + last_result: Optional[Dict[str, Any]] = None + counter_for_safety = 0 + while counter_for_safety < 100: + counter_for_safety += 1 + retry_state.next_attempt_requested = False + try: + result = _upload_file_via_v2_url( + url=url, + data=data, + logger=logger, + timeout=timeout, + proxy=proxy, + ssl=ssl, + ) + last_result = result + retry_response = self._retry_response_from_upload_result(result) + for handler in self.retry_handlers: + if handler.can_retry(state=retry_state, request=retry_request, response=retry_response): + if logger.level <= logging.DEBUG: + logger.info(f"A retry handler found: {type(handler).__name__} for POST {url}") + handler.prepare_for_next_attempt(state=retry_state, request=retry_request, response=retry_response) + break + if retry_state.next_attempt_requested is False: + return FileUploadV2Result( + status=result.get("status"), # type: ignore[arg-type] + body=result.get("body"), # type: ignore[arg-type] + ) + except HTTPError as e: + last_error = e + response_headers = dict(e.headers.items()) if e.headers is not None else {} + charset = "utf-8" + if e.headers is not None: + charset = e.headers.get_content_charset() or "utf-8" + response_body = e.read().decode(charset) + retry_response = RetryHttpResponse( + status_code=e.code, + headers={k: [v] for k, v in response_headers.items()}, + data=response_body.encode("utf-8") if response_body is not None else None, + ) + for handler in self.retry_handlers: + if handler.can_retry( + state=retry_state, + request=retry_request, + response=retry_response, + error=e, + ): + if logger.level <= logging.DEBUG: + logger.info(f"A retry handler found: {type(handler).__name__} for POST {url} - {e}") + handler.prepare_for_next_attempt( + state=retry_state, + request=retry_request, + response=retry_response, + error=e, + ) + break + if retry_state.next_attempt_requested is False: + raise e + except Exception as err: + last_error = err + logger.error(f"Failed to upload a file to Slack: {err}") + for handler in self.retry_handlers: + if handler.can_retry( + state=retry_state, + request=retry_request, + response=None, + error=err, + ): + if logger.level <= logging.DEBUG: + logger.info(f"A retry handler found: {type(handler).__name__} for POST {url} - {err}") + handler.prepare_for_next_attempt( + state=retry_state, + request=retry_request, + response=None, + error=err, + ) + logger.info(f"Going to retry the same request: POST {url}") + break + if retry_state.next_attempt_requested is False: + raise err + + if last_result is not None: + return FileUploadV2Result( + status=last_result.get("status"), # type: ignore[arg-type] + body=last_result.get("body"), # type: ignore[arg-type] + ) + raise last_error # type: ignore[misc] + + @staticmethod + def _retry_response_from_upload_result(result: Dict[str, Any]) -> RetryHttpResponse: + raw_headers = result.get("headers") + if raw_headers is None: + headers: Dict[str, Any] = {} + elif hasattr(raw_headers, "items"): + headers = dict(raw_headers.items()) + else: + headers = {} + body = result.get("body") + body_bytes = body.encode("utf-8") if isinstance(body, str) else body + return RetryHttpResponse( + status_code=result.get("status") or 0, + headers=headers, + data=body_bytes, ) # ================================================================= diff --git a/tests/slack_sdk/web/test_web_client_file_upload_retry.py b/tests/slack_sdk/web/test_web_client_file_upload_retry.py new file mode 100644 index 000000000..3f9e5c90c --- /dev/null +++ b/tests/slack_sdk/web/test_web_client_file_upload_retry.py @@ -0,0 +1,100 @@ +import logging +import threading +import unittest +from http.server import BaseHTTPRequestHandler, HTTPServer +from urllib.error import HTTPError + +from slack_sdk.http_retry import RetryHandler +from slack_sdk.http_retry.handler import default_interval_calculator +from slack_sdk.http_retry.interval_calculator import RetryIntervalCalculator +from slack_sdk.web import WebClient + + +class GatewayErrorRetryHandler(RetryHandler): + """Retry 5xx responses from the files.slack.com upload POST.""" + + def __init__( + self, + max_retry_count: int = 2, + interval_calculator: RetryIntervalCalculator = default_interval_calculator, + ): + super().__init__(max_retry_count, interval_calculator) + self.call_count = 0 + + def _can_retry( + self, + *, + state, + request, + response, + error, + ) -> bool: + self.call_count += 1 + return response is not None and response.status_code >= 500 + + +class _UploadHandler(BaseHTTPRequestHandler): + attempts = 0 + fail_times = 1 + fail_status = 504 + + def do_POST(self): + _UploadHandler.attempts += 1 + length = int(self.headers.get("Content-Length") or 0) + if length: + self.rfile.read(length) + if _UploadHandler.attempts <= _UploadHandler.fail_times: + self.send_response(_UploadHandler.fail_status) + self.send_header("Content-Type", "text/plain") + self.end_headers() + self.wfile.write(b"gateway timeout") + return + self.send_response(200) + self.send_header("Content-Type", "text/plain") + self.end_headers() + self.wfile.write(b"ok") + + def log_message(self, format, *args): + pass + + +class TestWebClient_FileUploadRetry(unittest.TestCase): + def setUp(self): + _UploadHandler.attempts = 0 + _UploadHandler.fail_times = 1 + _UploadHandler.fail_status = 504 + self.server = HTTPServer(("127.0.0.1", 0), _UploadHandler) + thread = threading.Thread(target=self.server.serve_forever, daemon=True) + thread.start() + host, port = self.server.server_address + self.upload_url = f"http://127.0.0.1:{port}/upload" + + def tearDown(self): + self.server.shutdown() + self.server.server_close() + + def _upload(self, client: WebClient): + return client._upload_file( + url=self.upload_url, + data=b"hello file", + logger=logging.getLogger(__name__), + timeout=5, + proxy=None, + ssl=None, + ) + + def test_upload_retries_gateway_error_with_handler(self): + retry_handler = GatewayErrorRetryHandler(max_retry_count=2) + client = WebClient(token="xoxb-test", retry_handlers=[retry_handler]) + result = self._upload(client) + self.assertEqual(200, result.status) + self.assertEqual("ok", result.body) + self.assertEqual(2, _UploadHandler.attempts) + self.assertGreaterEqual(retry_handler.call_count, 1) + + def test_upload_without_matching_handler_surfaces_error(self): + client = WebClient(token="xoxb-test", retry_handlers=[]) + with self.assertRaises(HTTPError) as cm: + self._upload(client) + self.assertEqual(504, cm.exception.code) + self.assertEqual(1, _UploadHandler.attempts)