fix(storage): add query parameters option to download (#1327)

Author: o-santi

src/storage/src/storage3/_async/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 AsyncClient, Headers, HTTPStatusError, Response
 from yarl import URL
@@ -439,7 +439,10 @@ async def list(
         return response.json()
 
     async def download(
-        self, path: str, options: Optional[DownloadOptions] = None
+        self,
+        path: str,
+        options: Optional[DownloadOptions] = None,
+        query_params: Optional[Dict[str, str]] = None,
     ) -> bytes:
         """
         Downloads a file.
@@ -449,20 +452,23 @@ async def download(
         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 = await 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/_async/test_client.py

@@ -283,6 +283,45 @@ async def test_client_upload(
     assert image_info.get("metadata", {}).get("mimetype") == file.mime_type
 
 
+async def test_client_upload_with_query(
+    storage_file_client: AsyncBucketProxy, file: FileForTesting
+) -> None:
+    """Ensure we can upload files to a bucket, even with query parameters"""
+    await storage_file_client.upload(
+        file.bucket_path, file.local_path, {"content-type": file.mime_type}
+    )
+
+    image = await storage_file_client.download(
+        file.bucket_path, query_params={"my-param": "test"}
+    )
+    files = await 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
+
+
+async def test_client_download_with_query_doesnt_lose_params(
+    storage_file_client: AsyncBucketProxy, 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
+        await 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
+
+
 async def test_client_update(
     storage_file_client: AsyncBucketProxy,
     two_files: list[FileForTesting],