firecube 0.1.4__py3-none-any.whl

firecube/ingestor/runtime/zarr/staged_metadata.py

@@ -0,0 +1,249 @@
+# 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.
+
+"""Metadata seeding for staged-write mode.
+
+When firecube writes through a local temp store before publishing to
+the final target (staged mode), the temp store starts empty and has no
+Zarr metadata.  This module copies zarr.json files from the existing
+final target into the temp store so that ``append_time_groups()`` reads
+the correct write cursor and shape — preventing metadata corruption where
+temp-store shape overwrites the cumulative target shape.
+
+Best-effort for fresh ingests (no final target yet); strict (raises
+``StagedMetadataError``) when resuming into an existing target.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any
+
+from firecube.core.storage.uri import StorageUri
+from firecube.ingestor.errors import StagedMetadataError
+
+if TYPE_CHECKING:
+    from firecube.core.storage.session import StorageSession
+
+log = logging.getLogger(__name__)
+
+
+def seed_staged_store_metadata(
+    *,
+    temp_store_uri: str,
+    final_target_uri: str,
+    groups: list[str] | None = None,
+    strict: bool = True,
+    session: StorageSession,
+    coordinate_arrays: list[str] | None = None,
+) -> dict[str, Any]:
+    """Copy zarr.json metadata files (and optionally named coordinate-array chunks) from final_target_uri into temp_store_uri.
+
+    Seeds the temp store with correct shape/dimension metadata so that
+    append_time_groups() starts write_cursor at the right position. By
+    default copies only zarr.json files (metadata).
+
+    When ``coordinate_arrays`` is a non-empty list of coord-array names,
+    also copies the chunk payloads of those arrays in every discovered
+    group that contains them. Used by the runtime hook to seed the
+    workspace timestamp coordinate so that plugin value-based
+    timestamp resolvers see real values instead of NaT on staged
+    re-ingest. Coordinate arrays missing in a group are tolerated (no-op).
+    Never copies data-array chunks.
+
+    All metadata reads from the final target use the driver-correct
+    ``session.fs()`` filesystem; ``session`` is REQUIRED.
+
+    When ``strict=True`` (default), unexpected exceptions are wrapped as
+    ``StagedMetadataError`` so callers that catch ``StagedMetadataError``
+    propagate hard failures while ``except Exception`` debug-suppressors do
+    not silently swallow them.
+
+    ``coordinate_arrays=None`` and ``coordinate_arrays=[]`` are
+    semantically identical and both preserve exact legacy metadata-only
+    behavior.
+
+    Returns {group: {"seeded": bool, "files": int}} for logging.
+    """
+    if groups is None:
+        groups = _discover_zarr_groups(session, final_target_uri)
+
+    results: dict[str, Any] = {}
+
+    for group in groups:
+        results[group] = {"seeded": False, "files": 0}
+        try:
+            _seed_group_via_session(
+                session=session,
+                temp_store_uri=temp_store_uri,
+                final_target_uri=final_target_uri,
+                group=group,
+                strict=strict,
+                results=results,
+                coordinate_arrays=coordinate_arrays,
+            )
+
+            if results[group]["files"] > 0:
+                results[group]["seeded"] = True
+                log.debug(
+                    "Seeded %d %s for group %s",
+                    results[group]["files"],
+                    "metadata + coord files" if coordinate_arrays else "zarr.json files",
+                    group,
+                )
+        except StagedMetadataError:
+            raise
+        except Exception as e:
+            if strict:
+                raise StagedMetadataError(
+                    f"Staged metadata seeding failed unexpectedly for group {group}: {e}"
+                ) from e
+            log.warning("Staged metadata seeding skipped for group %s: %s", group, e)
+
+    return results
+
+
+def _discover_zarr_groups(session: StorageSession, uri: str) -> list[str]:
+    """Discover top-level Zarr groups under ``uri`` from existing metadata."""
+    from firecube.core.uris import is_remote_target, local_path_from_target
+
+    fs = session.fs()
+    if is_remote_target(uri):
+        final_uri = StorageUri.parse(uri)
+    else:
+        final_uri = StorageUri.from_local_path(local_path_from_target(uri))
+
+    if not fs.exists(final_uri):
+        return []
+
+    final_root = final_uri.path.rstrip("/")
+    groups: set[str] = set()
+    for entry in fs.find(final_uri):
+        entry_path = entry.path.rstrip("/")
+        if not (entry_path.endswith("/zarr.json") or entry_path.endswith("/.zgroup")):
+            continue
+        rel = (
+            entry_path[len(final_root) + 1 :]
+            if entry_path.startswith(f"{final_root}/")
+            else entry_path.lstrip("/")
+        )
+        parent = rel.rsplit("/", 1)[0] if "/" in rel else ""
+        if not parent:
+            continue
+        groups.add(parent.split("/", 1)[0])
+
+    return sorted(groups)
+
+
+def _is_coordinate_chunk(path: str, coords: list[str] | None) -> bool:
+    """Return True if ``path`` is a chunk file under a named coord array.
+
+    Uses path-segment matching: splits on "/" and looks for the exact
+    sequence ``<coord>/c`` where ``<coord>`` is in ``coords``. Returns
+    False for the coord's ``zarr.json`` (already covered by the
+    metadata filter) and for arrays whose names contain a coord name
+    as a substring (e.g. ``timestamp_bounds`` for ``["timestamp"]``).
+
+    Args:
+        path: The entry path relative to or under the group root.
+              Trailing slashes are stripped.
+        coords: Coordinate array names to match (exact path segments).
+                ``None`` and ``[]`` always return False (legacy mode).
+
+    Examples:
+        _is_coordinate_chunk("data/timestamp/c/0", ["timestamp"]) -> True
+        _is_coordinate_chunk("data/timestamp/c/0/0", ["timestamp"]) -> True
+        _is_coordinate_chunk("data/timestamp/zarr.json", ["timestamp"]) -> False
+        _is_coordinate_chunk("data/timestamp_bounds/c/0", ["timestamp"]) -> False
+        _is_coordinate_chunk("data/val/c/0", ["timestamp"]) -> False
+        _is_coordinate_chunk("data/timestamp/c/0", []) -> False
+        _is_coordinate_chunk("data/timestamp/c/0", None) -> False
+    """
+    if not coords:
+        return False
+    parts = path.rstrip("/").split("/")
+    return any(parts[i] in coords and parts[i + 1] == "c" for i in range(len(parts) - 2))
+
+
+def _seed_group_via_session(
+    *,
+    session: StorageSession,
+    temp_store_uri: str,
+    final_target_uri: str,
+    group: str,
+    strict: bool,
+    results: dict[str, Any],
+    coordinate_arrays: list[str] | None = None,
+) -> None:
+    from firecube.core.uris import is_remote_target, local_path_from_target
+
+    fs = session.fs()
+    if is_remote_target(final_target_uri):
+        final_uri = StorageUri.parse(final_target_uri)
+    else:
+        final_uri = StorageUri.from_local_path(local_path_from_target(final_target_uri))
+    group_uri = final_uri.join(group.strip("/"))
+
+    if not fs.exists(group_uri):
+        return  # Fresh ingest — no final target yet
+
+    all_entries = fs.find(group_uri)
+    final_root = final_uri.path.rstrip("/")
+    files_to_copy = [
+        entry
+        for entry in all_entries
+        if entry.path.endswith("zarr.json")
+        or _is_coordinate_chunk(
+            entry.path[len(final_root) + 1 :]
+            if entry.path.startswith(f"{final_root}/")
+            else entry.path.lstrip("/"),
+            coordinate_arrays,
+        )
+    ]
+    if not files_to_copy:
+        return
+
+    temp_root = local_path_from_target(temp_store_uri)
+    for src_uri in files_to_copy:
+        rel = (
+            src_uri.path[len(final_root) + 1 :]
+            if src_uri.path.startswith(f"{final_root}/")
+            else src_uri.path.lstrip("/")
+        )
+        dst_path = temp_root / rel
+        if dst_path.exists():
+            results[group]["files"] += 1
+            continue
+        try:
+            dst_path.parent.mkdir(parents=True, exist_ok=True)
+            with fs.open(src_uri, "rb") as fh:
+                raw = fh.read()
+            if rel.endswith("zarr.json"):
+                payload = json.loads(raw)
+                if payload.get("node_type") == "array":
+                    attrs = payload.get("attributes")
+                    if isinstance(attrs, dict):
+                        # If a second write-once reserved attr is added later, refactor _reserved_attrs.py into RESERVED_SHAPE_ATTRS (kept on seed) vs RESERVED_RUN_STATE_ATTRS (stripped on seed), and switch this strip from a literal key to a set membership check. YAGNI today (only one such attr exists).
+                        attrs.pop("firecube_static_written", None)
+                dst_path.write_bytes(json.dumps(payload).encode())
+            else:
+                dst_path.write_bytes(raw)
+            results[group]["files"] += 1
+        except Exception as e:
+            if strict:
+                raise StagedMetadataError(
+                    f"Failed to seed staged metadata {src_uri.to_str()} -> {dst_path}: {e}"
+                ) from e
+            log.warning("Failed to seed staged metadata %s: %s", src_uri.to_str(), e)

firecube/ingestor/runtime/zarr/strategies/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+"""Concrete Zarr write strategy implementations."""

firecube/ingestor/runtime/zarr/strategies/append.py

@@ -0,0 +1,172 @@
+# 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.
+
+"""AppendStrategy — AppendWriteStrategy wrapping ``append_time_groups()``.
+
+Delegates entirely to the existing orchestrator in
+``firecube.ingestor.runtime.zarr.append``.  No append logic is reimplemented.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, cast
+
+import xarray as xr
+
+from firecube.core.uris import storage_uri_from_target
+
+if TYPE_CHECKING:
+    from firecube.core.config import StorageConfig
+    from firecube.core.filesystem.store_factory import ZarrStoreHandle
+    from firecube.core.storage.session import StorageSession
+
+
+def _session_for_store(store_uri: str, storage_config: StorageConfig) -> StorageSession:
+    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
+
+    uri = storage_uri_from_target(store_uri)
+    return StorageSession(
+        StorageBinding(
+            identity=ProductIdentity.from_uri(uri, format="zarr", product_name=store_uri),
+            driver=StorageDriverConfig.from_storage_config(storage_config),
+        )
+    )
+
+
+class AppendStrategy:
+    """Xarray-append write strategy for Zarr stores.
+
+    Wraps :func:`~firecube.ingestor.runtime.zarr.append.append_time_groups`
+    behind the ``AppendWriteStrategy`` Protocol so that
+    ``GenericZarrIngestor`` can treat append and region writes uniformly.
+
+    All static configuration is captured at construction time; only
+    per-batch dynamic inputs are passed to :meth:`write_groups`.
+    """
+
+    def __init__(
+        self,
+        *,
+        store: object,
+        store_uri: str | None = None,
+        store_handle: ZarrStoreHandle | None = None,
+        resume_target_uri: str | None = None,
+        arrays_for_group: Callable[[str], list[str]] | None = None,
+        chunk_shape: dict[str, int] | None = None,
+        shard_shape: dict[str, int] | None = None,
+        sharding: bool = False,
+        compression: bool | str = False,
+        consolidate: bool = False,
+        resume_existing: bool = False,
+        append_dim: str = "timestamp",
+        state_var_name: str = "firecube_timestamp_state",
+        state_deleted_value: int = 2,
+        logger: logging.Logger | None = None,
+        storage_config: StorageConfig | None = None,
+        session: StorageSession | None = None,
+    ) -> None:
+        # NOTE: append_dim is the already-resolved time dim name supplied by
+        # the caller (typically GenericZarrIngestor via _resolve_time_dim_name()).
+        # The strategy is host-free at runtime — never looks up the dim name
+        # from an ingestor object. This keeps the write protocol decoupled and
+        # the strategy testable in isolation.
+        self._store = store
+        self._store_uri = store_uri
+        self._store_handle = store_handle
+        self._resume_target_uri = resume_target_uri
+        self._storage_config = storage_config
+        self._arrays_for_group = arrays_for_group
+        self._chunk_shape = chunk_shape
+        self._shard_shape = shard_shape
+        self._sharding = sharding
+        self._compression = compression
+        self._consolidate = consolidate
+        self._resume_existing = resume_existing
+        self._append_dim = append_dim
+        self._state_var_name = state_var_name
+        self._state_deleted_value = state_deleted_value
+        self._logger = logger
+        self._session = session
+
+    def write_groups(
+        self,
+        *,
+        group_to_timestamps: Mapping[str, Sequence[Any]],
+        dataset_for_batch: Callable[[str, Sequence[Any]], xr.Dataset | None],
+        batch_size: int,
+        claim_for_group: Callable[[str], Any] | None = None,
+    ) -> dict[str, Any]:
+        """Delegate to ``append_time_groups()`` and return its metrics dict."""
+        from firecube.ingestor.runtime.zarr.append import append_time_groups
+
+        effective_store = self._store
+        effective_zarr_store: Any = self._store_handle
+        effective_session: StorageSession | None = self._session
+
+        if effective_zarr_store is None and self._storage_config and self._store_uri:
+            if effective_session is None:
+                effective_session = _session_for_store(self._store_uri, self._storage_config)
+            effective_zarr_store = effective_session.zarr.create_store(
+                uri=storage_uri_from_target(self._store_uri),
+                mode="a",
+            )
+
+        if effective_zarr_store is not None:
+            typed_handle = cast(Any, effective_zarr_store)
+            effective_store = typed_handle.store
+
+        if effective_zarr_store is None:
+            raise ValueError(
+                "AppendStrategy.write_groups requires either store_handle or "
+                "(storage_config + store_uri) to construct a ZarrStoreHandle."
+            )
+
+        resume_zarr_store: Any = None
+        resume_session: Any = None
+        if self._resume_target_uri and self._storage_config:
+            resume_session = _session_for_store(self._resume_target_uri, self._storage_config)
+            resume_zarr_store = resume_session.zarr.create_store(
+                uri=storage_uri_from_target(self._resume_target_uri),
+                mode="r",
+            )
+
+        append_time_groups_fn = cast(Any, append_time_groups)
+        append_dim = self._append_dim
+        return append_time_groups_fn(
+            store=effective_store,
+            zarr_store=effective_zarr_store,
+            group_to_timestamps=group_to_timestamps,
+            dataset_for_batch=dataset_for_batch,
+            resume_zarr_store=resume_zarr_store,
+            resume_session=resume_session,
+            arrays_for_group=self._arrays_for_group,
+            chunk_shape=self._chunk_shape,
+            shard_shape=self._shard_shape,
+            sharding=self._sharding,
+            compression=self._compression,
+            consolidate=self._consolidate,
+            resume_existing=self._resume_existing,
+            batch_size=batch_size,
+            append_dim=append_dim,
+            state_var_name=self._state_var_name,
+            state_deleted_value=self._state_deleted_value,
+            logger=self._logger,
+            claim_for_group=claim_for_group,
+            session=effective_session,
+        )