From 39f2489ed8bc14cbd24e7f94ee89a7e4e54a0eb4 Mon Sep 17 00:00:00 2001 From: Margubur Rahman Date: Wed, 12 Nov 2025 19:07:31 +0000 Subject: [PATCH] trying async download --- .../_experimental/asyncio/json/__init__.py | 0 .../_experimental/asyncio/json/_helpers.py | 128 ++++++++ .../asyncio/json/async_client.py | 285 ++++++++++++++++++ .../_experimental/asyncio/json/download.py | 90 ++++++ 4 files changed, 503 insertions(+) create mode 100644 google/cloud/storage/_experimental/asyncio/json/__init__.py create mode 100644 google/cloud/storage/_experimental/asyncio/json/_helpers.py create mode 100644 google/cloud/storage/_experimental/asyncio/json/async_client.py create mode 100644 google/cloud/storage/_experimental/asyncio/json/download.py diff --git a/google/cloud/storage/_experimental/asyncio/json/__init__.py b/google/cloud/storage/_experimental/asyncio/json/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/google/cloud/storage/_experimental/asyncio/json/_helpers.py b/google/cloud/storage/_experimental/asyncio/json/_helpers.py new file mode 100644 index 000000000..549a2e290 --- /dev/null +++ b/google/cloud/storage/_experimental/asyncio/json/_helpers.py @@ -0,0 +1,128 @@ +"""Async classes for holding the credentials, and connection""" + + +import os +import google.api_core + +from google.cloud.client import _ClientProjectMixin +from google.cloud.client import _CREDENTIALS_REFRESH_TIMEOUT +from google.cloud.client import _GOOGLE_AUTH_CREDENTIALS_HELP +from google.auth.transport import _aiohttp_requests as async_requests +from google.cloud.storage import retry as storage_retry +from pickle import PicklingError +from google.auth import _default_async +from google.cloud.storage import _http +from google.api_core import retry_async + + +DEFAULT_ASYNC_RETRY = retry_async.AsyncRetry(predicate=storage_retry._should_retry) + +# While the parent class is synchronous, we know that downloads does not use any synchronous network calls +# originated from this class. We still require some helper methods from the parent, so we are keeping the class +# structure as is, with minimal modifications. +class AsyncConnection(_http.Connection): + def get_api_base_url_for_mtls(self, api_base_url=None): + if api_base_url: + return api_base_url + + env = os.getenv("GOOGLE_API_USE_MTLS_ENDPOINT", "auto") + if env == "always": + url_to_use = self.API_BASE_MTLS_URL + elif env == "never": + url_to_use = self.API_BASE_URL + else: + if self.ALLOW_AUTO_SWITCH_TO_MTLS_URL: + url_to_use = ( + self.API_BASE_MTLS_URL if self._client._async_http._is_mtls else self.API_BASE_URL + ) + else: + url_to_use = self.API_BASE_URL + return url_to_use + + +class Client: + SCOPE = None + + def __init__(self, credentials=None, _http=None, client_options=None): + if isinstance(client_options, dict): + client_options = google.api_core.client_options.from_dict(client_options) + if client_options is None: + client_options = google.api_core.client_options.ClientOptions() + + if credentials and client_options.credentials_file: + raise google.api_core.exceptions.DuplicateCredentialArgs( + "'credentials' and 'client_options.credentials_file' are mutually exclusive." + ) + + if client_options.api_key: + raise ValueError("'client_options.api_key' currently not supported in async version.") + + if credentials and not isinstance( + credentials, google.auth._credentials_async.Credentials + ): + raise ValueError(_GOOGLE_AUTH_CREDENTIALS_HELP) + + scopes = client_options.scopes or self.SCOPE + + if not _http and credentials is None: + if client_options.credentials_file: + async_creds, _ = _default_async.load_credentials_from_file( + client_options.credentials_file, scopes=scopes + ) + else: + # This particular google.auth._default_async.default_async falls back to sync credentials in case of + # ComputeEngineCredentials, for experimentation - one need to do `gcloud auth application-default login` + # to ensure that credential do not fall back to Compute Engine. + async_creds, _ = _default_async.default_async(scopes=scopes) + else: + async_creds = credentials + + self._async_credentials = google.auth._credentials_async.with_scopes_if_required( + async_creds, scopes=scopes + ) + + if client_options.quota_project_id: + self._async_credentials = self._async_credentials.with_quota_project( + client_options.quota_project_id + ) + + self._async_http_internal = _http + self._client_cert_source = client_options.client_cert_source + + def __getstate__(self): + raise PicklingError( + "\n".join( + [ + "Pickling client objects is explicitly not supported.", + "Clients have non-trivial state that is local and unpickleable.", + ] + ) + ) + + @property + def _async_http(self): + if self._async_http_internal is None: + self._async_http_internal = async_requests.AuthorizedSession( + self._async_credentials, + refresh_timeout=_CREDENTIALS_REFRESH_TIMEOUT, + ) + # self._async_http_internal.configure_mtls_channel(self._client_cert_source) + return self._async_http_internal + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + del exc_type, exc_val, exc_tb + if self._async_http_internal is not None: + await self._async_http_internal.close() + + +class ClientWithProjectAsync(Client, _ClientProjectMixin): + _SET_PROJECT = True + + def __init__(self, project=None, credentials=None, client_options=None, _http=None): + _ClientProjectMixin.__init__(self, project=project, credentials=credentials) + Client.__init__( + self, credentials=credentials, client_options=client_options, _http=_http + ) \ No newline at end of file diff --git a/google/cloud/storage/_experimental/asyncio/json/async_client.py b/google/cloud/storage/_experimental/asyncio/json/async_client.py new file mode 100644 index 000000000..70d6b274e --- /dev/null +++ b/google/cloud/storage/_experimental/asyncio/json/async_client.py @@ -0,0 +1,285 @@ +"""Async client for SDK downloads""" + +import os +import asyncio +import google.api_core.client_options + + +from google.auth._credentials_async import AnonymousCredentials +from google.cloud.storage._experimental.asyncio.json import _helpers +from google.cloud.storage._experimental.asyncio.json import download +from google.cloud.storage._helpers import _get_environ_project +from google.cloud.storage._helpers import _get_storage_emulator_override +from google.cloud.storage._helpers import _get_api_endpoint_override +from google.cloud.storage._helpers import _use_client_cert +from google.cloud.storage._helpers import _DEFAULT_SCHEME +from google.cloud.storage._helpers import _STORAGE_HOST_TEMPLATE +from google.cloud.storage._helpers import _DEFAULT_UNIVERSE_DOMAIN +from google.cloud.storage import blob + + + +_marker = object() + +class AsyncClient(_helpers.ClientWithProjectAsync): + + SCOPE = ( + "https://www.googleapis.com/auth/devstorage.full_control", + "https://www.googleapis.com/auth/devstorage.read_only", + "https://www.googleapis.com/auth/devstorage.read_write", + ) + + def __init__( + self, + project=_marker, + credentials=None, + _http=None, + client_info=None, + client_options=None, + use_auth_w_custom_endpoint=True, + extra_headers={}, + *, + api_key=None, + ): + self._base_connection = None + + if project is None: + no_project = True + project = "" + else: + no_project = False + + if project is _marker: + project = None + + self._initial_client_info = client_info + self._initial_client_options = client_options + self._extra_headers = extra_headers + + connection_kw_args = {"client_info": client_info} + + if api_key: + if client_options and not isinstance(client_options, dict): + client_options.api_key = api_key + else: + if not client_options: + client_options = {} + client_options["api_key"] = api_key + + if client_options: + if isinstance(client_options, dict): + client_options = google.api_core.client_options.from_dict( + client_options + ) + + if client_options and client_options.universe_domain: + self._universe_domain = client_options.universe_domain + else: + self._universe_domain = None + + storage_emulator_override = _get_storage_emulator_override() + api_endpoint_override = _get_api_endpoint_override() + + if client_options and client_options.api_endpoint: + api_endpoint = client_options.api_endpoint + + elif storage_emulator_override: + api_endpoint = storage_emulator_override + + elif api_endpoint_override: + api_endpoint = api_endpoint_override + + elif self._universe_domain: + + if _use_client_cert(): + raise ValueError( + 'The "GOOGLE_API_USE_CLIENT_CERTIFICATE" env variable is ' + 'set to "true" and a non-default universe domain is ' + "configured. mTLS is not supported in any universe other than" + "googleapis.com." + ) + api_endpoint = _DEFAULT_SCHEME + _STORAGE_HOST_TEMPLATE.format( + universe_domain=self._universe_domain + ) + + else: + api_endpoint = None + + connection_kw_args["api_endpoint"] = api_endpoint + + self._is_emulator_set = True if storage_emulator_override else False + + if connection_kw_args["api_endpoint"] is not None: + if self._is_emulator_set or not use_auth_w_custom_endpoint: + if credentials is None: + credentials = AnonymousCredentials() + if project is None: + project = _get_environ_project() + if project is None: + no_project = True + project = "" + + super(AsyncClient, self).__init__( + project=project, + credentials=credentials, + client_options=client_options, + _http=_http, + ) + + + if self._async_credentials.universe_domain != self.universe_domain: + raise ValueError( + "The configured universe domain ({client_ud}) does not match " + "the universe domain found in the credentials ({cred_ud}). If " + "you haven't configured the universe domain explicitly, " + "`googleapis.com` is the default.".format( + client_ud=self.universe_domain, + cred_ud=self._async_credentials.universe_domain, + ) + ) + + if no_project: + self.project = None + + connection = _helpers.AsyncConnection(self, **connection_kw_args) + connection.extra_headers = extra_headers + self._connection = connection + + @property + def universe_domain(self): + return self._universe_domain or _DEFAULT_UNIVERSE_DOMAIN + + async def _perform_download( + self, + transport, + file_obj, + download_url, + headers, + start=None, + end=None, + timeout=None, + checksum="md5", + retry=_helpers.DEFAULT_ASYNC_RETRY, + sequential_read=False, + ): + extra_attributes = blob._get_opentelemetry_attributes_from_url(download_url) + extra_attributes["upload.checksum"] = f"{checksum}" + + download_obj = download.DownloadAsync( + download_url, + stream=file_obj, + headers=headers, + start=start, + end=end, + checksum=checksum, + retry=retry, + sequential_read=sequential_read, + ) + await download_obj.consume(transport, timeout=timeout) + + def _check_if_sliced_download_is_eligible(self, obj_size, checksum): + if obj_size < 1024*1024*1024: + return False + # Need to support checksum validations for parallel downloads. + return checksum==None + + async def download_to_file( + self, + blob_obj, + filename, + start=None, + end=None, + if_etag_match=None, + if_etag_not_match=None, + if_generation_match=None, + if_generation_not_match=None, + if_metageneration_match=None, + if_metageneration_not_match=None, + timeout=None, + checksum="md5", + retry=_helpers.DEFAULT_ASYNC_RETRY, + sequential_read=False, + ): + download_url = blob_obj._get_download_url( + self, + if_generation_match=if_generation_match, + if_generation_not_match=if_generation_not_match, + if_metageneration_match=if_metageneration_match, + if_metageneration_not_match=if_metageneration_not_match, + ) + headers = blob._get_encryption_headers(blob_obj._encryption_key) + headers["accept-encoding"] = "gzip" + blob._add_etag_match_headers( + headers, + if_etag_match=if_etag_match, + if_etag_not_match=if_etag_not_match, + ) + headers = { + **blob._get_default_headers(self._connection.user_agent, command=None), + **headers, + **self._extra_headers, + } + + transport = self._async_http + if not blob_obj.size: + blob_obj.reload() + obj_size = blob_obj.size + try: + if not sequential_read and self._check_if_sliced_download_is_eligible(obj_size, checksum): # 1GB + print("Sliced Download Preferred, and Starting...") + _parts = 5 + chunks_offset = [0] + [obj_size//_parts]*(_parts-1) + [obj_size - obj_size//_parts*(_parts-1)] + for i in range(1, _parts+1): + chunks_offset[i]+=chunks_offset[i-1] + + with open(filename, 'wb') as _: pass # trunacates the file to zero, and keeps the file. + + tasks, file_handles = [], [] + try: + for idx in range(_parts): + file_handle = open(filename, 'r+b') + file_handle.seek(chunks_offset[idx]) + tasks.append( + self._perform_download( + transport, + file_handle, + download_url, + headers, + chunks_offset[idx], + chunks_offset[idx+1]-1, + timeout=timeout, + checksum=checksum, + retry=retry, + sequential_read=sequential_read, + ) + ) + file_handles.append(file_handle) + await asyncio.gather(*tasks) + finally: + for file_handle in file_handles: + file_handle.close() + else: + print("Sequential Download Preferred, and Starting...") + with open(filename, "wb") as file_obj: + await self._perform_download( + transport, + file_obj, + download_url, + headers, + start, + end, + timeout=timeout, + checksum=checksum, + retry=retry, + sequential_read=sequential_read, + ) + except (blob.DataCorruption, blob.NotFound): + os.remove(filename) + raise + except blob.InvalidResponse as exc: + blob._raise_from_invalid_response(exc) + + updated = blob_obj.updated + if updated is not None: + mtime = updated.timestamp() + os.utime(filename, (mtime, mtime)) \ No newline at end of file diff --git a/google/cloud/storage/_experimental/asyncio/json/download.py b/google/cloud/storage/_experimental/asyncio/json/download.py new file mode 100644 index 000000000..67ad0185b --- /dev/null +++ b/google/cloud/storage/_experimental/asyncio/json/download.py @@ -0,0 +1,90 @@ +"""Async based download code""" + +import http +import aiohttp + +from google.cloud.storage._experimental.asyncio.json._helpers import DEFAULT_ASYNC_RETRY +from google.cloud.storage._media.requests import _request_helpers +from google.cloud.storage._media import _download +from google.cloud.storage._media import _helpers +from google.cloud.storage._media.requests import download as storage_download + + +class DownloadAsync(_request_helpers.RequestsMixin, _download.Download): + + def __init__( + self, + media_url, + stream=None, + start=None, + end=None, + headers=None, + checksum="md5", + retry=DEFAULT_ASYNC_RETRY, + sequential_read=False, + ): + super(DownloadAsync, self).__init__( + media_url, stream=stream, start=start, end=end, headers=headers, checksum=checksum, retry=retry + ) + self.sequential_read = sequential_read + + @staticmethod + def _get_status_code(response): + return response.status + + async def _write_to_stream(self, response): + if self._expected_checksum or self._checksum_object: + # Presevre it across calls. + expected_checksum = self._expected_checksum + checksum_object = self._checksum_object + if not self.sequential_read: + # If we've not set expected checksum, or checksum object yet, and if it is not + # sequential download, API would not return us hash value for each chunk. + # We could ideally compute the crc32c checksum for each chunk, and later combine them + # and check, However for prototype not implementing it. + expected_checksum = None + checksum_object = _helpers._DoNothingHash() + self._expected_checksum = expected_checksum + self._checksum_object = checksum_object + else: + # Sequential read, so fetch the hash from the headers. + expected_checksum, checksum_object = _helpers._get_expected_checksum( + response, self._get_headers, self.media_url, checksum_type=self.checksum + ) + self._expected_checksum = expected_checksum + self._checksum_object = checksum_object + + async with response: + chunk_size = 4096 * 32 + async for chunk in response.content.iter_chunked(chunk_size): + self._stream.write(chunk) # for some reason, aiofiles shows worse performance. + self._bytes_downloaded += len(chunk) + checksum_object.update(chunk) + + if ( + expected_checksum is not None + and response.status != http.client.PARTIAL_CONTENT + ): + actual_checksum = _helpers.prepare_checksum_digest(checksum_object.digest()) + + if actual_checksum != expected_checksum: + raise storage_download.DataCorruption('Corrupted download!') + + async def consume( + self, + transport, + timeout=aiohttp.ClientTimeout(total=None, sock_read=300), + ): + method, _, payload, headers = self._prepare_request() + request_kwargs = { + "data": payload, + "headers": headers, + "timeout": timeout, + } + async def retriable_request(): + url = self.media_url + result = await transport.request(method, url, **request_kwargs) + await self._write_to_stream(result) + return result + + return await _request_helpers.wait_and_retry(retriable_request, self._retry_strategy)