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(self._rust_client, "get", key, byte_range)
448 else:
449 return blob.download_as_bytes(
450 start=byte_range.offset, end=byte_range.offset + byte_range.size - 1, single_shot_download=True
451 )
452 else:
453 if self._rust_client:
454 return run_async_rust_client_method(self._rust_client, "get", key)
455 else:
456 return blob.download_as_bytes(single_shot_download=True)
457
458 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
459
460 def _copy_object(self, src_path: str, dest_path: str) -> int:
461 src_bucket, src_key = split_path(src_path)
462 dest_bucket, dest_key = split_path(dest_path)
463 self._refresh_gcs_client_if_needed()
464
465 src_object = self._get_object_metadata(src_path)
466
467 def _invoke_api() -> int:
468 source_bucket_obj = self._gcs_client.bucket(src_bucket)
469 source_blob = source_bucket_obj.blob(src_key)
470
471 destination_bucket_obj = self._gcs_client.bucket(dest_bucket)
472 destination_blob = destination_bucket_obj.blob(dest_key)
473
474 rewrite_tokens = [None]
475 while len(rewrite_tokens) > 0:
476 rewrite_token = rewrite_tokens.pop()
477 next_rewrite_token, _, _ = destination_blob.rewrite(source=source_blob, token=rewrite_token)
478 if next_rewrite_token is not None:
479 rewrite_tokens.append(next_rewrite_token)
480
481 return src_object.content_length
482
483 return self._translate_errors(_invoke_api, operation="COPY", bucket=src_bucket, key=src_key)
484
485 def _delete_object(self, path: str, if_match: Optional[str] = None) -> None:
486 bucket, key = split_path(path)
487 self._refresh_gcs_client_if_needed()
488
489 def _invoke_api() -> None:
490 bucket_obj = self._gcs_client.bucket(bucket)
491 blob = bucket_obj.blob(key)
492
493 # If if_match is provided, use it as a precondition
494 if if_match:
495 generation = int(if_match)
496 blob.delete(if_generation_match=generation)
497 else:
498 # No if_match check needed, just delete
499 blob.delete()
500
501 return self._translate_errors(_invoke_api, operation="DELETE", bucket=bucket, key=key)
502
503 def _delete_objects(self, paths: list[str]) -> None:
504 if not paths:
505 return
506
507 by_bucket: dict[str, list[str]] = {}
508 for p in paths:
509 bucket, key = split_path(p)
510 by_bucket.setdefault(bucket, []).append(key)
511 self._refresh_gcs_client_if_needed()
512
513 GCS_BATCH_LIMIT = 100
514
515 def _invoke_api() -> None:
516 for bucket, keys in by_bucket.items():
517 bucket_obj = self._gcs_client.bucket(bucket)
518 for i in range(0, len(keys), GCS_BATCH_LIMIT):
519 chunk = keys[i : i + GCS_BATCH_LIMIT]
520 with self._gcs_client.batch(raise_exception=False) as batch:
521 for k in chunk:
522 bucket_obj.blob(k).delete()
523 if not hasattr(batch, "_responses"):
524 raise RuntimeError("GCS batch delete did not expose responses.")
525 for response in batch._responses:
526 status_code = response.status_code
527 if 200 <= status_code < 300 or status_code == 404:
528 continue
529 if status_code in {408, 429} or 500 <= status_code < 600:
530 raise RetryableError(
531 f"GCS batch delete failed with status_code: {status_code}, response: {response.text}"
532 )
533 raise RuntimeError(
534 f"GCS batch delete failed with status_code: {status_code}, response: {response.text}"
535 )
536
537 bucket_desc = "(" + "|".join(by_bucket) + ")"
538 key_desc = "(" + "|".join(str(len(keys)) for keys in by_bucket.values()) + " keys)"
539 self._translate_errors(_invoke_api, operation="DELETE_MANY", bucket=bucket_desc, key=key_desc)
540
541 def _is_dir(self, path: str) -> bool:
542 # Ensure the path ends with '/' to mimic a directory
543 path = self._append_delimiter(path)
544
545 bucket, key = split_path(path)
546 self._refresh_gcs_client_if_needed()
547
548 def _invoke_api() -> bool:
549 bucket_obj = self._gcs_client.bucket(bucket)
550 # List objects with the given prefix
551 blobs = bucket_obj.list_blobs(
552 prefix=key,
553 delimiter="/",
554 )
555 # Check if there are any contents or common prefixes
556 return any(True for _ in blobs) or any(True for _ in blobs.prefixes)
557
558 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=key)
559
560 def _make_symlink(self, path: str, target: str) -> None:
561 bucket_name, key = split_path(path)
562 target_bucket, target_key = split_path(target)
563 if bucket_name != target_bucket:
564 raise ValueError(f"Cannot create cross-bucket symlink: '{bucket_name}' -> '{target_bucket}'.")
565 relative_target = ObjectMetadata.encode_symlink_target(key, target_key)
566 self._refresh_gcs_client_if_needed()
567
568 def _invoke_api() -> None:
569 bucket = self._gcs_client.bucket(bucket_name)
570 blob = bucket.blob(key)
571 blob.metadata = {"msc-symlink-target": relative_target}
572 blob.upload_from_string(b"")
573
574 self._translate_errors(_invoke_api, operation="PUT", bucket=bucket_name, key=key)
575
576 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata:
577 bucket, key = split_path(path)
578 if path.endswith("/") or (bucket and not key):
579 # If path ends with "/" or empty key name is provided, then assume it's a "directory",
580 # which metadata is not guaranteed to exist for cases such as
581 # "virtual prefix" that was never explicitly created.
582 if self._is_dir(path):
583 return ObjectMetadata(
584 key=path, type="directory", content_length=0, last_modified=AWARE_DATETIME_MIN, etag=None
585 )
586 else:
587 raise FileNotFoundError(f"Directory {path} does not exist.")
588 else:
589 self._refresh_gcs_client_if_needed()
590
591 def _invoke_api() -> ObjectMetadata:
592 bucket_obj = self._gcs_client.bucket(bucket)
593 blob = bucket_obj.get_blob(key)
594 if not blob:
595 raise NotFound(f"Blob {key} not found in bucket {bucket}")
596 user_metadata = dict(blob.metadata) if blob.metadata else None
597 symlink_target = user_metadata.get("msc-symlink-target") if user_metadata else None
598 return ObjectMetadata(
599 key=path,
600 content_length=blob.size or 0,
601 content_type=blob.content_type,
602 last_modified=blob.updated or AWARE_DATETIME_MIN,
603 etag=str(blob.generation),
604 metadata=user_metadata,
605 symlink_target=symlink_target,
606 )
607
608 try:
609 return self._translate_errors(_invoke_api, operation="HEAD", bucket=bucket, key=key)
610 except FileNotFoundError as error:
611 if strict:
612 # If the object does not exist on the given path, we will append a trailing slash and
613 # check if the path is a directory.
614 path = self._append_delimiter(path)
615 if self._is_dir(path):
616 return ObjectMetadata(
617 key=path,
618 type="directory",
619 content_length=0,
620 last_modified=AWARE_DATETIME_MIN,
621 )
622 raise error
623
624 def _list_objects(
625 self,
626 path: str,
627 start_after: Optional[str] = None,
628 end_at: Optional[str] = None,
629 include_directories: bool = False,
630 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
631 ) -> Iterator[ObjectMetadata]:
632 bucket, prefix = split_path(path)
633
634 # Get the prefix of the start_after and end_at paths relative to the bucket.
635 if start_after:
636 _, start_after = split_path(start_after)
637 if end_at:
638 _, end_at = split_path(end_at)
639
640 self._refresh_gcs_client_if_needed()
641
642 def _invoke_api() -> Iterator[ObjectMetadata]:
643 bucket_obj = self._gcs_client.bucket(bucket)
644 if include_directories:
645 blobs = bucket_obj.list_blobs(
646 prefix=prefix,
647 # This is ≥ instead of >.
648 start_offset=start_after,
649 delimiter="/",
650 )
651 else:
652 blobs = bucket_obj.list_blobs(
653 prefix=prefix,
654 # This is ≥ instead of >.
655 start_offset=start_after,
656 )
657
658 # GCS guarantees lexicographical order.
659 for blob in blobs:
660 key = blob.name
661 if (start_after is None or start_after < key) and (end_at is None or key <= end_at):
662 if key.endswith("/"):
663 if include_directories:
664 yield ObjectMetadata(
665 key=os.path.join(bucket, key.rstrip("/")),
666 type="directory",
667 content_length=0,
668 last_modified=blob.updated,
669 )
670 else:
671 user_metadata = blob.metadata
672 symlink_target = user_metadata.get("msc-symlink-target") if user_metadata else None
673 yield ObjectMetadata(
674 key=os.path.join(bucket, key),
675 content_length=blob.size,
676 content_type=blob.content_type,
677 last_modified=blob.updated,
678 etag=blob.etag,
679 symlink_target=symlink_target,
680 )
681 elif start_after != key:
682 return
683
684 # The directories must be accessed last.
685 if include_directories:
686 for directory in blobs.prefixes:
687 prefix_key = directory.rstrip("/")
688 # Filter by start_after and end_at if specified
689 if (start_after is None or start_after < prefix_key) and (end_at is None or prefix_key <= end_at):
690 yield ObjectMetadata(
691 key=os.path.join(bucket, prefix_key),
692 type="directory",
693 content_length=0,
694 last_modified=AWARE_DATETIME_MIN,
695 )
696
697 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=prefix)
698
699 @property
700 def supports_parallel_listing(self) -> bool:
701 return True
702
703 def _upload_file(self, remote_path: str, f: Union[str, IO], attributes: Optional[dict[str, str]] = None) -> int:
704 bucket, key = split_path(remote_path)
705 file_size: int = 0
706 self._refresh_gcs_client_if_needed()
707
708 if isinstance(f, str):
709 file_size = os.path.getsize(f)
710
711 # Upload small files
712 if file_size <= self._multipart_threshold:
713 if self._rust_client and not attributes:
714 run_async_rust_client_method(self._rust_client, "upload", f, key)
715 else:
716 with open(f, "rb") as fp:
717 self._put_object(remote_path, fp.read(), attributes=attributes)
718 return file_size
719
720 # Upload large files using transfer manager
721 def _invoke_api() -> int:
722 if self._rust_client and not attributes:
723 run_async_rust_client_method(self._rust_client, "upload_multipart_from_file", f, key)
724 else:
725 bucket_obj = self._gcs_client.bucket(bucket)
726 blob = bucket_obj.blob(key)
727 # GCS will raise an error if blob.metadata is None
728 validated_attributes = validate_attributes(attributes)
729 if validated_attributes is not None:
730 blob.metadata = validated_attributes
731 transfer_manager.upload_chunks_concurrently(
732 f,
733 blob,
734 chunk_size=self._multipart_chunksize,
735 max_workers=self._max_concurrency,
736 worker_type=transfer_manager.THREAD,
737 )
738
739 return file_size
740
741 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
742 else:
743 f.seek(0, io.SEEK_END)
744 file_size = f.tell()
745 f.seek(0)
746
747 # Upload small files
748 if file_size <= self._multipart_threshold:
749 if isinstance(f, io.StringIO):
750 self._put_object(remote_path, f.read().encode("utf-8"), attributes=attributes)
751 else:
752 self._put_object(remote_path, f.read(), attributes=attributes)
753 return file_size
754
755 # Upload large files using transfer manager
756 def _invoke_api() -> int:
757 bucket_obj = self._gcs_client.bucket(bucket)
758 blob = bucket_obj.blob(key)
759 validated_attributes = validate_attributes(attributes)
760 if validated_attributes:
761 blob.metadata = validated_attributes
762 if isinstance(f, io.StringIO):
763 mode = "w"
764 else:
765 mode = "wb"
766
767 # transfer manager does not support uploading a file object
768 with tempfile.NamedTemporaryFile(mode=mode, delete=False, prefix=".") as fp:
769 temp_file_path = fp.name
770 fp.write(f.read())
771
772 transfer_manager.upload_chunks_concurrently(
773 temp_file_path,
774 blob,
775 chunk_size=self._multipart_chunksize,
776 max_workers=self._max_concurrency,
777 worker_type=transfer_manager.THREAD,
778 )
779
780 os.unlink(temp_file_path)
781
782 return file_size
783
784 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
785
786 def _download_file(self, remote_path: str, f: Union[str, IO], metadata: Optional[ObjectMetadata] = None) -> int:
787 self._refresh_gcs_client_if_needed()
788
789 if metadata is None:
790 metadata = self._get_object_metadata(remote_path)
791
792 bucket, key = split_path(remote_path)
793
794 if isinstance(f, str):
795 if os.path.dirname(f):
796 safe_makedirs(os.path.dirname(f))
797 # Download small files
798 if metadata.content_length <= self._multipart_threshold:
799 if self._rust_client:
800 run_async_rust_client_method(self._rust_client, "download", key, f)
801 else:
802 with tempfile.NamedTemporaryFile(mode="wb", delete=False, dir=os.path.dirname(f), prefix=".") as fp:
803 temp_file_path = fp.name
804 fp.write(self._get_object(remote_path))
805 os.rename(src=temp_file_path, dst=f)
806 return metadata.content_length
807
808 # Download large files using transfer manager
809 def _invoke_api() -> int:
810 bucket_obj = self._gcs_client.bucket(bucket)
811 blob = bucket_obj.blob(key)
812 if self._rust_client:
813 run_async_rust_client_method(self._rust_client, "download_multipart_to_file", key, f)
814 else:
815 with tempfile.NamedTemporaryFile(mode="wb", delete=False, dir=os.path.dirname(f), prefix=".") as fp:
816 temp_file_path = fp.name
817 transfer_manager.download_chunks_concurrently(
818 blob,
819 temp_file_path,
820 chunk_size=self._io_chunksize,
821 max_workers=self._max_concurrency,
822 worker_type=transfer_manager.THREAD,
823 )
824 os.rename(src=temp_file_path, dst=f)
825
826 return metadata.content_length
827
828 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
829 else:
830 # Download small files
831 if metadata.content_length <= self._multipart_threshold:
832 response = self._get_object(remote_path)
833 # Python client returns `bytes`, but Rust client returns an object that implements the buffer protocol,
834 # so we need to check whether `.decode()` is available.
835 if isinstance(f, io.StringIO):
836 if hasattr(response, "decode"):
837 f.write(response.decode("utf-8"))
838 else:
839 f.write(codecs.decode(memoryview(response), "utf-8"))
840 else:
841 f.write(response)
842 return metadata.content_length
843
844 # Download large files using transfer manager
845 def _invoke_api() -> int:
846 bucket_obj = self._gcs_client.bucket(bucket)
847 blob = bucket_obj.blob(key)
848
849 # transfer manager does not support downloading to a file object
850 with tempfile.NamedTemporaryFile(mode="wb", delete=False, prefix=".") as fp:
851 temp_file_path = fp.name
852 transfer_manager.download_chunks_concurrently(
853 blob,
854 temp_file_path,
855 chunk_size=self._io_chunksize,
856 max_workers=self._max_concurrency,
857 worker_type=transfer_manager.THREAD,
858 )
859
860 if isinstance(f, io.StringIO):
861 with open(temp_file_path, "r") as fp:
862 f.write(fp.read())
863 else:
864 with open(temp_file_path, "rb") as fp:
865 f.write(fp.read())
866
867 os.unlink(temp_file_path)
868
869 return metadata.content_length
870
871 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)