"""Resolution of a dataset URI to the object store that serves it.
A backend claims one URI scheme. The empty scheme is the local filesystem,
whose "URIs" are ordinary absolute paths.
Factories import their implementation module inside the call body so that
building the registry never re-enters ``signal_dataset.storage``. Note that this
is not what keeps optional cloud SDKs out of ``import signal_dataset``: the
storage facade exports every built-in store eagerly, so those modules do load.
What defers the SDK is each backend's client loader, called only when a store is
constructed without an injected client.
"""
from __future__ import annotations
import logging
import os
import threading
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol
from urllib.parse import urlparse
from signal_dataset.storage.contracts import ObjectStore
from signal_dataset.storage.paths import gcs_parts, s3_parts
logger = logging.getLogger(__name__)
LOCAL_SCHEME = ""
GCS_SCHEME = "gs"
S3_SCHEME = "s3"
class ObjectStoreFactory(Protocol):
"""Builds a store for one URI under a backend's scheme."""
def __call__(self, uri: str, *, root_is_directory: bool) -> ObjectStore: ...
[docs]
@dataclass(frozen=True, slots=True, kw_only=True)
class ObjectStoreBackend:
"""One storage service, keyed by the URI scheme it claims."""
scheme: str
filesystem: bool
"""Whether this backend's URIs are operating-system paths."""
factory: ObjectStoreFactory
validate_root: Callable[[str], None]
[docs]
class ObjectStoreRegistry:
"""Maps URI schemes to backends.
Registration is additive and guarded: re-registering a scheme raises unless
``replace`` is passed, so an import cycle or a duplicated plugin cannot
silently redirect an existing scheme. There is deliberately no way to
unregister.
"""
def __init__(self) -> None:
self._backends: dict[str, ObjectStoreBackend] = {}
self._schemes: frozenset[str] = frozenset()
self._lock = threading.Lock()
[docs]
def register(self, backend: ObjectStoreBackend, *, replace: bool = False) -> None:
with self._lock:
if backend.scheme in self._backends and not replace:
raise ValueError(f"storage scheme is already registered: {backend.scheme!r}")
self._backends[backend.scheme] = backend
self._schemes = frozenset(self._backends)
logger.debug("Registered storage backend", extra={"scheme": backend.scheme})
[docs]
def backend_for(self, uri: str) -> ObjectStoreBackend:
scheme = urlparse(uri).scheme
backend = self._backends.get(scheme)
if backend is None:
known = ", ".join(sorted(repr(name) for name in self._schemes))
raise ValueError(f"unsupported storage URI scheme: {scheme!r} (registered: {known})")
return backend
[docs]
def schemes(self) -> frozenset[str]:
# Reads are lock-free: the snapshot is immutable and rebound atomically
# under the lock in register(). This runs once per manifest entry.
return self._schemes
#: Colon-separated absolute path prefixes that are mounted buckets.
#:
#: An environment variable rather than an argument, because "this machine has a
#: gcsfuse mount at /mnt/gcs" is a property of the deployment, not of the code.
#: It belongs next to the gcsfuse invocation in a launcher or Dockerfile, and it
#: leaves the training script identical whether it reads from a mount or from
#: gs:// directly.
#:
#: It is also the only mechanism that reaches everywhere. DatasetView.publish and
#: materialize take no object_store argument at all and build their stores through
#: this factory, so a capabilities= argument could never get to them.
MOUNT_ROOTS_VARIABLE = "SIGNAL_DATASET_MOUNT_ROOTS"
def mounted_roots(environ: Mapping[str, str] | None = None) -> tuple[str, ...]:
values = os.environ if environ is None else environ
declared = values.get(MOUNT_ROOTS_VARIABLE, "")
return tuple(
str(Path(entry).absolute()) for entry in declared.split(os.pathsep) if entry.strip()
)
def is_mounted(path: str, *, environ: Mapping[str, str] | None = None) -> bool:
"""Whether `path` lies under a declared mount root."""
resolved = Path(path).absolute()
for root in mounted_roots(environ):
if resolved == Path(root) or Path(root) in resolved.parents:
return True
return False
def local_object_store(uri: str, *, root_is_directory: bool) -> ObjectStore:
from signal_dataset.storage.backends.local import (
MOUNTED_BUCKET,
LocalObjectStore,
)
root = uri if root_is_directory else str(Path(uri).parent)
capabilities = MOUNTED_BUCKET if is_mounted(root) else None
return LocalObjectStore(root=root, capabilities=capabilities)
def gcs_object_store(uri: str, *, root_is_directory: bool) -> ObjectStore:
from signal_dataset.storage.backends.gcs import GCSObjectStore
del uri, root_is_directory # a GCS client is not bound to one location
return GCSObjectStore()
def s3_object_store(uri: str, *, root_is_directory: bool) -> ObjectStore:
from signal_dataset.storage.backends.s3 import S3ObjectStore
del uri, root_is_directory # an S3 client is not bound to one location
return S3ObjectStore()
def validate_local_root(uri: str) -> None:
del uri # any non-empty, non-root path is acceptable
def validate_gcs_root(uri: str) -> None:
gcs_parts(uri, allow_empty_name=True)
def validate_s3_root(uri: str) -> None:
s3_parts(uri, allow_empty_name=True)
def build_default_registry() -> ObjectStoreRegistry:
"""Build a registry holding the backends that ship with the package."""
registry = ObjectStoreRegistry()
registry.register(
ObjectStoreBackend(
scheme=LOCAL_SCHEME,
filesystem=True,
factory=local_object_store,
validate_root=validate_local_root,
)
)
registry.register(
ObjectStoreBackend(
scheme=GCS_SCHEME,
filesystem=False,
factory=gcs_object_store,
validate_root=validate_gcs_root,
)
)
registry.register(
ObjectStoreBackend(
scheme=S3_SCHEME,
filesystem=False,
factory=s3_object_store,
validate_root=validate_s3_root,
)
)
return registry
_DEFAULT_REGISTRY: ObjectStoreRegistry | None = None
_DEFAULT_LOCK = threading.Lock()
def default_registry() -> ObjectStoreRegistry:
"""Return the process-wide registry, building it on first use."""
global _DEFAULT_REGISTRY
with _DEFAULT_LOCK:
if _DEFAULT_REGISTRY is None:
_DEFAULT_REGISTRY = build_default_registry()
return _DEFAULT_REGISTRY
[docs]
def object_store_for(
uri: str,
*,
root_is_directory: bool = False,
registry: ObjectStoreRegistry | None = None,
) -> ObjectStore:
"""Return a store able to serve ``uri``."""
resolved = default_registry() if registry is None else registry
return resolved.backend_for(uri).factory(uri, root_is_directory=root_is_directory)
def is_absolute_uri(uri: str, *, registry: ObjectStoreRegistry | None = None) -> bool:
"""Whether a persisted path is a full location rather than dataset-relative.
Absolute means a leading ``/`` or a registered scheme in full ``scheme://``
form. Both halves matter, because ``safe_relative_path`` permits ``:``:
``"a:b/c"`` is a legal relative path whose scheme-looking prefix is ``"a"``,
and so is ``"gs:notascheme"``, whose prefix is a registered scheme. Treating
either as a location would resolve it to the wrong object.
The prefix is split by hand rather than with ``urlparse`` on purpose. This
replaced a ``str.startswith`` test, which cannot fail; ``urlparse`` raises
``ValueError`` on a malformed netloc such as ``"gs://[oops"``. One caller is
the persisted-format deserializer, whose job is to classify corrupt input
rather than leak a urllib error — and answering "relative" there would be
worse still, because ``safe_relative_path`` rewrites that value to
``"gs:/[oops"``, silently naming a different object.
"""
if uri.startswith("/"):
return True
scheme, separator, _ = uri.partition("://")
if not separator or not scheme:
return False
resolved = default_registry() if registry is None else registry
return scheme in resolved.schemes()