firecube 0.1.4__py3-none-any.whl

firecube/core/storage/transfer.py

@@ -0,0 +1,142 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Cross-endpoint file transfer primitives.
+
+StorageSession owns one endpoint binding. Transfers between two endpoints
+(e.g., local filesystem → remote S3) require explicit per-endpoint
+sessions. This module provides stateless transfer primitives that open
+each endpoint with its own driver-correct filesystem.
+"""
+
+from __future__ import annotations
+
+import shutil
+from typing import TYPE_CHECKING, Any, cast
+
+from firecube.core.storage.uri import StorageUri
+
+if TYPE_CHECKING:
+    from firecube.core.storage.driver_config import StorageDriverConfig
+    from firecube.core.storage.session import StorageSession
+
+
+def session_for_uri(
+    uri: StorageUri,
+    driver_config: StorageDriverConfig,
+    *,
+    format: str = "tensogram",
+) -> StorageSession:
+    """Construct a minimal StorageSession bound to ``uri``'s endpoint.
+
+    Used by callers (e.g., CLI archive create/restore) that need to
+    transfer to/from a URI that is NOT the same endpoint as their
+    primary product session. The returned session is single-purpose:
+    its product_uri is the input ``uri`` itself, and credentials/region
+    come from ``driver_config``.
+
+    For obstore (which is bucket-scoped), this guarantees the right
+    S3Store is constructed. For fsspec, it ensures consistent driver
+    configuration even when bucket differs from the product session.
+
+    ``format`` defaults to ``"tensogram"`` (the primary archive transfer
+    artifact). Caller may override with any value in
+    ``ProductIdentity.VALID_FORMATS`` if the transfer target is a
+    different format. Format is metadata-only and does not affect
+    transfer behavior; it satisfies ``ProductIdentity`` validation.
+    """
+    from firecube.core.product.identity import ProductIdentity
+    from firecube.core.storage.binding import StorageBinding
+    from firecube.core.storage.session import StorageSession
+
+    identity = ProductIdentity(
+        product_uri=uri,
+        control_root_uri=uri.parent().join(".firecube"),
+        product_name=uri.path.rstrip("/").rsplit("/", 1)[-1] or "transfer",
+        format=format,
+    )
+    return StorageSession(StorageBinding(identity=identity, driver=driver_config))
+
+
+def copy_file(
+    src: StorageUri,
+    dst: StorageUri,
+    *,
+    source_session: StorageSession | None = None,
+    target_session: StorageSession | None = None,
+) -> None:
+    """Copy a single file between two storage endpoints.
+
+    Strict semantics:
+    - If src is remote: source_session is REQUIRED and must be bound to src's endpoint
+    - If src is local: source_session is ignored; opened via local Path
+    - If dst is remote: target_session is REQUIRED and must be bound to dst's endpoint
+    - If dst is local: target_session is ignored; opened via local Path
+
+    Raises ValueError if a remote endpoint lacks its required session.
+    Raises ValueError if a provided session's binding doesn't match the URI's endpoint.
+
+    Validation order: ALL session/binding checks complete BEFORE any file
+    handles are opened. This guarantees mismatched-binding failures are
+    surfaced loudly without leaking file handles or partial writes.
+    """
+    from firecube.core.uris import local_path_from_target
+
+    # PHASE 1: validate sessions+bindings BEFORE any I/O
+    if src.is_remote():
+        if source_session is None:
+            raise ValueError(f"remote source {src.to_str()} requires source_session")
+        _validate_session_binding(source_session, src, role="source")
+    if dst.is_remote():
+        if target_session is None:
+            raise ValueError(f"remote destination {dst.to_str()} requires target_session")
+        _validate_session_binding(target_session, dst, role="target")
+
+    # PHASE 2: open both handles (validation already passed)
+    if src.is_remote():
+        assert source_session is not None  # narrowed by Phase 1
+        src_handle = cast(Any, source_session.fs()).open(src, "rb")
+    else:
+        src_path = local_path_from_target(src.to_str())
+        src_handle = src_path.open("rb")
+
+    if dst.is_remote():
+        assert target_session is not None  # narrowed by Phase 1
+        dst_handle = cast(Any, target_session.fs()).open(dst, "wb")
+    else:
+        dst_path = local_path_from_target(dst.to_str())
+        dst_path.parent.mkdir(parents=True, exist_ok=True)
+        dst_handle = dst_path.open("wb")
+
+    # PHASE 3: copy
+    try:
+        shutil.copyfileobj(src_handle, dst_handle)
+    finally:
+        src_handle.close()
+        dst_handle.close()
+
+
+def _validate_session_binding(
+    session: StorageSession,
+    uri: StorageUri,
+    *,
+    role: str,
+) -> None:
+    """Ensure session is bound to the same protocol+authority as the URI."""
+    bound_uri = session.product.product_uri
+    if bound_uri.protocol != uri.protocol or bound_uri.authority != uri.authority:
+        raise ValueError(
+            f"{role}_session is bound to {bound_uri.protocol}://{bound_uri.authority}, "
+            f"but URI is {uri.protocol}://{uri.authority}"
+        )

firecube/core/storage/uri.py

@@ -0,0 +1,173 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Canonical storage URI value object.
+
+Canonical URI Contract:
+- Protocol is always lowercased. Supported protocols are: s3, gs, file, memory.
+- Authority is required for s3/gs (bucket) and absent for file/memory.
+- file://localhost/... normalizes localhost to absent authority.
+- Path always starts with /, collapses duplicate slashes, and strips trailing /
+  except for root (/).
+- Query strings and fragments are not supported.
+- Bare paths without URI schemes are rejected; use file:///... for local paths.
+- Percent-encoded characters are preserved as-is. Storage adapters own decoding.
+- to_str() emits canonical round-trippable strings:
+  file:///abs/path, memory:///test/path, s3://bucket/path, gs://bucket/path.
+"""
+
+from __future__ import annotations
+
+import posixpath
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from urllib.parse import urlparse
+
+SUPPORTED_PROTOCOLS = frozenset(
+    {
+        "s3",
+        # gs retained for planned future GCS support; not currently advertised in CLI choices
+        "gs",
+        "file",
+        "memory",
+    }
+)
+REMOTE_PROTOCOLS = frozenset({"s3", "gs"})
+LOCAL_PROTOCOLS = frozenset({"file", "memory"})
+
+
+def _validate_protocol(protocol: str) -> str:
+    normalized = str(protocol or "").strip().lower()
+    if not normalized:
+        raise ValueError("a URI scheme is required (use file:///… or s3://…)")
+    if normalized not in SUPPORTED_PROTOCOLS:
+        supported = ", ".join(sorted(SUPPORTED_PROTOCOLS))
+        raise ValueError(f"unsupported protocol: {normalized!r}; supported: {supported}")
+    return normalized
+
+
+def _normalize_path(path: str) -> str:
+    candidate = str(path or "")
+    if not candidate.startswith("/"):
+        candidate = f"/{candidate}"
+    candidate = re.sub(r"/+", "/", candidate)
+    collapsed = posixpath.normpath(candidate)
+    if collapsed == ".":
+        return "/"
+    return collapsed.rstrip("/") or "/"
+
+
+@dataclass(frozen=True, slots=True)
+class StorageUri:
+    protocol: str
+    authority: str | None
+    path: str
+    _root_had_slash: bool = field(default=True, init=False, repr=False, compare=False)
+
+    def __post_init__(self) -> None:
+        protocol = _validate_protocol(self.protocol)
+        authority = str(self.authority) if self.authority is not None else None
+        if authority == "":
+            authority = None
+        if protocol == "file" and authority == "localhost":
+            authority = None
+        path = _normalize_path(self.path)
+
+        if protocol in REMOTE_PROTOCOLS and authority is None:
+            raise ValueError("authority (bucket) is required for s3/gs URIs")
+        if protocol in LOCAL_PROTOCOLS and authority is not None:
+            raise ValueError(f"authority is not supported for {protocol} URIs")
+
+        object.__setattr__(self, "protocol", protocol)
+        object.__setattr__(self, "authority", authority)
+        object.__setattr__(self, "path", path)
+
+    @classmethod
+    def parse(cls, raw: str) -> StorageUri:
+        candidate = str(raw or "")
+        if not candidate.strip():
+            raise ValueError("StorageUri input must not be empty")
+
+        parsed = urlparse(candidate)
+        if not parsed.scheme:
+            if candidate.startswith("/"):
+                raise ValueError("bare paths are not supported; use file:///…")
+            raise ValueError("a URI scheme is required (use file:///… or s3://…)")
+        if parsed.query:
+            raise ValueError("query strings not supported in storage URIs")
+        if parsed.fragment:
+            raise ValueError("fragments not supported in storage URIs")
+
+        protocol = _validate_protocol(parsed.scheme)
+        authority = parsed.netloc or None
+        if protocol == "file" and authority == "localhost":
+            authority = None
+        if protocol in REMOTE_PROTOCOLS and authority is None:
+            raise ValueError("missing authority: authority (bucket) is required for s3/gs URIs")
+
+        uri = cls(protocol=protocol, authority=authority, path=parsed.path)
+        object.__setattr__(uri, "_root_had_slash", bool(parsed.path))
+        return uri
+
+    @classmethod
+    def from_local_path(cls, abs_path: str | Path) -> StorageUri:
+        path = Path(abs_path)
+        if not path.is_absolute():
+            raise ValueError("from_local_path requires an absolute path")
+        return cls(protocol="file", authority=None, path=path.as_posix())
+
+    def join(self, *segments: str) -> StorageUri:
+        cleaned_segments = [segment.strip("/") for segment in segments if segment.strip("/")]
+        if not cleaned_segments:
+            return self
+        joined_path = posixpath.join(self.path.rstrip("/"), *cleaned_segments)
+        return StorageUri(
+            protocol=self.protocol,
+            authority=self.authority,
+            path=joined_path,
+        )
+
+    def with_protocol(self, p: str) -> StorageUri:
+        protocol = _validate_protocol(p)
+        if protocol in REMOTE_PROTOCOLS and self.authority is None:
+            raise ValueError("authority (bucket) is required for s3/gs URIs")
+        if protocol in LOCAL_PROTOCOLS:
+            return StorageUri(protocol=protocol, authority=None, path=self.path)
+        return StorageUri(protocol=protocol, authority=self.authority, path=self.path)
+
+    def parent(self) -> StorageUri:
+        if self.path == "/":
+            return self
+        parent_path = posixpath.dirname(self.path)
+        parent = StorageUri(
+            protocol=self.protocol,
+            authority=self.authority,
+            path=parent_path or "/",
+        )
+        if parent.path == "/" and self.protocol in REMOTE_PROTOCOLS:
+            object.__setattr__(parent, "_root_had_slash", False)
+        return parent
+
+    def is_remote(self) -> bool:
+        return self.protocol in REMOTE_PROTOCOLS
+
+    def to_str(self) -> str:
+        if self.authority is None:
+            return f"{self.protocol}://{self.path}"
+        if self.path == "/":
+            if not self._root_had_slash:
+                return f"{self.protocol}://{self.authority}"
+            return f"{self.protocol}://{self.authority}/"
+        return f"{self.protocol}://{self.authority}{self.path}"

firecube/core/tensogram/__init__.py

@@ -0,0 +1,37 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tensogram integration for Firecube — optional dependency."""
+
+from firecube.core.tensogram.schema import (
+    ARCHIVE_VERSION,
+    KEY_ARCHIVE_VERSION,
+    KEY_GROUP,
+    KEY_ROLE,
+    ROLE_CONTROLPLANE,
+    ROLE_DATA,
+    make_controlplane_meta,
+    make_data_meta,
+)
+
+__all__ = [
+    "ARCHIVE_VERSION",
+    "KEY_ARCHIVE_VERSION",
+    "KEY_GROUP",
+    "KEY_ROLE",
+    "ROLE_CONTROLPLANE",
+    "ROLE_DATA",
+    "make_controlplane_meta",
+    "make_data_meta",
+]

firecube/core/tensogram/_compat.py

@@ -0,0 +1,36 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Optional tensogram dependency with graceful fallback."""
+
+from __future__ import annotations
+
+try:
+    import tensogram  # pyright: ignore[reportMissingImports]
+    import tensogram_xarray  # pyright: ignore[reportMissingImports]
+
+    HAS_TENSOGRAM = True
+except ImportError:
+    HAS_TENSOGRAM = False
+    tensogram = None  # type: ignore[assignment]
+    tensogram_xarray = None  # type: ignore[assignment]
+
+
+def require_tensogram(operation: str) -> None:
+    """Raise ImportError with install instructions if tensogram is not available."""
+    if not HAS_TENSOGRAM:
+        raise ImportError(
+            f"{operation} requires the tensogram extras. "
+            "Install with: uv pip install 'firecube[tensogram]'"
+        )

firecube/core/tensogram/controlplane_codec.py

@@ -0,0 +1,230 @@
+# Copyright 2025-2026 EUMETSAT
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Control-plane serialization codec for Firecube .tgm archives.
+
+Serializes .firecube/ state into a numpy uint8 array for embedding in .tgm
+archives, and deserializes/restores it on the restore path.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections import defaultdict
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from firecube.core.controlplane.manager import ChunkManager
+
+_log = logging.getLogger(__name__)
+
+
+def serialize_controlplane(
+    manager: ChunkManager,
+    product: str,
+    *,
+    group: str | None = None,
+) -> tuple[dict[str, Any], np.ndarray]:
+    """Serialize control-plane state for a product into a uint8 numpy array.
+
+    When *group* is specified only spans for that group are included.
+    Claims (write locks) are always included in full — they have no group field.
+    Run records are included only for runs referenced by the result set.
+
+    Returns
+    -------
+    tuple[dict, np.ndarray]
+        ``(descriptor_dict, uint8_array)`` suitable for ``TensogramFile.append()``.
+    """
+    list_kwargs: dict[str, Any] = {"product": product, "chunk_type": "span"}
+    if group is not None:
+        list_kwargs["meta"] = {"group": group}
+
+    spans = manager.list_chunks(**list_kwargs)
+
+    span_records: list[dict[str, Any]] = []
+    run_ids_seen: set[str] = set()
+    for chunk in spans:
+        span_meta = chunk.meta or {}
+        run_id = span_meta.get("run_id")
+        if run_id:
+            run_ids_seen.add(run_id)
+        span_records.append(
+            {
+                "key": chunk.key,
+                "type": chunk.chunk_type,
+                "status": chunk.status,
+                "meta": span_meta,
+                "record": chunk.record or {},
+            }
+        )
+
+    runs = manager.list_runs(product=product)
+    run_records: list[dict[str, Any]] = []
+    for run in runs:
+        if run_ids_seen and run.run_id not in run_ids_seen:
+            continue
+        run_records.append(
+            {
+                "run_id": run.run_id,
+                "status": run.status,
+                "started_at": run.started_at,
+                "updated_at": run.updated_at,
+                "completed_at": run.completed_at,
+                "events": run.events,
+                "parts": run.parts,
+                "product": run.product,
+            }
+        )
+
+    claim_records: list[dict[str, Any]] = []
+    try:
+        claims = manager.list_claims(product=product)
+        claim_records.extend(
+            {
+                "domain": claim.domain,
+                "owner_id": claim.owner_id,
+                "acquired_at": claim.acquired_at,
+            }
+            for claim in claims
+        )
+    except Exception as exc:
+        _log.debug("Could not collect claims for %s: %s", product, exc)
+
+    state: dict[str, Any] = {
+        "schema_version": "v1",
+        "product": product,
+        "group_filter": group,
+        "spans": span_records,
+        "runs": run_records,
+        "claims": claim_records,
+    }
+
+    data_bytes = json.dumps(state).encode("utf-8")
+    arr = np.frombuffer(data_bytes, dtype=np.uint8).copy()
+    descriptor: dict[str, Any] = {
+        "type": "ntensor",
+        "shape": [len(arr)],
+        "dtype": "uint8",
+        "compression": "none",
+    }
+    return descriptor, arr
+
+
+def deserialize_controlplane(data: np.ndarray) -> dict[str, Any]:
+    """Deserialize a control-plane uint8 array back into a state dict."""
+    data_bytes = data.astype(np.uint8).tobytes()
+    return json.loads(data_bytes.decode("utf-8"))
+
+
+def restore_controlplane(
+    state: dict[str, Any],
+    target_path: str,
+    *,
+    storage_config: Any | None = None,
+) -> None:
+    """Write .firecube/ control-plane state to a target path.
+
+    Replays run and span records as WAL events under
+    ``<target_path>/.firecube/`` using ``RunEventWriter``.
+
+    This is best-effort: individual run failures are logged as warnings
+    and do not prevent other runs from being restored.
+
+    Parameters
+    ----------
+    state:
+        Deserialized control-plane state dict from ``deserialize_controlplane()``.
+    target_path:
+        Root path of the zarr store to restore ``.firecube/`` into.
+    storage_config:
+        Optional ``StorageConfig`` used to wire the underlying filesystem with
+        the configured endpoint and credentials. When ``None``, ambient defaults
+        apply — required for remote ``target_path`` values such as ``s3://...``.
+    """
+    from firecube.core.controlplane.events import RunEventWriter
+    from firecube.core.filesystem import StorageFilesystemFull, create_filesystem
+    from firecube.core.product.identity import ProductIdentity
+    from firecube.core.storage.binding import StorageBinding
+    from firecube.core.storage.driver_config import StorageDriverConfig
+    from firecube.core.storage.uri import StorageUri
+
+    target_uri = (
+        StorageUri.parse(target_path)
+        if "://" in target_path
+        else StorageUri.from_local_path(Path(target_path).resolve())
+    )
+    control_uri = target_uri.join(".firecube")
+    driver = StorageDriverConfig.from_storage_config_or_default(storage_config)
+    fs = cast(
+        StorageFilesystemFull,
+        create_filesystem(
+            StorageBinding(
+                identity=ProductIdentity.from_uri(target_uri, "zarr", product_name=target_path),
+                driver=driver,
+            )
+        ),
+    )
+
+    product = state.get("product", "unknown")
+    spans = state.get("spans", [])
+    runs = state.get("runs", [])
+
+    if not spans and not runs:
+        _log.info("No control-plane records to restore for %s", product)
+        return
+
+    spans_by_run: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for span in spans:
+        run_id = span.get("meta", {}).get("run_id", "unknown")
+        spans_by_run[run_id].append(span)
+
+    runs_by_id = {r["run_id"]: r for r in runs}
+
+    for run_id, run_spans in spans_by_run.items():
+        run_info = runs_by_id.get(run_id, {})
+        try:
+            writer = RunEventWriter(
+                fs=fs,
+                control_uri=control_uri,
+                product=product,
+                run_id=run_id,
+                resume_meta=run_info if run_info else None,
+            )
+            for span in run_spans:
+                writer.append(
+                    event_type="span_recorded",
+                    record=span.get("record", {}),
+                    meta=span.get("meta", {}),
+                )
+            original_status = run_info.get("status", "complete")
+            writer.finalize(status=original_status)
+        except Exception as exc:
+            _log.warning(
+                "Could not restore WAL events for run %s in %s: %s",
+                run_id,
+                product,
+                exc,
+            )
+
+    _log.info(
+        "Restored .firecube/ state for %s: %d spans across %d runs",
+        product,
+        len(spans),
+        len(runs_by_id),
+    )