"""Google Cloud Storage backend using object generations and preconditions."""
from __future__ import annotations
from pathlib import Path
from typing import Any, cast
from signal_dataset.errors import PublicationCollisionError
from signal_dataset.storage.backends.gcs import errors
from signal_dataset.storage.backends.gcs.client import load_google_storage
from signal_dataset.storage.contracts import ObjectInfo, ObjectStore, ObjectVersion
from signal_dataset.storage.paths import gcs_parts
[docs]
class GCSObjectStore(ObjectStore):
def __init__(
self,
client: Any | None = None,
*,
precondition_failed: type[Exception] | None = None,
api_error: type[Exception] | tuple[type[Exception], ...] | None = None,
) -> None:
if client is None:
storage, google_precondition_failed, transport_errors = load_google_storage()
client = storage.Client()
precondition_failed = google_precondition_failed
api_error = transport_errors
self._client = client
# An empty tuple is a valid `except` target that catches nothing, so a
# caller injecting a client without a failure type needs no placeholder.
self._precondition_failed: type[Exception] | tuple[()] = precondition_failed or ()
if api_error is None:
api_error = errors.loaded_transport_errors()
self._api_error: type[Exception] | tuple[type[Exception], ...] = api_error or ()
[docs]
def read(self, uri: str, *, generation: int | None = None) -> bytes:
bucket_name, object_name = gcs_parts(uri)
blob = self._client.bucket(bucket_name).blob(object_name, generation=generation)
try:
return cast(bytes, blob.download_as_bytes(checksum="auto"))
except self._api_error as exc:
raise errors.translate(exc, uri) from exc
[docs]
def read_version(self, uri: str) -> ObjectVersion:
bucket_name, object_name = gcs_parts(uri)
try:
blob = self._client.bucket(bucket_name).get_blob(object_name)
except self._api_error as exc:
raise errors.translate(exc, uri) from exc
if blob is None or blob.generation is None:
raise FileNotFoundError(uri)
generation = int(blob.generation)
return ObjectVersion(self.read(uri, generation=generation), generation)
[docs]
def info(self, uri: str, *, generation: int | None = None) -> ObjectInfo:
bucket_name, object_name = gcs_parts(uri)
try:
blob = self._client.bucket(bucket_name).get_blob(object_name, generation=generation)
except self._api_error as exc:
raise errors.translate(exc, uri) from exc
if blob is None or blob.generation is None or blob.size is None:
raise FileNotFoundError(uri)
return ObjectInfo(int(blob.size), int(blob.generation))
[docs]
def create(self, uri: str, data: bytes, *, content_type: str) -> int:
bucket_name, object_name = gcs_parts(uri)
blob = self._client.bucket(bucket_name).blob(object_name)
try:
blob.upload_from_string(
data,
content_type=content_type,
if_generation_match=0,
checksum="auto",
)
except self._precondition_failed as exc:
raise PublicationCollisionError(uri) from exc
except self._api_error as exc:
raise errors.translate(exc, uri) from exc
if blob.generation is None:
raise RuntimeError("GCS upload response did not include the created object generation")
return int(blob.generation)
[docs]
def create_file(self, uri: str, path: Path, *, content_type: str) -> int:
bucket_name, object_name = gcs_parts(uri)
blob = self._client.bucket(bucket_name).blob(object_name)
try:
blob.upload_from_filename(
str(path),
content_type=content_type,
if_generation_match=0,
checksum="auto",
)
except self._precondition_failed as exc:
raise PublicationCollisionError(uri) from exc
except self._api_error as exc:
raise errors.translate(exc, uri) from exc
if blob.generation is None:
raise RuntimeError("GCS upload response omitted object generation")
return int(blob.generation)
[docs]
def compare_and_swap(
self,
uri: str,
data: bytes,
*,
expected_generation: int | None,
content_type: str,
) -> int:
bucket_name, object_name = gcs_parts(uri)
blob = self._client.bucket(bucket_name).blob(object_name)
match = 0 if expected_generation is None else expected_generation
try:
blob.upload_from_string(
data,
content_type=content_type,
if_generation_match=match,
checksum="auto",
)
except self._precondition_failed as exc:
raise PublicationCollisionError(uri) from exc
except self._api_error as exc:
raise errors.translate(exc, uri) from exc
if blob.generation is None:
raise RuntimeError("GCS upload response omitted object generation")
return int(blob.generation)
[docs]
def list(self, prefix_uri: str) -> list[str]:
bucket_name, object_prefix = gcs_parts(prefix_uri, allow_empty_name=True)
try:
return sorted(
f"gs://{bucket_name}/{blob.name}"
for blob in self._client.list_blobs(bucket_name, prefix=object_prefix)
)
except self._api_error as exc:
raise errors.translate(exc, prefix_uri) from exc
[docs]
def generation(self, uri: str) -> int:
bucket_name, object_name = gcs_parts(uri)
try:
blob = self._client.bucket(bucket_name).get_blob(object_name)
except self._api_error as exc:
raise errors.translate(exc, uri) from exc
if blob is None or blob.generation is None:
raise FileNotFoundError(uri)
return int(blob.generation)