1# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2# SPDX-License-Identifier: Apache-2.0
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import codecs
17import copy
18import io
19import json
20import logging
21import os
22import tempfile
23from collections.abc import Callable, Iterator
24from typing import IO, Any, Optional, TypeVar, Union
25
26from google.api_core.exceptions import GoogleAPICallError, NotFound
27from google.auth import credentials as auth_credentials
28from google.auth import identity_pool
29from google.cloud import storage
30from google.cloud.storage import transfer_manager
31from google.cloud.storage.exceptions import InvalidResponse
32from google.oauth2 import service_account
33from google.oauth2.credentials import Credentials as OAuth2Credentials
34
35from multistorageclient_rust import RustClient, RustClientError, RustRetryableError
36
37from ..constants import DEFAULT_MAX_POOL_CONNECTIONS
38from ..rust_utils import parse_retry_config, run_async_rust_client_method
39from ..telemetry import Telemetry
40from ..types import (
41 AWARE_DATETIME_MIN,
42 Credentials,
43 CredentialsProvider,
44 NotModifiedError,
45 ObjectMetadata,
46 PreconditionFailedError,
47 Range,
48 RetryableError,
49 SymlinkHandling,
50)
51from ..utils import (
52 safe_makedirs,
53 split_path,
54 validate_attributes,
55)
56from .base import BaseStorageProvider
57
58_T = TypeVar("_T")
59
60PROVIDER = "gcs"
61
62MiB = 1024 * 1024
63
64DEFAULT_MULTIPART_THRESHOLD = 64 * MiB
65DEFAULT_MULTIPART_CHUNKSIZE = 32 * MiB
66DEFAULT_IO_CHUNKSIZE = 32 * MiB
67PYTHON_MAX_CONCURRENCY = 8
68
69logger = logging.getLogger(__name__)
70
71
[docs]
72class StringTokenSupplier(identity_pool.SubjectTokenSupplier):
73 """
74 Supply a string token to the Google Identity Pool.
75 """
76
77 def __init__(self, token: str):
78 self._token = token
79
[docs]
80 def get_subject_token(self, context, request):
81 return self._token
82
83
[docs]
84class GoogleIdentityPoolCredentialsProvider(CredentialsProvider):
85 """
86 A concrete implementation of the :py:class:`multistorageclient.types.CredentialsProvider` that provides Google's identity pool credentials.
87 """
88
89 def __init__(self, audience: str, token_supplier: str):
90 """
91 Initializes the :py:class:`GoogleIdentityPoolCredentialsProvider` with the audience and token supplier.
92
93 :param audience: The audience for the Google Identity Pool.
94 :param token_supplier: The token supplier for the Google Identity Pool.
95 """
96 self._audience = audience
97 self._token_supplier = token_supplier
98
[docs]
99 def get_credentials(self) -> Credentials:
100 return Credentials(
101 access_key="",
102 secret_key="",
103 token="",
104 expiration=None,
105 custom_fields={"audience": self._audience, "token": self._token_supplier},
106 )
107
[docs]
108 def refresh_credentials(self) -> None:
109 pass
110
111
[docs]
112class GoogleServiceAccountCredentialsProvider(CredentialsProvider):
113 """
114 A concrete implementation of the :py:class:`multistorageclient.types.CredentialsProvider` that provides Google's service account credentials.
115 """
116
117 #: Google service account private key file contents.
118 _info: dict[str, Any]
119
120 def __init__(self, file: Optional[str] = None, info: Optional[dict[str, Any]] = None):
121 """
122 Initializes the :py:class:`GoogleServiceAccountCredentialsProvider` with either a path to a
123 `Google service account private key <https://docs.cloud.google.com/iam/docs/keys-create-delete#creating>`_ file
124 or the file contents.
125
126 :param file: Path to a Google service account private key file.
127 :param info: Google service account private key file contents.
128 """
129 if all(_ is None for _ in (file, info)) or all(_ is not None for _ in (file, info)):
130 raise ValueError("Must specify exactly one of file or info")
131
132 if file is not None:
133 with open(file, "r") as f:
134 self._info = json.load(f)
135 elif info is not None:
136 self._info = copy.deepcopy(info)
137
[docs]
138 def get_credentials(self) -> Credentials:
139 return Credentials(
140 access_key="",
141 secret_key="",
142 token=None,
143 expiration=None,
144 custom_fields={"info": copy.deepcopy(self._info)},
145 )
146
[docs]
147 def refresh_credentials(self) -> None:
148 pass
149
150
[docs]
151class GoogleStorageProvider(BaseStorageProvider):
152 """
153 A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with Google Cloud Storage.
154 """
155
156 def __init__(
157 self,
158 project_id: str = os.getenv("GOOGLE_CLOUD_PROJECT_ID", ""),
159 endpoint_url: str = "",
160 base_path: str = "",
161 credentials_provider: Optional[CredentialsProvider] = None,
162 config_dict: Optional[dict[str, Any]] = None,
163 telemetry_provider: Optional[Callable[[], Telemetry]] = None,
164 **kwargs: Any,
165 ):
166 """
167 Initializes the :py:class:`GoogleStorageProvider` with the project ID and optional credentials provider.
168
169 :param project_id: The Google Cloud project ID.
170 :param endpoint_url: The custom endpoint URL for the GCS service.
171 :param base_path: The root prefix path within the bucket where all operations will be scoped.
172 :param credentials_provider: The provider to retrieve GCS credentials.
173 :param config_dict: Resolved MSC config.
174 :param telemetry_provider: A function that provides a telemetry instance.
175 :param max_pool_connections: Maximum connection pool size for the Rust client.
176 """
177 super().__init__(
178 base_path=base_path,
179 provider_name=PROVIDER,
180 config_dict=config_dict,
181 telemetry_provider=telemetry_provider,
182 )
183
184 self._project_id = project_id
185 self._endpoint_url = endpoint_url
186 self._credentials_provider = credentials_provider
187 self._skip_signature = kwargs.get("skip_signature", False)
188 self._gcs_client = self._create_gcs_client()
189 self._multipart_threshold = kwargs.get("multipart_threshold", DEFAULT_MULTIPART_THRESHOLD)
190 self._multipart_chunksize = kwargs.get("multipart_chunksize", DEFAULT_MULTIPART_CHUNKSIZE)
191 self._io_chunksize = kwargs.get("io_chunksize", DEFAULT_IO_CHUNKSIZE)
192 self._max_concurrency = kwargs.get("max_concurrency", PYTHON_MAX_CONCURRENCY)
193 self._rust_client = None
194 if "rust_client" in kwargs:
195 # Inherit the rust client options from the kwargs
196 rust_client_options = copy.deepcopy(kwargs["rust_client"])
197 if "max_pool_connections" in kwargs:
198 rust_client_options["max_pool_connections"] = kwargs["max_pool_connections"]
199 if "max_concurrency" in kwargs:
200 rust_client_options["max_concurrency"] = kwargs["max_concurrency"]
201 if "multipart_chunksize" in kwargs:
202 rust_client_options["multipart_chunksize"] = kwargs["multipart_chunksize"]
203 if "read_timeout" in kwargs:
204 rust_client_options["read_timeout"] = kwargs["read_timeout"]
205 if "connect_timeout" in kwargs:
206 rust_client_options["connect_timeout"] = kwargs["connect_timeout"]
207 self._rust_client = self._create_rust_client(rust_client_options)
208
209 def _create_gcs_client(self) -> storage.Client:
210 client_options = {}
211 if self._endpoint_url:
212 client_options["api_endpoint"] = self._endpoint_url
213
214 # Use anonymous credentials for public buckets when skip_signature is enabled
215 if self._skip_signature:
216 return storage.Client(
217 project=self._project_id,
218 credentials=auth_credentials.AnonymousCredentials(),
219 client_options=client_options,
220 )
221
222 if self._credentials_provider:
223 if isinstance(self._credentials_provider, GoogleIdentityPoolCredentialsProvider):
224 audience = self._credentials_provider.get_credentials().get_custom_field("audience")
225 token = self._credentials_provider.get_credentials().get_custom_field("token")
226
227 # Use Workload Identity Federation (WIF)
228 identity_pool_credentials = identity_pool.Credentials(
229 audience=audience,
230 subject_token_type="urn:ietf:params:oauth:token-type:id_token",
231 subject_token_supplier=StringTokenSupplier(token),
232 )
233 return storage.Client(
234 project=self._project_id, credentials=identity_pool_credentials, client_options=client_options
235 )
236 elif isinstance(self._credentials_provider, GoogleServiceAccountCredentialsProvider):
237 # Use service account key.
238 service_account_credentials = service_account.Credentials.from_service_account_info(
239 info=self._credentials_provider.get_credentials().get_custom_field("info")
240 )
241 return storage.Client(
242 project=self._project_id, credentials=service_account_credentials, client_options=client_options
243 )
244 else:
245 # Use OAuth 2.0 token
246 token = self._credentials_provider.get_credentials().token
247 creds = OAuth2Credentials(token=token)
248 return storage.Client(project=self._project_id, credentials=creds, client_options=client_options)
249 else:
250 return storage.Client(project=self._project_id, client_options=client_options)
251
252 def _create_rust_client(self, rust_client_options: Optional[dict[str, Any]] = None):
253 if self._endpoint_url:
254 logger.warning("Rust client for GCS does not support customized endpoint URL, skipping rust client")
255 return None
256
257 configs = dict(rust_client_options) if rust_client_options else {}
258
259 # Extract and parse retry configuration
260 retry_config = parse_retry_config(configs)
261
262 if "application_credentials" not in configs and os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
263 configs["application_credentials"] = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
264 if "service_account_key" not in configs and os.getenv("GOOGLE_SERVICE_ACCOUNT_KEY"):
265 configs["service_account_key"] = os.getenv("GOOGLE_SERVICE_ACCOUNT_KEY")
266 if "service_account_path" not in configs and os.getenv("GOOGLE_SERVICE_ACCOUNT"):
267 configs["service_account_path"] = os.getenv("GOOGLE_SERVICE_ACCOUNT")
268 if "service_account_path" not in configs and os.getenv("GOOGLE_SERVICE_ACCOUNT_PATH"):
269 configs["service_account_path"] = os.getenv("GOOGLE_SERVICE_ACCOUNT_PATH")
270
271 if self._skip_signature and "skip_signature" not in configs:
272 configs["skip_signature"] = True
273
274 if "bucket" not in configs:
275 bucket, _ = split_path(self._base_path)
276 configs["bucket"] = bucket
277
278 if "max_pool_connections" not in configs:
279 configs["max_pool_connections"] = DEFAULT_MAX_POOL_CONNECTIONS
280
281 if self._credentials_provider:
282 if isinstance(self._credentials_provider, GoogleIdentityPoolCredentialsProvider):
283 # Workload Identity Federation (WIF) is not supported by the rust client:
284 # https://github.com/apache/arrow-rs-object-store/issues/258
285 logger.warning("Rust client for GCS doesn't support Workload Identity Federation, skipping rust client")
286 return None
287 if isinstance(self._credentials_provider, GoogleServiceAccountCredentialsProvider):
288 # Use service account key.
289 configs["service_account_key"] = json.dumps(
290 self._credentials_provider.get_credentials().get_custom_field("info")
291 )
292 return RustClient(
293 provider=PROVIDER,
294 configs=configs,
295 retry=retry_config,
296 )
297 try:
298 return RustClient(
299 provider=PROVIDER,
300 configs=configs,
301 credentials_provider=self._credentials_provider,
302 retry=retry_config,
303 )
304 except Exception as e:
305 logger.warning(f"Failed to create rust client for GCS: {e}, falling back to Python client")
306 return None
307
308 def _refresh_gcs_client_if_needed(self) -> None:
309 """
310 Refreshes the GCS client if the current credentials are expired.
311 """
312 if self._credentials_provider:
313 credentials = self._credentials_provider.get_credentials()
314 if credentials.is_expired():
315 self._credentials_provider.refresh_credentials()
316 self._gcs_client = self._create_gcs_client()
317
318 def _translate_errors(
319 self,
320 func: Callable[[], _T],
321 operation: str,
322 bucket: str,
323 key: str,
324 ) -> _T:
325 """
326 Translates errors like timeouts and client errors.
327
328 :param func: The function that performs the actual GCS operation.
329 :param operation: The type of operation being performed (e.g., "PUT", "GET", "DELETE").
330 :param bucket: The name of the GCS bucket involved in the operation.
331 :param key: The key of the object within the GCS bucket.
332
333 :return: The result of the GCS operation, typically the return value of the `func` callable.
334 """
335 try:
336 return func()
337 except GoogleAPICallError as error:
338 status_code = error.code if error.code else -1
339 error_info = f"status_code: {status_code}, message: {error.message}"
340 if status_code == 404:
341 raise FileNotFoundError(f"Object {bucket}/{key} does not exist.") # pylint: disable=raise-missing-from
342 elif status_code == 412:
343 raise PreconditionFailedError(
344 f"Failed to {operation} object(s) at {bucket}/{key}. {error_info}"
345 ) from error
346 elif status_code == 304:
347 # for if_none_match with a specific etag condition.
348 raise NotModifiedError(f"Object {bucket}/{key} has not been modified.") from error
349 else:
350 raise RuntimeError(f"Failed to {operation} object(s) at {bucket}/{key}. {error_info}") from error
351 except InvalidResponse as error:
352 response_text = error.response.text
353 error_details = f"error: {error}, error_response_text: {response_text}"
354 # Check for NoSuchUpload within the response text
355 if "NoSuchUpload" in response_text:
356 raise RetryableError(f"Multipart upload failed for {bucket}/{key}, {error_details}") from error
357 else:
358 raise RuntimeError(f"Failed to {operation} object(s) at {bucket}/{key}. {error_details}") from error
359 except RustRetryableError as error:
360 raise RetryableError(
361 f"Failed to {operation} object(s) at {bucket}/{key} due to retryable error from Rust. "
362 f"error_type: {type(error).__name__}"
363 ) from error
364 except RustClientError as error:
365 message = error.args[0]
366 status_code = error.args[1]
367 if status_code == 404:
368 raise FileNotFoundError(f"Object {bucket}/{key} does not exist. {message}") from error
369 elif status_code == 403:
370 raise PermissionError(
371 f"Permission denied to {operation} object(s) at {bucket}/{key}. {message}"
372 ) from error
373 else:
374 raise RetryableError(
375 f"Failed to {operation} object(s) at {bucket}/{key}. {message}. status_code: {status_code}"
376 ) from error
377 except RetryableError:
378 raise
379 except Exception as error:
380 error_details = str(error)
381 raise RuntimeError(
382 f"Failed to {operation} object(s) at {bucket}/{key}. error_type: {type(error).__name__}, {error_details}"
383 ) from error
384
385 def _put_object(
386 self,
387 path: str,
388 body: bytes,
389 if_match: Optional[str] = None,
390 if_none_match: Optional[str] = None,
391 attributes: Optional[dict[str, str]] = None,
392 ) -> int:
393 """
394 Uploads an object to Google Cloud Storage.
395
396 :param path: The path to the object to upload.
397 :param body: The content of the object to upload.
398 :param if_match: Optional ETag to match against the object.
399 :param if_none_match: Optional ETag to match against the object.
400 :param attributes: Optional attributes to attach to the object.
401 """
402 bucket, key = split_path(path)
403 self._refresh_gcs_client_if_needed()
404
405 def _invoke_api() -> int:
406 bucket_obj = self._gcs_client.bucket(bucket)
407 blob = bucket_obj.blob(key)
408
409 kwargs = {}
410
411 if if_match:
412 kwargs["if_generation_match"] = int(if_match) # 412 error code
413 if if_none_match:
414 if if_none_match == "*":
415 raise NotImplementedError("if_none_match='*' is not supported for GCS")
416 else:
417 kwargs["if_generation_not_match"] = int(if_none_match) # 304 error code
418
419 validated_attributes = validate_attributes(attributes)
420 if validated_attributes:
421 blob.metadata = validated_attributes
422
423 if (
424 self._rust_client
425 # Rust client doesn't support creating objects with trailing /, see https://github.com/apache/arrow-rs/issues/7026
426 and not path.endswith("/")
427 and not kwargs
428 and not validated_attributes
429 ):
430 run_async_rust_client_method(self._rust_client, "put", key, body)
431 else:
432 blob.upload_from_string(body, **kwargs)
433
434 return len(body)
435
436 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
437
438 def _get_object(self, path: str, byte_range: Optional[Range] = None) -> bytes:
439 bucket, key = split_path(path)
440 self._refresh_gcs_client_if_needed()
441
442 def _invoke_api() -> bytes:
443 bucket_obj = self._gcs_client.bucket(bucket)
444 blob = bucket_obj.blob(key)
445 if byte_range:
446 if self._rust_client:
447 return run_async_rust_client_method(
448 self._rust_client, "get", key, byte_range.offset, byte_range.offset + byte_range.size - 1
449 )
450 else:
451 return blob.download_as_bytes(
452 start=byte_range.offset, end=byte_range.offset + byte_range.size - 1, single_shot_download=True
453 )
454 else:
455 if self._rust_client:
456 return run_async_rust_client_method(self._rust_client, "get", key)
457 else:
458 return blob.download_as_bytes(single_shot_download=True)
459
460 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
461
462 def _copy_object(self, src_path: str, dest_path: str) -> int:
463 src_bucket, src_key = split_path(src_path)
464 dest_bucket, dest_key = split_path(dest_path)
465 self._refresh_gcs_client_if_needed()
466
467 src_object = self._get_object_metadata(src_path)
468
469 def _invoke_api() -> int:
470 source_bucket_obj = self._gcs_client.bucket(src_bucket)
471 source_blob = source_bucket_obj.blob(src_key)
472
473 destination_bucket_obj = self._gcs_client.bucket(dest_bucket)
474 destination_blob = destination_bucket_obj.blob(dest_key)
475
476 rewrite_tokens = [None]
477 while len(rewrite_tokens) > 0:
478 rewrite_token = rewrite_tokens.pop()
479 next_rewrite_token, _, _ = destination_blob.rewrite(source=source_blob, token=rewrite_token)
480 if next_rewrite_token is not None:
481 rewrite_tokens.append(next_rewrite_token)
482
483 return src_object.content_length
484
485 return self._translate_errors(_invoke_api, operation="COPY", bucket=src_bucket, key=src_key)
486
487 def _delete_object(self, path: str, if_match: Optional[str] = None) -> None:
488 bucket, key = split_path(path)
489 self._refresh_gcs_client_if_needed()
490
491 def _invoke_api() -> None:
492 bucket_obj = self._gcs_client.bucket(bucket)
493 blob = bucket_obj.blob(key)
494
495 # If if_match is provided, use it as a precondition
496 if if_match:
497 generation = int(if_match)
498 blob.delete(if_generation_match=generation)
499 else:
500 # No if_match check needed, just delete
501 blob.delete()
502
503 return self._translate_errors(_invoke_api, operation="DELETE", bucket=bucket, key=key)
504
505 def _delete_objects(self, paths: list[str]) -> None:
506 if not paths:
507 return
508
509 by_bucket: dict[str, list[str]] = {}
510 for p in paths:
511 bucket, key = split_path(p)
512 by_bucket.setdefault(bucket, []).append(key)
513 self._refresh_gcs_client_if_needed()
514
515 GCS_BATCH_LIMIT = 100
516
517 def _invoke_api() -> None:
518 for bucket, keys in by_bucket.items():
519 bucket_obj = self._gcs_client.bucket(bucket)
520 for i in range(0, len(keys), GCS_BATCH_LIMIT):
521 chunk = keys[i : i + GCS_BATCH_LIMIT]
522 with self._gcs_client.batch(raise_exception=False) as batch:
523 for k in chunk:
524 bucket_obj.blob(k).delete()
525 if not hasattr(batch, "_responses"):
526 raise RuntimeError("GCS batch delete did not expose responses.")
527 for response in batch._responses:
528 status_code = response.status_code
529 if 200 <= status_code < 300 or status_code == 404:
530 continue
531 if status_code in {408, 429} or 500 <= status_code < 600:
532 raise RetryableError(
533 f"GCS batch delete failed with status_code: {status_code}, response: {response.text}"
534 )
535 raise RuntimeError(
536 f"GCS batch delete failed with status_code: {status_code}, response: {response.text}"
537 )
538
539 bucket_desc = "(" + "|".join(by_bucket) + ")"
540 key_desc = "(" + "|".join(str(len(keys)) for keys in by_bucket.values()) + " keys)"
541 self._translate_errors(_invoke_api, operation="DELETE_MANY", bucket=bucket_desc, key=key_desc)
542
543 def _is_dir(self, path: str) -> bool:
544 # Ensure the path ends with '/' to mimic a directory
545 path = self._append_delimiter(path)
546
547 bucket, key = split_path(path)
548 self._refresh_gcs_client_if_needed()
549
550 def _invoke_api() -> bool:
551 bucket_obj = self._gcs_client.bucket(bucket)
552 # List objects with the given prefix
553 blobs = bucket_obj.list_blobs(
554 prefix=key,
555 delimiter="/",
556 )
557 # Check if there are any contents or common prefixes
558 return any(True for _ in blobs) or any(True for _ in blobs.prefixes)
559
560 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=key)
561
562 def _make_symlink(self, path: str, target: str) -> None:
563 bucket_name, key = split_path(path)
564 target_bucket, target_key = split_path(target)
565 if bucket_name != target_bucket:
566 raise ValueError(f"Cannot create cross-bucket symlink: '{bucket_name}' -> '{target_bucket}'.")
567 relative_target = ObjectMetadata.encode_symlink_target(key, target_key)
568 self._refresh_gcs_client_if_needed()
569
570 def _invoke_api() -> None:
571 bucket = self._gcs_client.bucket(bucket_name)
572 blob = bucket.blob(key)
573 blob.metadata = {"msc-symlink-target": relative_target}
574 blob.upload_from_string(b"")
575
576 self._translate_errors(_invoke_api, operation="PUT", bucket=bucket_name, key=key)
577
578 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata:
579 bucket, key = split_path(path)
580 if path.endswith("/") or (bucket and not key):
581 # If path ends with "/" or empty key name is provided, then assume it's a "directory",
582 # which metadata is not guaranteed to exist for cases such as
583 # "virtual prefix" that was never explicitly created.
584 if self._is_dir(path):
585 return ObjectMetadata(
586 key=path, type="directory", content_length=0, last_modified=AWARE_DATETIME_MIN, etag=None
587 )
588 else:
589 raise FileNotFoundError(f"Directory {path} does not exist.")
590 else:
591 self._refresh_gcs_client_if_needed()
592
593 def _invoke_api() -> ObjectMetadata:
594 bucket_obj = self._gcs_client.bucket(bucket)
595 blob = bucket_obj.get_blob(key)
596 if not blob:
597 raise NotFound(f"Blob {key} not found in bucket {bucket}")
598 user_metadata = dict(blob.metadata) if blob.metadata else None
599 symlink_target = user_metadata.get("msc-symlink-target") if user_metadata else None
600 return ObjectMetadata(
601 key=path,
602 content_length=blob.size or 0,
603 content_type=blob.content_type,
604 last_modified=blob.updated or AWARE_DATETIME_MIN,
605 etag=str(blob.generation),
606 metadata=user_metadata,
607 symlink_target=symlink_target,
608 )
609
610 try:
611 return self._translate_errors(_invoke_api, operation="HEAD", bucket=bucket, key=key)
612 except FileNotFoundError as error:
613 if strict:
614 # If the object does not exist on the given path, we will append a trailing slash and
615 # check if the path is a directory.
616 path = self._append_delimiter(path)
617 if self._is_dir(path):
618 return ObjectMetadata(
619 key=path,
620 type="directory",
621 content_length=0,
622 last_modified=AWARE_DATETIME_MIN,
623 )
624 raise error
625
626 def _list_objects(
627 self,
628 path: str,
629 start_after: Optional[str] = None,
630 end_at: Optional[str] = None,
631 include_directories: bool = False,
632 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
633 ) -> Iterator[ObjectMetadata]:
634 bucket, prefix = split_path(path)
635
636 # Get the prefix of the start_after and end_at paths relative to the bucket.
637 if start_after:
638 _, start_after = split_path(start_after)
639 if end_at:
640 _, end_at = split_path(end_at)
641
642 self._refresh_gcs_client_if_needed()
643
644 def _invoke_api() -> Iterator[ObjectMetadata]:
645 bucket_obj = self._gcs_client.bucket(bucket)
646 if include_directories:
647 blobs = bucket_obj.list_blobs(
648 prefix=prefix,
649 # This is ≥ instead of >.
650 start_offset=start_after,
651 delimiter="/",
652 )
653 else:
654 blobs = bucket_obj.list_blobs(
655 prefix=prefix,
656 # This is ≥ instead of >.
657 start_offset=start_after,
658 )
659
660 # GCS guarantees lexicographical order.
661 for blob in blobs:
662 key = blob.name
663 if (start_after is None or start_after < key) and (end_at is None or key <= end_at):
664 if key.endswith("/"):
665 if include_directories:
666 yield ObjectMetadata(
667 key=os.path.join(bucket, key.rstrip("/")),
668 type="directory",
669 content_length=0,
670 last_modified=blob.updated,
671 )
672 else:
673 user_metadata = blob.metadata
674 symlink_target = user_metadata.get("msc-symlink-target") if user_metadata else None
675 yield ObjectMetadata(
676 key=os.path.join(bucket, key),
677 content_length=blob.size,
678 content_type=blob.content_type,
679 last_modified=blob.updated,
680 etag=blob.etag,
681 symlink_target=symlink_target,
682 )
683 elif start_after != key:
684 return
685
686 # The directories must be accessed last.
687 if include_directories:
688 for directory in blobs.prefixes:
689 prefix_key = directory.rstrip("/")
690 # Filter by start_after and end_at if specified
691 if (start_after is None or start_after < prefix_key) and (end_at is None or prefix_key <= end_at):
692 yield ObjectMetadata(
693 key=os.path.join(bucket, prefix_key),
694 type="directory",
695 content_length=0,
696 last_modified=AWARE_DATETIME_MIN,
697 )
698
699 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=prefix)
700
701 @property
702 def supports_parallel_listing(self) -> bool:
703 return True
704
705 def _upload_file(self, remote_path: str, f: Union[str, IO], attributes: Optional[dict[str, str]] = None) -> int:
706 bucket, key = split_path(remote_path)
707 file_size: int = 0
708 self._refresh_gcs_client_if_needed()
709
710 if isinstance(f, str):
711 file_size = os.path.getsize(f)
712
713 # Upload small files
714 if file_size <= self._multipart_threshold:
715 if self._rust_client and not attributes:
716 run_async_rust_client_method(self._rust_client, "upload", f, key)
717 else:
718 with open(f, "rb") as fp:
719 self._put_object(remote_path, fp.read(), attributes=attributes)
720 return file_size
721
722 # Upload large files using transfer manager
723 def _invoke_api() -> int:
724 if self._rust_client and not attributes:
725 run_async_rust_client_method(self._rust_client, "upload_multipart_from_file", f, key)
726 else:
727 bucket_obj = self._gcs_client.bucket(bucket)
728 blob = bucket_obj.blob(key)
729 # GCS will raise an error if blob.metadata is None
730 validated_attributes = validate_attributes(attributes)
731 if validated_attributes is not None:
732 blob.metadata = validated_attributes
733 transfer_manager.upload_chunks_concurrently(
734 f,
735 blob,
736 chunk_size=self._multipart_chunksize,
737 max_workers=self._max_concurrency,
738 worker_type=transfer_manager.THREAD,
739 )
740
741 return file_size
742
743 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
744 else:
745 f.seek(0, io.SEEK_END)
746 file_size = f.tell()
747 f.seek(0)
748
749 # Upload small files
750 if file_size <= self._multipart_threshold:
751 if isinstance(f, io.StringIO):
752 self._put_object(remote_path, f.read().encode("utf-8"), attributes=attributes)
753 else:
754 self._put_object(remote_path, f.read(), attributes=attributes)
755 return file_size
756
757 # Upload large files using transfer manager
758 def _invoke_api() -> int:
759 bucket_obj = self._gcs_client.bucket(bucket)
760 blob = bucket_obj.blob(key)
761 validated_attributes = validate_attributes(attributes)
762 if validated_attributes:
763 blob.metadata = validated_attributes
764 if isinstance(f, io.StringIO):
765 mode = "w"
766 else:
767 mode = "wb"
768
769 # transfer manager does not support uploading a file object
770 with tempfile.NamedTemporaryFile(mode=mode, delete=False, prefix=".") as fp:
771 temp_file_path = fp.name
772 fp.write(f.read())
773
774 transfer_manager.upload_chunks_concurrently(
775 temp_file_path,
776 blob,
777 chunk_size=self._multipart_chunksize,
778 max_workers=self._max_concurrency,
779 worker_type=transfer_manager.THREAD,
780 )
781
782 os.unlink(temp_file_path)
783
784 return file_size
785
786 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
787
788 def _download_file(self, remote_path: str, f: Union[str, IO], metadata: Optional[ObjectMetadata] = None) -> int:
789 self._refresh_gcs_client_if_needed()
790
791 if metadata is None:
792 metadata = self._get_object_metadata(remote_path)
793
794 bucket, key = split_path(remote_path)
795
796 if isinstance(f, str):
797 if os.path.dirname(f):
798 safe_makedirs(os.path.dirname(f))
799 # Download small files
800 if metadata.content_length <= self._multipart_threshold:
801 if self._rust_client:
802 run_async_rust_client_method(self._rust_client, "download", key, f)
803 else:
804 with tempfile.NamedTemporaryFile(mode="wb", delete=False, dir=os.path.dirname(f), prefix=".") as fp:
805 temp_file_path = fp.name
806 fp.write(self._get_object(remote_path))
807 os.rename(src=temp_file_path, dst=f)
808 return metadata.content_length
809
810 # Download large files using transfer manager
811 def _invoke_api() -> int:
812 bucket_obj = self._gcs_client.bucket(bucket)
813 blob = bucket_obj.blob(key)
814 if self._rust_client:
815 run_async_rust_client_method(self._rust_client, "download_multipart_to_file", key, f)
816 else:
817 with tempfile.NamedTemporaryFile(mode="wb", delete=False, dir=os.path.dirname(f), prefix=".") as fp:
818 temp_file_path = fp.name
819 transfer_manager.download_chunks_concurrently(
820 blob,
821 temp_file_path,
822 chunk_size=self._io_chunksize,
823 max_workers=self._max_concurrency,
824 worker_type=transfer_manager.THREAD,
825 )
826 os.rename(src=temp_file_path, dst=f)
827
828 return metadata.content_length
829
830 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
831 else:
832 # Download small files
833 if metadata.content_length <= self._multipart_threshold:
834 response = self._get_object(remote_path)
835 # Python client returns `bytes`, but Rust client returns an object that implements the buffer protocol,
836 # so we need to check whether `.decode()` is available.
837 if isinstance(f, io.StringIO):
838 if hasattr(response, "decode"):
839 f.write(response.decode("utf-8"))
840 else:
841 f.write(codecs.decode(memoryview(response), "utf-8"))
842 else:
843 f.write(response)
844 return metadata.content_length
845
846 # Download large files using transfer manager
847 def _invoke_api() -> int:
848 bucket_obj = self._gcs_client.bucket(bucket)
849 blob = bucket_obj.blob(key)
850
851 # transfer manager does not support downloading to a file object
852 with tempfile.NamedTemporaryFile(mode="wb", delete=False, prefix=".") as fp:
853 temp_file_path = fp.name
854 transfer_manager.download_chunks_concurrently(
855 blob,
856 temp_file_path,
857 chunk_size=self._io_chunksize,
858 max_workers=self._max_concurrency,
859 worker_type=transfer_manager.THREAD,
860 )
861
862 if isinstance(f, io.StringIO):
863 with open(temp_file_path, "r") as fp:
864 f.write(fp.read())
865 else:
866 with open(temp_file_path, "rb") as fp:
867 f.write(fp.read())
868
869 os.unlink(temp_file_path)
870
871 return metadata.content_length
872
873 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)