From 4d93493bd1406fce8abe297ced85e50ac8a81f03 Mon Sep 17 00:00:00 2001 From: Michael Sasser Date: Sat, 28 Sep 2024 16:02:26 +0200 Subject: [PATCH] fix: Break out the the sync status code handler from `_request` of the API handler into `handle_sync_response_status_code` and re-use it in `streamed_download` --- matrixctl/handlers/api.py | 143 +++++++++++++++++++++++++------------- news/835.bugfix.rst | 1 + 2 files changed, 94 insertions(+), 50 deletions(-) create mode 100644 news/835.bugfix.rst diff --git a/matrixctl/handlers/api.py b/matrixctl/handlers/api.py index c07e4614..255fc17b 100644 --- a/matrixctl/handlers/api.py +++ b/matrixctl/handlers/api.py @@ -49,6 +49,17 @@ HTTP_RETURN_CODE_302: int = 302 HTTP_RETURN_CODE_404: int = 404 +DEFAULT_SUCCESS_CODES: tuple[int, ...] = ( + 200, + 201, + 202, + 203, + 204, + 205, + 206, + 207, + 226, +) logger = logging.getLogger(__name__) @@ -95,17 +106,7 @@ class RequestBuilder: headers: dict[str, str] = {} # noqa: RUF012 concurrent_limit: int = 4 timeout: float = 5.0 # seconds - success_codes: tuple[int, ...] = ( - 200, - 201, - 202, - 203, - 204, - 205, - 206, - 207, - 226, - ) + success_codes: tuple[int, ...] = DEFAULT_SUCCESS_CODES @property def headers_with_auth(self) -> dict[str, str]: @@ -541,36 +542,27 @@ async def gen_async_request() -> list[httpx.Response] | httpx.Response: return asyncio.run(gen_async_request()) -def _request(request_config: RequestBuilder) -> httpx.Response: - """Send an synchronous request to the synapse API and receive a response. +def handle_sync_response_status_code( + response: httpx.Response, + success_codes: tuple[int, ...] | None = None, +) -> None: + """Handle the response status code of a synchronous request. Attributes ---------- - req : matrixctl.handlers.api.RequestBuilder - An instance of an RequestBuilder + response : https.Response + The response of the synchronous request. + + success_codes : tuple of int, optional + A tuple of success codes. For example: 200, 201, 202, 203. + If the parameter is not set, the default success codes are used. Returns ------- - response : httpx.Response - Returns the response - + None + The function either returns None or raises an exception. """ - - logger.debug("repr: %s", repr(request_config)) - - # There is some weird stuff going on in httpx. It is set to None by default - with httpx.Client(http2=True, timeout=request_config.timeout) as client: - response: httpx.Response = client.request( - method=request_config.method, - data=request_config.data, # type: ignore # noqa: PGH003 - json=request_config.json, - content=request_config.content, # type: ignore # noqa: PGH003 - url=str(request_config), - params=request_config.params, - headers=request_config.headers_with_auth, - follow_redirects=False, - ) - + success_codes = success_codes or DEFAULT_SUCCESS_CODES if response.status_code == HTTP_RETURN_CODE_302: logger.critical( "The api request resulted in an redirect (302). " @@ -595,10 +587,13 @@ def _request(request_config: RequestBuilder) -> httpx.Response: ) sys.exit(1) - logger.debug("JSON response: %s", response.json()) + try: + logger.debug("JSON response: %s", response.json()) + except httpx.ResponseNotRead: + logger.debug("Response: %s", response.read()) logger.debug("Response Status Code: %d", response.status_code) - if response.status_code not in request_config.success_codes: + if response.status_code not in success_codes: with suppress(Exception): if response.json()["errcode"] == "M_UNKNOWN_TOKEN": logger.critical( @@ -610,6 +605,38 @@ def _request(request_config: RequestBuilder) -> httpx.Response: sys.exit(1) raise InternalResponseError(payload=response) + +def _request(request_config: RequestBuilder) -> httpx.Response: + """Send an synchronous request to the synapse API and receive a response. + + Attributes + ---------- + req : matrixctl.handlers.api.RequestBuilder + An instance of an RequestBuilder + + Returns + ------- + response : httpx.Response + Returns the response + + """ + + logger.debug("repr: %s", repr(request_config)) + + # There is some weird stuff going on in httpx. It is set to None by default + with httpx.Client(http2=True, timeout=request_config.timeout) as client: + response: httpx.Response = client.request( + method=request_config.method, + data=request_config.data, # type: ignore # noqa: PGH003 + json=request_config.json, + content=request_config.content, # type: ignore # noqa: PGH003 + url=str(request_config), + params=request_config.params, + headers=request_config.headers_with_auth, + follow_redirects=False, + ) + handle_sync_response_status_code(response, request_config.success_codes) + return response @@ -720,6 +747,7 @@ def streamed_download( path_download_file: Path = Path(download_file.name) logger.debug("Temporary file: %s", path_download_file) extension: str | None = None + content_length: int | None = None try: with httpx.stream( method=request_config.method, @@ -732,7 +760,18 @@ def streamed_download( timeout=request_config.timeout, follow_redirects=False, ) as response: - total: int = int(response.headers["Content-Length"]) + handle_sync_response_status_code( + response, + request_config.success_codes, + ) + try: + content_length = int(response.headers["Content-Length"]) + logger.debug("Content-Length: %s", content_length) + except KeyError as err: + logger.warning( + "Response did not include Content-Length. Error: %s", + err, + ) try: content_type: str = response.headers["Content-Type"] if content_type: @@ -746,22 +785,26 @@ def streamed_download( "Response did not include Content-Type. Error: %s", err, ) - with rich.progress.Progress( - "[progress.percentage]{task.percentage:>3.0f}%", - rich.progress.BarColumn(bar_width=None), - rich.progress.DownloadColumn(), - rich.progress.TransferSpeedColumn(), - ) as progress: - download_task: rich.progress.TaskID = progress.add_task( - "Download", - total=total, - ) + if content_length is None: for chunk in response.iter_bytes(): download_file.write(chunk) - progress.update( - download_task, - completed=response.num_bytes_downloaded, + else: + with rich.progress.Progress( + "[progress.percentage]{task.percentage:>3.0f}%", + rich.progress.BarColumn(bar_width=None), + rich.progress.DownloadColumn(), + rich.progress.TransferSpeedColumn(), + ) as progress: + download_task: rich.progress.TaskID = progress.add_task( + "Download", + total=content_length, ) + for chunk in response.iter_bytes(): + download_file.write(chunk) + progress.update( + download_task, + completed=response.num_bytes_downloaded, + ) except Exception as err: # skipcq: PYL-W0703 raise InternalResponseError(payload=response) from err diff --git a/news/835.bugfix.rst b/news/835.bugfix.rst new file mode 100644 index 00000000..3bf7680f --- /dev/null +++ b/news/835.bugfix.rst @@ -0,0 +1 @@ +Make `streamed_download` handle response status codes.