python(feat): let download_file forward extra headers to rest_client.get

python/lib/sift_client/_internal/util/file.py

@@ -94,6 +94,7 @@ def download_file(
     *,
     rest_client: RestClient,
     show_progress: bool = False,
+    extra_headers: dict[str, str] | None = None,
 ) -> Path:
     """Download a file from a URL in streaming 4 MiB chunks.
 
@@ -103,6 +104,13 @@ def download_file(
         rest_client: The SDK rest client to use for the download.
         show_progress: If True, display a progress bar during download.
             Defaults to False.
+        extra_headers: Optional additional headers to include on the request.
+            Merged on top of the internal ``Authorization: None`` strip, so a
+            caller-supplied ``Authorization`` (if any) wins. Useful for
+            presigned URLs whose signature covers a specific ``Host`` header
+            that differs from the URL's netloc (e.g. an S3 endpoint accessed
+            through a container-internal alias while the URL was signed
+            against the browser-facing loopback authority).
 
     Returns:
         The path to the downloaded file.
@@ -111,8 +119,13 @@ def download_file(
         requests.HTTPError: If the download request fails.
     """
     output_path.parent.mkdir(parents=True, exist_ok=True)
-    # Strip the session's default Authorization header, presigned URLs carry their own auth
-    with rest_client.get(signed_url, stream=True, headers={"Authorization": None}) as response:
+    # Strip the session's default Authorization header (presigned URLs carry
+    # their own auth). Any ``extra_headers`` merge on top so callers can
+    # override the Host header, add a range, etc.
+    headers: dict[str, str | None] = {"Authorization": None}
+    if extra_headers:
+        headers.update(extra_headers)
+    with rest_client.get(signed_url, stream=True, headers=headers) as response:
         response.raise_for_status()
         total_bytes = int(response.headers.get("Content-Length", 0)) or None
         with alive_bar(

python/lib/sift_client/_tests/_internal/util/test_file.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+from sift_client._internal.util.file import download_file
+
+
+class _FakeResponse:
+    """Minimal ``requests.Response`` stand-in for the streamed-download tests."""
+
+    def __init__(self, chunks: list[bytes]):
+        self._chunks = chunks
+        self.headers: dict[str, str] = {}
+
+    def raise_for_status(self) -> None: ...
+
+    def iter_content(self, chunk_size: int):
+        return iter(self._chunks)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        return False
+
+
+def _install_rest_client_stub(response: _FakeResponse) -> tuple[MagicMock, dict]:
+    """Return a rest_client mock whose ``.get(...)`` records call kwargs."""
+    captured: dict = {}
+    rest_client = MagicMock()
+
+    def _fake_get(url, **kwargs):
+        captured["url"] = url
+        captured.update(kwargs)
+        return response
+
+    rest_client.get.side_effect = _fake_get
+    return rest_client, captured
+
+
+class TestDownloadFileExtraHeaders:
+    """``extra_headers`` merges on top of the internal ``Authorization: None``
+    strip so callers can override the Host header (or add a range, etc.) for
+    presigned URLs whose signature covers a specific ``Host`` distinct from the
+    URL's authority. Without this the caller has to bypass ``download_file``
+    entirely for signed-URL edge cases (e.g. local minio via a container-
+    internal alias) and reimplement streaming + auth stripping from scratch.
+    """
+
+    def test_no_extra_headers_strips_only_authorization(self, tmp_path):
+        rest_client, captured = _install_rest_client_stub(_FakeResponse([b"hello ", b"world"]))
+        target = tmp_path / "out.bin"
+
+        download_file(
+            "https://example.com/file?sig=1",
+            target,
+            rest_client=rest_client,
+        )
+
+        assert captured["headers"] == {"Authorization": None}
+        assert target.read_bytes() == b"hello world"
+
+    def test_extra_headers_are_forwarded(self, tmp_path):
+        rest_client, captured = _install_rest_client_stub(_FakeResponse([b"payload"]))
+        target = tmp_path / "out.bin"
+
+        download_file(
+            "https://example.com/file?sig=1",
+            target,
+            rest_client=rest_client,
+            extra_headers={"Host": "signed.example.com", "X-Trace": "abc"},
+        )
+
+        assert captured["headers"] == {
+            "Authorization": None,
+            "Host": "signed.example.com",
+            "X-Trace": "abc",
+        }
+        assert target.read_bytes() == b"payload"
+
+    def test_extra_headers_can_override_authorization_strip(self, tmp_path):
+        # The strip is a default, not a hard invariant -- if a caller has a
+        # concrete reason to pass an Authorization on the download request
+        # (e.g. a proxy-scoped bearer), let their value win.
+        rest_client, captured = _install_rest_client_stub(_FakeResponse([b"payload"]))
+        target = tmp_path / "out.bin"
+
+        download_file(
+            "https://example.com/file?sig=1",
+            target,
+            rest_client=rest_client,
+            extra_headers={"Authorization": "Bearer proxy-token"},
+        )
+
+        assert captured["headers"]["Authorization"] == "Bearer proxy-token"