Stream the final file to the cache

Author: notatallshaw

src/pip/_internal/network/cache.py

@@ -3,10 +3,11 @@
 from __future__ import annotations
 
 import os
+import shutil
 from collections.abc import Generator
 from contextlib import contextmanager
 from datetime import datetime
-from typing import BinaryIO
+from typing import Any, BinaryIO, Callable
 
 from pip._vendor.cachecontrol.cache import SeparateBodyBaseCache
 from pip._vendor.cachecontrol.caches import SeparateBodyFileCache
@@ -72,12 +73,13 @@ def get(self, key: str) -> bytes | None:
             with open(metadata_path, "rb") as f:
                 return f.read()
 
-    def _write(self, path: str, data: bytes) -> None:
+    def _write_to_file(self, path: str, writer_func: Callable[[BinaryIO], Any]) -> None:
+        """Common file writing logic with proper permissions and atomic replacement."""
         with suppressed_cache_errors():
             ensure_dir(os.path.dirname(path))
 
             with adjacent_tmp_file(path) as f:
-                f.write(data)
+                writer_func(f)
                 # Inherit the read/write permissions of the cache directory
                 # to enable multi-user cache use-cases.
                 mode = (
@@ -93,6 +95,12 @@ def _write(self, path: str, data: bytes) -> None:
 
             replace(f.name, path)
 
+    def _write(self, path: str, data: bytes) -> None:
+        self._write_to_file(path, lambda f: f.write(data))
+
+    def _write_from_io(self, path: str, source_file: BinaryIO) -> None:
+        self._write_to_file(path, lambda f: shutil.copyfileobj(source_file, f))
+
     def set(
         self, key: str, value: bytes, expires: int | datetime | None = None
     ) -> None:
@@ -118,3 +126,8 @@ def get_body(self, key: str) -> BinaryIO | None:
     def set_body(self, key: str, body: bytes) -> None:
         path = self._get_cache_path(key) + ".body"
         self._write(path, body)
+
+    def set_body_from_io(self, key: str, body_file: BinaryIO) -> None:
+        """Set the body of the cache entry from a file object."""
+        path = self._get_cache_path(key) + ".body"
+        self._write_from_io(path, body_file)

src/pip/_internal/network/download.py

@@ -21,7 +21,7 @@
 from pip._internal.exceptions import IncompleteDownloadError, NetworkConnectionError
 from pip._internal.models.index import PyPI
 from pip._internal.models.link import Link
-from pip._internal.network.cache import is_from_cache
+from pip._internal.network.cache import SafeFileCache, is_from_cache
 from pip._internal.network.session import CacheControlAdapter, PipSession
 from pip._internal.network.utils import HEADERS, raise_for_status, response_chunks
 from pip._internal.utils.misc import format_size, redact_auth_from_url, splitext
@@ -284,6 +284,14 @@ def _cache_resumed_download(
             )
             return
 
+        # Check SafeFileCache is being used
+        if not isinstance(adapter.cache, SafeFileCache):
+            logger.debug(
+                "Skipping resume download caching: "
+                "cache doesn't support separate body storage"
+            )
+            return
+
         synthetic_request = PreparedRequest()
         synthetic_request.prepare(method="GET", url=url, headers={})
 
@@ -300,15 +308,17 @@ def _cache_resumed_download(
             preload_content=False,
         )
 
-        # Use the cache controller to store this as a complete response
+        # Stream the file to cache
+        cache_url = adapter.controller.cache_url(url)
+        adapter.cache.set(
+            cache_url,
+            adapter.controller.serializer.dumps(
+                synthetic_request, synthetic_response, b""
+            ),
+        )
         download.output_file.flush()
         with open(download.output_file.name, "rb") as f:
-            adapter.controller.cache_response(
-                synthetic_request,
-                synthetic_response,
-                body=f.read(),
-                status_codes=(200, 203, 300, 301, 308),
-            )
+            adapter.cache.set_body_from_io(cache_url, f)
 
         logger.debug(
             "Cached resumed download as complete response for future use: %s", url

tests/unit/test_network_download.py

@@ -9,6 +9,7 @@
 
 from pip._internal.exceptions import IncompleteDownloadError
 from pip._internal.models.link import Link
+from pip._internal.network.cache import SafeFileCache
 from pip._internal.network.download import (
     Downloader,
     _get_http_response_size,
@@ -374,6 +375,13 @@ def test_resumed_download_caching(tmpdir: Path) -> None:
     mock_adapter = MagicMock(spec=CacheControlAdapter)
     mock_controller = MagicMock()
     mock_adapter.controller = mock_controller
+    mock_controller.cache_url = MagicMock(return_value="cache_key")
+    mock_controller.serializer = MagicMock()
+    mock_controller.serializer.dumps = MagicMock(return_value=b"serialized_data")
+
+    # Mock the cache to be a SafeFileCache
+    mock_cache = MagicMock(spec=SafeFileCache)
+    mock_adapter.cache = mock_cache
 
     # Create a mock for the session adapters
     adapters_mock = MagicMock()
@@ -392,20 +400,18 @@ def test_resumed_download_caching(tmpdir: Path) -> None:
             expected_bytes = b"0cfa7e9d-1868-4dd7-9fb3-f2561d5dfd89"
             assert downloaded_bytes == expected_bytes
 
-        # Verify that cache_response was called for the resumed download
-        mock_controller.cache_response.assert_called_once()
+        # Verify that cache.set was called for metadata
+        mock_cache.set.assert_called_once()
+
+        # Verify that set_body_from_io was called for streaming the body
+        mock_cache.set_body_from_io.assert_called_once()
 
-        # Get the call arguments to verify the cached content
-        call_args = mock_controller.cache_response.call_args
-        assert call_args is not None
+        # Verify the call arguments
+        set_call_args = mock_cache.set.call_args
+        assert set_call_args[0][0] == "cache_key"  # First argument should be cache_key
 
-        # Extract positional and keyword arguments
-        args, kwargs = call_args
-        request, response = args
-        body = kwargs.get("body")
-        status_codes = kwargs.get("status_codes")
+        set_body_call_args = mock_cache.set_body_from_io.call_args
 
-        assert body == expected_bytes, "Cached body should match complete file content"
-        assert response.status == 200, "Cached response should have status 200"
-        assert request.url == link.url_without_fragment
-        assert 200 in status_codes
+        assert set_body_call_args[0][0] == "cache_key"
+        assert hasattr(set_body_call_args[0][1], "read")
+        assert set_body_call_args[0][1].name == filepath