cdm_data_loaders.utils.file_transfer.async_client ¶

class AsyncFileDownloader(BaseFileDownloader):
    """Asynchronous downloader interface."""

    def __init__(  # noqa: PLR0913
        self,
        client: httpx.AsyncClient | None = None,
        max_attempts: int = RETRY_DEFAULTS["max_attempts"],
        min_backoff: int = RETRY_DEFAULTS["min_backoff"],
        max_backoff: int = RETRY_DEFAULTS["max_backoff"],
        chunk_size: int = 8192,
        max_concurrency: int | None = None,
    ) -> None:
        """Initialise an async download client.

        :param client: an httpx.Client object, defaults to None
        :type client: httpx.Client | None, optional
        :param max_attempts: how many times to retry the request defaults to 5
        :type max_attempts: int, optional
        :param min_backoff: minimum backoff for retries, defaults to 1
        :type min_backoff: int, optional
        :param max_backoff: maximum backoff for retries, defaults to 30
        :type max_backoff: int, optional
        :param chunk_size: chunk size for the downloader, defaults to 8192
        :type chunk_size: int, optional
        :param max_concurrency: maximum concurrency for the downloader, defaults to None
        :type max_concurrency: int, optional
        """
        super().__init__(client or get_async_httpx_client(), chunk_size=chunk_size)
        self.semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None

        self._retry = AsyncRetrying(**DownloadCore.retry_kwargs(max_attempts, min_backoff, max_backoff))

    async def download(
        self,
        url: str,
        destination: str | Path,
        expected_checksum: str | None = None,
        checksum_fn: str | None = "sha256",
        extra_headers: dict[str, Any] | None = None,
    ) -> Path | None:
        """Download a file from ``url`` to ``destination`` on disk.

        Returns the file path if successful; None otherwise.
        304 / Not Modified responses return None.

        :param url: URL to download from
        :type url: str
        :param destination: where to save the file to
        :type destination: str | Path
        :param expected_checksum: expected checksum (if known), defaults to None
        :type expected_checksum: str | None, optional
        :param checksum_fn: function used for calculating the checksum, defaults to "sha256"
        :type checksum_fn: str, optional
        :param extra_headers: allow extra headers to be passed, defaults to None
        :type extra_headers: dict[str, Any] | None, optional
        :return: either the path to the downloaded file or None if no file was downloaded
        :rtype: Path | None
        """
        destination, checksum_fn, extra_headers = DownloadCore.validate_args(destination, checksum_fn, extra_headers)

        async with self.semaphore or contextlib.nullcontext():
            return await self._download_with_retry(
                url,
                destination,
                expected_checksum,
                checksum_fn,
                extra_headers,
            )

    async def _download_with_retry(self, *args) -> Path | None:
        async for attempt in self._retry:
            with attempt:
                return await self._download_once(*args)
        # this should never be raised under normal circumstances
        # tenacity should never exhaust silently
        msg = "Iterator exhausted, unreachable code reached!"
        raise RuntimeError(msg)

    async def _download_once(
        self,
        url: str,
        destination: Path,
        expected_checksum: str | None,
        checksum_fn: str,
        extra_headers: dict[str, Any],
    ) -> Path | None:
        """Core download function.

        :param url: URL to download from
        :type url: str
        :param destination: where to save the file to
        :type destination: Path
        :param expected_checksum: expected checksum (if known), defaults to None
        :type expected_checksum: str | None, optional
        :param checksum_fn: function used for calculating the checksum
        :type checksum_fn: str
        :param extra_headers: any extra headers to pass to the request
        :type extra_headers: dict[str, Any]
        :raises DownloadError: for unrecoverable download errors
        :raises ChecksumMismatchError: if the checksum does not match expected_checksum
        :return: either the path to the downloaded file or None if no file was downloaded
        :rtype: Path | None
        """
        try:
            async with self.client.stream("GET", url, headers=extra_headers) as response:
                if not DownloadCore.check_should_download(response, url, logger):
                    return None

                chunks = response.aiter_bytes(chunk_size=self.chunk_size)

                await DownloadCore.write_and_hash_async(
                    url,
                    destination,
                    chunks,
                    expected_checksum=expected_checksum,
                    checksum_fn=checksum_fn,
                )

        except Exception as exc:  # noqa: BLE001
            DownloadCore.handle_download_exception(exc, url, logger)

        return DownloadCore.log_download_success(url, destination, logger)