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