Cache final file from resume retry process

Author: notatallshaw

src/pip/_internal/network/download.py

@@ -11,15 +11,18 @@
 from http import HTTPStatus
 from typing import BinaryIO
 
+from pip._vendor.requests import PreparedRequest
 from pip._vendor.requests.models import Response
+from pip._vendor.urllib3 import HTTPResponse as URLlib3Response
+from pip._vendor.urllib3._collections import HTTPHeaderDict
 from pip._vendor.urllib3.exceptions import ReadTimeoutError
 
 from pip._internal.cli.progress_bars import get_download_progress_renderer
 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.session import PipSession
+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
 
@@ -250,6 +253,67 @@ def _attempt_resumes_or_redownloads(
             os.remove(download.output_file.name)
             raise IncompleteDownloadError(download)
 
+        # If we successfully completed the download via resume, manually cache it
+        # as a complete response to enable future caching
+        if download.reattempts > 0:
+            self._cache_resumed_download(download, first_resp)
+
+    def _cache_resumed_download(
+        self, download: _FileDownload, original_response: Response
+    ) -> None:
+        """
+        Manually cache a file that was successfully downloaded via resume retries.
+
+        cachecontrol doesn't cache 206 (Partial Content) responses, since they
+        are not complete files. This method manually adds the final file to the
+        cache as though it was downloaded in a single request, so that future
+        requests can use the cache.
+        """
+        url = download.link.url_without_fragment
+        if url.startswith("https://"):
+            adapter = self._session.adapters["https://"]
+        elif url.startswith("http://"):
+            adapter = self._session.adapters["http://"]
+        else:
+            return
+
+        # Check if the adapter is the CacheControlAdapter (i.e. caching is enabled)
+        if not isinstance(adapter, CacheControlAdapter):
+            logger.debug(
+                "Skipping resume download caching: no cache controller for %s", url
+            )
+            return
+
+        synthetic_request = PreparedRequest()
+        synthetic_request.prepare(method="GET", url=url, headers={})
+
+        synthetic_response_headers = HTTPHeaderDict()
+        for key, value in original_response.headers.items():
+            if key.lower() not in ["content-range", "content-length"]:
+                synthetic_response_headers[key] = value
+        synthetic_response_headers["content-length"] = str(download.size)
+
+        synthetic_response = URLlib3Response(
+            body="",
+            headers=synthetic_response_headers,
+            status=200,
+            preload_content=False,
+        )
+
+        # Use the cache controller to store this as a complete response
+        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),
+            )
+
+        logger.debug(
+            "Cached resumed download as complete response for future use: %s", url
+        )
+
     def _http_get_resume(
         self, download: _FileDownload, should_match: Response
     ) -> Response:

tests/unit/test_network_download.py

@@ -16,7 +16,7 @@
     parse_content_disposition,
     sanitize_content_filename,
 )
-from pip._internal.network.session import PipSession
+from pip._internal.network.session import CacheControlAdapter, PipSession
 from pip._internal.network.utils import HEADERS
 
 from tests.lib.requests_mocks import MockResponse
@@ -350,3 +350,62 @@ def test_downloader(
 
     # Make sure that the downloader makes additional requests for resumption
     _http_get_mock.assert_has_calls(calls)
+
+
+def test_resumed_download_caching(tmpdir: Path) -> None:
+    """Test that resumed downloads are cached properly for future use."""
+    session = PipSession()
+    link = Link("http://example.com/foo.tgz")
+    downloader = Downloader(session, "on", resume_retries=5)
+
+    # Mock an incomplete download followed by a successful resume
+    incomplete_resp = MockResponse(b"0cfa7e9d-1868-4dd7-9fb3-")
+    incomplete_resp.headers = {"content-length": "36"}
+    incomplete_resp.status_code = 200
+
+    resume_resp = MockResponse(b"f2561d5dfd89")
+    resume_resp.headers = {"content-length": "12"}
+    resume_resp.status_code = 206
+
+    responses = [incomplete_resp, resume_resp]
+    _http_get_mock = MagicMock(side_effect=responses)
+
+    # Mock the session's adapters to have a cache controller
+    mock_adapter = MagicMock(spec=CacheControlAdapter)
+    mock_controller = MagicMock()
+    mock_adapter.controller = mock_controller
+
+    # Create a mock for the session adapters
+    adapters_mock = MagicMock()
+    adapters_mock.__getitem__ = MagicMock(return_value=mock_adapter)
+
+    with (
+        patch.object(Downloader, "_http_get", _http_get_mock),
+        patch.object(session, "adapters", adapters_mock),
+    ):
+
+        filepath, _ = downloader(link, str(tmpdir))
+
+        # Verify the file was downloaded correctly
+        with open(filepath, "rb") as downloaded_file:
+            downloaded_bytes = downloaded_file.read()
+            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()
+
+        # Get the call arguments to verify the cached content
+        call_args = mock_controller.cache_response.call_args
+        assert call_args is not None
+
+        # Extract positional and keyword arguments
+        args, kwargs = call_args
+        request, response = args
+        body = kwargs.get("body")
+        status_codes = kwargs.get("status_codes")
+
+        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