firecube 0.1.4__py3-none-any.whl

firecube/core/tensogram/restore.py

@@ -0,0 +1,303 @@
+# 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.
+
+"""Core Tensogram (.tgm) -> Zarr restore."""
+
+from __future__ import annotations
+
+import logging
+import os
+import shutil
+import tempfile
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, cast
+
+import xarray as xr
+
+from firecube.core.config import StorageConfig
+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.session import StorageSession
+from firecube.core.uris import is_remote_target, local_path_from_target, parse_target
+
+from ._compat import require_tensogram
+from .controlplane_codec import deserialize_controlplane, restore_controlplane
+from .schema import ARCHIVE_VERSION, KEY_ARCHIVE_VERSION, ROLE_CONTROLPLANE, ROLE_DATA
+
+logger = logging.getLogger(__name__)
+
+
+def _make_archive_session(target: str, storage_config: Any | None) -> StorageSession:
+    target_uri = parse_target(target)
+    if storage_config is None:
+        target_path = Path(target_uri.path).resolve()
+        storage_config = StorageConfig(storage_type="local")
+        storage_config.target_path = str(target_path.parent)  # type: ignore[attr-defined]
+
+    driver_config = StorageDriverConfig.from_storage_config(storage_config)
+    return StorageSession(
+        StorageBinding(
+            identity=ProductIdentity.from_uri(target_uri, "zarr", product_name=target),
+            driver=driver_config,
+        )
+    )
+
+
+def _prepare_target(target: str, *, overwrite: bool) -> None:
+    if is_remote_target(target):
+        return
+
+    target_uri = parse_target(target)
+    local_target_path = Path(target_uri.path).resolve()
+    if not local_target_path.exists():
+        return
+
+    # A directory that contains ONLY .firecube/ (control-plane state from claim
+    # acquisition) is not a Zarr store yet — do not treat it as an existing store.
+    # Only raise if the directory contains actual Zarr metadata (zarr.json or .zarray).
+    if local_target_path.is_dir():
+        has_zarr_metadata = (local_target_path / "zarr.json").exists() or (
+            local_target_path / ".zarray"
+        ).exists()
+        if not has_zarr_metadata:
+            return
+
+    if not overwrite:
+        raise FileExistsError(
+            f"Target Zarr already exists: {target}. Use overwrite=True to replace it."
+        )
+
+    if local_target_path.is_dir():
+        shutil.rmtree(local_target_path)
+    else:
+        local_target_path.unlink()
+
+
+def _reconstruct_compressor(config: dict[str, Any] | None) -> Any | None:
+    """Reconstruct a numcodecs compressor from a config dict."""
+    if config is None:
+        return None
+
+    try:
+        import numcodecs
+
+        return numcodecs.get_codec(config)
+    except Exception:
+        logger.debug("Could not reconstruct compressor from config: %s", config, exc_info=True)
+        return None
+
+
+def _build_encoding(base_entries: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
+    encoding: dict[str, dict[str, Any]] = {}
+    for base_entry in base_entries:
+        var_name = base_entry.get("name")
+        if not var_name:
+            continue
+
+        variable_encoding: dict[str, Any] = {}
+        zarr_chunks = base_entry.get("zarr_chunks")
+        if zarr_chunks:
+            variable_encoding["chunks"] = tuple(zarr_chunks)
+
+        if "zarr_compressor" in base_entry:
+            variable_encoding["compressor"] = _reconstruct_compressor(
+                cast(dict[str, Any] | None, base_entry.get("zarr_compressor"))
+            )
+
+        if "zarr_fill_value" in base_entry:
+            variable_encoding["fill_value"] = base_entry.get("zarr_fill_value")
+
+        if variable_encoding:
+            encoding[str(var_name)] = variable_encoding
+
+    return encoding
+
+
+def _open_message_dataset(
+    tensogram_file: Any,
+    message_index: int,
+    *,
+    coord_names: list[str],
+) -> tuple[xr.Dataset, str]:
+    raw_message = tensogram_file.read_message(message_index)
+    with tempfile.NamedTemporaryFile(suffix=".tgm", delete=False) as tmp:
+        tmp.write(raw_message)
+        tmp_path = tmp.name
+
+    ds = xr.open_dataset(tmp_path, engine="tensogram")
+    vars_to_promote = [coord_name for coord_name in coord_names if coord_name in ds.data_vars]
+    if vars_to_promote:
+        ds = ds.set_coords(vars_to_promote)
+    return ds, tmp_path
+
+
+def _restore_data_message(
+    tensogram_file: Any,
+    message_index: int,
+    *,
+    target: str,
+    group_name: str | None,
+    session: StorageSession,
+    first_group: bool,
+) -> dict[str, Any]:
+    msg_meta = tensogram_file.file_decode_metadata(message_index)
+    base_entries = list(msg_meta.base or [])
+    extra = (msg_meta.extra or {}) if hasattr(msg_meta, "extra") else {}
+    firecube_meta = extra.get("firecube", {})
+    coord_names = list(cast(list[str], firecube_meta.get("coordinates", [])))
+    encoding = _build_encoding(base_entries)
+
+    ds, tmp_path = _open_message_dataset(
+        tensogram_file,
+        message_index,
+        coord_names=coord_names,
+    )
+    try:
+        mode = "w" if first_group else "a"
+
+        valid_names = set(ds.data_vars) | set(ds.coords)
+        encoding = {k: v for k, v in encoding.items() if k in valid_names}
+        session.zarr.write_dataset(
+            ds,
+            parse_target(target),
+            mode=mode,
+            group=group_name,
+            encoding=encoding,
+        )
+        return {
+            "group": group_name,
+            "variables": list(ds.data_vars),
+            "coordinates": list(ds.coords),
+        }
+    finally:
+        ds.close()
+        os.unlink(tmp_path)
+
+
+def tgm_to_zarr(
+    source: str,
+    target: str,
+    *,
+    overwrite: bool = False,
+    storage_config: Any | None = None,
+    session: StorageSession | None = None,
+    on_group_restored: Callable[[str | None], None] | None = None,
+) -> dict[str, Any]:
+    """Restore a Tensogram .tgm archive to a Zarr store.
+
+    ``source`` and ``target`` accept either a bare absolute path or a
+    canonical URI (e.g. ``file:///…``, ``s3://…``).
+
+    Args:
+        source: Path to .tgm file (local only for now).
+        target: Output Zarr store path (local or s3://).
+        overwrite: If True, overwrite an existing Zarr store.
+        storage_config: Optional storage config used to derive remote
+            ``storage_options``.
+
+    Returns:
+        Summary dict describing the restore outcome.
+
+    Raises:
+        ImportError: If tensogram extras are not installed.
+        FileExistsError: If target exists and overwrite is False.
+    """
+    require_tensogram("tgm_to_zarr")
+    import tensogram as _tensogram  # pyright: ignore[reportMissingImports]
+
+    tensogram = cast(Any, _tensogram)
+    _prepare_target(target, overwrite=overwrite)
+    session = session or _make_archive_session(target, storage_config)
+
+    source_path = str(local_path_from_target(source))
+
+    with tensogram.TensogramFile.open(source_path) as tensogram_file:  # pyright: ignore[reportAttributeAccessIssue]
+        count = tensogram_file.message_count()
+        if count == 0:
+            return {
+                "source": source,
+                "target": target,
+                "group": None,
+                "groups": [],
+                "variables": [],
+                "coordinates": [],
+                "restored": 0,
+                "messages_restored": 0,
+                "controlplane_restored": False,
+            }
+
+        first_meta = tensogram_file.file_decode_metadata(0)
+        first_extra = (first_meta.extra or {}) if hasattr(first_meta, "extra") else {}
+        firecube_meta = first_extra.get("firecube", {})
+        archive_version = firecube_meta.get(KEY_ARCHIVE_VERSION)
+        if archive_version != ARCHIVE_VERSION:
+            raise ValueError(
+                f"Unsupported Firecube archive version: {archive_version!r}. "
+                f"Expected {ARCHIVE_VERSION!r}."
+            )
+
+        restored_groups: list[str] = []
+        restored_variables: list[str] = []
+        restored_coordinates: list[str] = []
+        controlplane_restored = False
+        first_group = True
+
+        for i in range(count):
+            msg_meta = tensogram_file.file_decode_metadata(i)
+            extra = (msg_meta.extra or {}) if hasattr(msg_meta, "extra") else {}
+            firecube_meta = extra.get("firecube", {})
+            role = firecube_meta.get("role", ROLE_DATA)
+
+            if role == ROLE_CONTROLPLANE:
+                msg = tensogram_file.decode_message(i)
+                if not msg.objects:
+                    continue
+                cp_array = msg.objects[0][1]
+                cp_state = deserialize_controlplane(cp_array)
+                restore_controlplane(cp_state, target, storage_config=storage_config)
+                controlplane_restored = True
+                continue
+
+            group_name = cast(str | None, firecube_meta.get("group") or None)
+            restored = _restore_data_message(
+                tensogram_file,
+                i,
+                target=target,
+                group_name=group_name,
+                session=session,
+                first_group=first_group,
+            )
+            first_group = False
+
+            if restored["group"] is not None:
+                restored_groups.append(cast(str, restored["group"]))
+            restored_variables.extend(cast(list[str], restored["variables"]))
+            restored_coordinates.extend(cast(list[str], restored["coordinates"]))
+
+            if on_group_restored is not None:
+                on_group_restored(cast(str | None, restored["group"]))
+
+        return {
+            "source": source,
+            "target": target,
+            "group": None,
+            "groups": restored_groups,
+            "variables": sorted(set(restored_variables)),
+            "coordinates": sorted(set(restored_coordinates)),
+            "restored": len(restored_groups),
+            "messages_restored": len(restored_groups),
+            "controlplane_restored": controlplane_restored,
+        }

firecube/core/tensogram/schema.py

@@ -0,0 +1,93 @@
+# 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.
+
+"""Archive schema constants and metadata helpers for Firecube .tgm archives."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import Any
+
+import firecube
+
+# Archive format version.
+ARCHIVE_VERSION: str = "v1"
+
+# Message role constants (stored in firecube.role within message metadata)
+ROLE_DATA: str = "data"
+ROLE_CONTROLPLANE: str = "controlplane"
+
+# Metadata key names (for documentation / avoiding magic strings)
+KEY_GROUP: str = "group"
+KEY_ROLE: str = "role"
+KEY_ARCHIVE_VERSION: str = "archive_version"
+
+
+def make_data_meta(
+    group: str,
+    base: list[dict[str, Any]],
+    *,
+    source_uri: str,
+    compression: str,
+    coordinates: list[str] | None = None,
+    start_date: str | None = None,
+    end_date: str | None = None,
+    extra_attrs: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Build a complete metadata dict for a data message.
+
+    Returns a dict suitable for passing to TensogramFile.append().
+    """
+    # Tensogram >=0.18 metadata frames are free-form; the wire version lives in
+    # the file preamble, not here. The legacy top-level "version" key (the
+    # removed GlobalMetadata.version field from tensogram <0.18) is not written.
+    meta: dict[str, Any] = {
+        "base": base,
+        "firecube": {
+            "version": getattr(firecube, "__version__", "unknown"),
+            "archive_version": ARCHIVE_VERSION,
+            "role": ROLE_DATA,
+            "group": group,
+            "source_uri": source_uri,
+            "archived_at": datetime.now(UTC).isoformat(),
+            "compression": compression,
+        },
+    }
+    if coordinates is not None:
+        meta["firecube"]["coordinates"] = coordinates
+    if start_date is not None:
+        meta["firecube"]["start_date"] = start_date
+    if end_date is not None:
+        meta["firecube"]["end_date"] = end_date
+    if extra_attrs:
+        meta.update(extra_attrs)
+    return meta
+
+
+def make_controlplane_meta(product: str) -> dict[str, Any]:
+    """Build a metadata dict for a control-plane message.
+
+    Returns a dict suitable for passing to TensogramFile.append().
+    The control-plane message has no group — it covers all groups for the product.
+    """
+    return {
+        "base": [{"name": "controlplane", "dim_names": ["bytes"]}],
+        "firecube": {
+            "version": getattr(firecube, "__version__", "unknown"),
+            "archive_version": ARCHIVE_VERSION,
+            "role": ROLE_CONTROLPLANE,
+            "product": product,
+            "archived_at": datetime.now(UTC).isoformat(),
+        },
+    }

firecube/core/uris.py

@@ -0,0 +1,182 @@
+# 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.
+
+"""URI/path helpers used across core, CLI, and plugins.
+
+This module centralizes protocol detection and local path normalization so we
+don't duplicate `fsspec.utils.infer_storage_options(...)` logic in multiple
+places (storage, filesystem, CLI).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from fsspec.utils import infer_storage_options
+
+if TYPE_CHECKING:
+    from firecube.core.storage.uri import StorageUri
+
+_log = logging.getLogger(__name__)
+_URI_SCHEME_RE = re.compile(r"^[a-zA-Z][a-zA-Z0-9+\-.]*[:]{0,1}[/]{1,2}")
+_WINDOWS_ABSOLUTE_PATH_RE = re.compile(r"^[a-zA-Z]:[/\\]")
+_MALFORMED_URI_SCHEMELESS_RE = re.compile(r"^([a-zA-Z][a-zA-Z0-9+\-.]*)//")
+_COMMON_URI_SCHEMES = {
+    "abfs",
+    "adl",
+    "az",
+    "file",
+    "ftp",
+    "gcs",
+    "gs",
+    "hdfs",
+    "http",
+    "https",
+    "memory",
+    "s3",
+    "sftp",
+}
+
+
+def _looks_like_uri(target: str) -> bool:
+    """Return True if target appears to be an intentional URI scheme (even if malformed)."""
+    return bool(_URI_SCHEME_RE.match(target))
+
+
+def _looks_like_windows_absolute_path(target: str) -> bool:
+    """Return True for Windows drive-letter absolute paths on any host OS."""
+    return bool(_WINDOWS_ABSOLUTE_PATH_RE.match(target))
+
+
+def _is_malformed_uri_candidate(target: str, protocol: str, path: str) -> bool:
+    """Return True when a URI-like target resolved only as a plain local path."""
+    if protocol != "file" or path != target or _looks_like_windows_absolute_path(target):
+        return False
+
+    if target.startswith("://") or "://" in target:
+        return True
+
+    if re.match(r"^[a-zA-Z][a-zA-Z0-9+\-.]*:/", target):
+        return True
+
+    schemeless_match = _MALFORMED_URI_SCHEMELESS_RE.match(target)
+    return bool(schemeless_match and schemeless_match.group(1).lower() in _COMMON_URI_SCHEMES)
+
+
+def parse_uri(target: str) -> dict[str, str]:
+    """Parse a URI/path into a minimal fsspec-like structure.
+
+    Returns a dict with:
+      - protocol: "file", "s3", ...
+      - path: protocol-specific path (e.g. "bucket/prefix/key" for s3, "/abs/path" for file)
+    """
+    target = str(target or "")
+    if target.startswith("file://"):
+        # Treat file:// as local. fsspec parsing will give us an absolute path if possible.
+        try:
+            parsed = infer_storage_options(target)
+        except Exception:
+            if _looks_like_uri(target):
+                raise ValueError(f"Malformed URI: {target!r}") from None
+            return {"protocol": "file", "path": target.replace("file://", "", 1)}
+    else:
+        try:
+            parsed = infer_storage_options(target)
+        except Exception:
+            if _looks_like_uri(target):
+                raise ValueError(f"Malformed URI: {target!r}") from None
+            return {"protocol": "file", "path": target}
+
+    protocol = parsed.get("protocol") or "file"
+    if isinstance(protocol, (list, tuple)):
+        protocol = protocol[0] if protocol else "file"
+    path = str(parsed.get("path") or "")
+    if _is_malformed_uri_candidate(target, str(protocol), path):
+        raise ValueError(f"Malformed URI: {target!r}")
+    _log.debug("parse_uri(%r) -> protocol=%s path=%s", target, protocol, path)
+    return {"protocol": str(protocol), "path": path}
+
+
+def infer_target_protocol(target: str) -> str:
+    """Infer the fsspec protocol for a target string.
+
+    Returns "file" for local paths and file:// URIs.
+    """
+    return parse_uri(target)["protocol"]
+
+
+def is_remote_target(target: str) -> bool:
+    """Return True when `target` is a non-local URI (e.g. s3://...)."""
+    protocol = infer_target_protocol(target)
+    return bool(protocol and protocol != "file")
+
+
+def local_path_from_target(target: str) -> Path:
+    """Resolve a target string into an absolute local path.
+
+    Supports relative paths and file:// URIs.
+    """
+    target = str(target or "")
+    if target.startswith("file://"):
+        from firecube.core.storage.uri import StorageUri
+
+        uri = StorageUri.parse(target)
+        return Path(uri.path)
+    p = Path(target)  # firecube: STORAGE-URI — intentional: this IS the local-path adapter
+    return p if p.is_absolute() else (Path.cwd() / p)
+
+
+def storage_uri_from_target(target: str) -> StorageUri:
+    from firecube.core.storage.uri import StorageUri  # local import to avoid circular
+
+    if is_remote_target(target) or str(target).startswith("file://"):
+        return StorageUri.parse(target)
+    return StorageUri.from_local_path(local_path_from_target(target))
+
+
+def is_windows_absolute_path(target: str) -> bool:
+    """Best-effort check for Windows drive-letter absolute paths."""
+    target = str(target or "")
+    return bool(os.name == "nt" and ":" in target)
+
+
+def parse_target(target: StorageUri | str) -> StorageUri:
+    """Parse a target string tolerant of bare local paths.
+
+    Accepts:
+      - StorageUri: returned as-is
+      - file:///… URI: parsed strictly via StorageUri.parse
+      - s3:// gs:// memory:// URIs: parsed strictly via StorageUri.parse
+      - Bare absolute local path (/abs/path): coerced via StorageUri.from_local_path
+      - Bare relative path: rejected (use file:/// or absolute path)
+
+    This is the tolerant boundary helper for public APIs that accept
+    user-supplied target strings (CLI args, plugin entry points,
+    legacy-shaped function signatures). Internal code should use
+    StorageUri.parse directly (strict).
+    """
+    from firecube.core.storage.uri import StorageUri  # local import to avoid circular
+
+    if isinstance(target, StorageUri):
+        return target
+    s = str(target)
+    if "://" in s:
+        return StorageUri.parse(s)
+    if s.startswith("/"):
+        return StorageUri.from_local_path(s)
+    raise ValueError(f"target must be an absolute path or URI: {target!r}")

firecube/core/workspaces.py

@@ -0,0 +1,68 @@
+# 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.
+
+"""Generic runtime helpers for ingestion plugins.
+
+These helpers are intentionally optional. Plugins that need a workspace and/or
+temporary directory root can use them to avoid duplicating path wiring logic.
+"""
+
+from __future__ import annotations
+
+import tempfile
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+
+def resolve_workspace_root(
+    options: Mapping[str, Any],
+    *,
+    prefix: str,
+    option_key: str = "workspace",
+) -> tuple[Path, bool]:
+    """Return (workspace_root, autocreated).
+
+    If `options[option_key]` is present, the workspace is treated as
+    user-provided (autocreated=False). Otherwise, a temporary directory is
+    created with the given prefix.
+    """
+    user_value = options.get(option_key)
+    if user_value:
+        return Path(user_value), False
+    return Path(tempfile.mkdtemp(prefix=f"{prefix}_")), True
+
+
+def ensure_workspace_dirs(
+    workspace_root: Path,
+    *,
+    temp_subdir: str = "tmp",
+) -> tuple[Path, Path]:
+    """Ensure workspace directory exists and return (workspace_root, temp_root)."""
+    workspace_root.mkdir(parents=True, exist_ok=True)
+    temp_root = workspace_root / temp_subdir
+    temp_root.mkdir(parents=True, exist_ok=True)
+    return workspace_root, temp_root
+
+
+def ensure_batch_workspace(
+    workspace_root: Path,
+    *,
+    batch_id: str,
+    prefix: str = "pipeline",
+    temp_subdir: str = "tmp",
+) -> tuple[Path, Path]:
+    """Return (batch_workspace, temp_root) as a subdirectory under workspace_root."""
+    batch_workspace = workspace_root / f"{prefix}_{batch_id}"
+    return ensure_workspace_dirs(batch_workspace, temp_subdir=temp_subdir)