fix(storage): catch bad responses from server (#1344)

Author: o-santi

src/storage/src/storage3/_async/file_api.py

@@ -77,10 +77,14 @@ async def _request(
             )
             response.raise_for_status()
         except HTTPStatusError as exc:
-            resp = exc.response.json()
-            raise StorageApiError(
-                resp["message"], resp["error"], resp["statusCode"]
-            ) from exc
+            try:
+                resp = exc.response.json()
+                raise StorageApiError(
+                    resp["message"], resp["error"], resp["statusCode"]
+                ) from exc
+            except KeyError as err:
+                message = f"Unable to parse error message: {resp.text}"
+                raise StorageApiError(message, "InternalError", 400) from err
 
         # close the resource before returning the response
         if files and "file" in files and isinstance(files["file"][1], BufferedReader):

src/storage/src/storage3/_sync/file_api.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass, field
 from io import BufferedReader, FileIO
 from pathlib import Path
-from typing import Any, List, Literal, Optional, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 from httpx import Client, Headers, HTTPStatusError, Response
 from yarl import URL
@@ -77,10 +77,14 @@ def _request(
             )
             response.raise_for_status()
         except HTTPStatusError as exc:
-            resp = exc.response.json()
-            raise StorageApiError(
-                resp["message"], resp["error"], resp["statusCode"]
-            ) from exc
+            try:
+                resp = exc.response.json()
+                raise StorageApiError(
+                    resp["message"], resp["error"], resp["statusCode"]
+                ) from exc
+            except KeyError as err:
+                message = f"Unable to parse error message: {resp.text}"
+                raise StorageApiError(message, "InternalError", 400) from err
 
         # close the resource before returning the response
         if files and "file" in files and isinstance(files["file"][1], BufferedReader):
@@ -436,7 +440,12 @@ def list(
         )
         return response.json()
 
-    def download(self, path: str, options: Optional[DownloadOptions] = None) -> bytes:
+    def download(
+        self,
+        path: str,
+        options: Optional[DownloadOptions] = None,
+        query_params: Optional[Dict[str, str]] = None,
+    ) -> bytes:
         """
         Downloads a file.
 
@@ -445,20 +454,23 @@ def download(self, path: str, options: Optional[DownloadOptions] = None) -> byte
         path
             The file path to be downloaded, including the path and file name. For example `folder/image.png`.
         """
-        url_options = options or {}
+        url_options = options or DownloadOptions()
         render_path = (
             ["render", "image", "authenticated"]
             if url_options.get("transform")
             else ["object"]
         )
 
-        transform_options = url_options.get("transform") or {}
+        transform_options = url_options.get("transform") or TransformOptions()
 
         path_parts = relative_path_to_parts(path)
         response = self._request(
             "GET",
             [*render_path, self.id, *path_parts],
-            query_params=transform_to_dict(transform_options),
+            query_params={
+                **transform_to_dict(transform_options),
+                **(query_params or {}),
+            },
         )
         return response.content
 

src/storage/tests/_sync/test_client.py

@@ -283,6 +283,45 @@ def test_client_upload(
     assert image_info.get("metadata", {}).get("mimetype") == file.mime_type
 
 
+def test_client_upload_with_query(
+    storage_file_client: SyncBucketProxy, file: FileForTesting
+) -> None:
+    """Ensure we can upload files to a bucket, even with query parameters"""
+    storage_file_client.upload(
+        file.bucket_path, file.local_path, {"content-type": file.mime_type}
+    )
+
+    image = storage_file_client.download(
+        file.bucket_path, query_params={"my-param": "test"}
+    )
+    files = storage_file_client.list(file.bucket_folder)
+    image_info = next((f for f in files if f.get("name") == file.name), None)
+
+    assert image == file.file_content
+    assert image_info is not None
+    assert image_info.get("metadata", {}).get("mimetype") == file.mime_type
+
+
+def test_client_download_with_query_doesnt_lose_params(
+    storage_file_client: SyncBucketProxy, file: FileForTesting
+) -> None:
+    """Ensure query params aren't lost"""
+    from yarl import URL
+
+    params = {"my-param": "test"}
+    mock_response = Mock()
+    with patch.object(HttpxClient, "request") as mock_request:
+        mock_request.return_value = mock_response
+        storage_file_client.download(file.bucket_path, query_params=params)
+        expected_url = storage_file_client._base_url.joinpath(
+            "object", storage_file_client.id, *URL(file.bucket_path).parts
+        ).with_query(params)
+        actual_url = mock_request.call_args[0][1]
+
+        assert URL(actual_url).query == params
+        assert str(expected_url) == actual_url
+
+
 def test_client_update(
     storage_file_client: SyncBucketProxy,
     two_files: list[FileForTesting],