diff --git a/tests/test_30_engine_ecmwf.py b/tests/test_30_engine_ecmwf.py new file mode 100644 index 0000000..cf17557 --- /dev/null +++ b/tests/test_30_engine_ecmwf.py @@ -0,0 +1,90 @@ +import os +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import xarray as xr + +from xarray_ecmwf import engine_ecmwf + +# several processes can build the same .zarr schema sample at once: each one +# writes its own copy under a temporary name, and the first copy renamed into +# place is the one everybody opens. +REQUEST = {"dataset": "dummy-dataset", "variable": ["dummy_variable"]} + + +# every download of the request gets the same file name, as with CDS +class DummyClient: + def __init__(self, client_kwargs: dict[str, Any]) -> None: + pass + + def submit_and_wait_on_result(self, request: dict[str, Any]) -> Any: + return request + + def get_filename(self, result: Any) -> str: + return "request.grib" + + def download(self, result: Any, target: str | None = None) -> str: + assert target is not None + Path(target).touch() + return target + + +def open_dummy(path: str, **kwargs: Any) -> xr.Dataset: + return xr.Dataset({"t": ("time", np.zeros(3))}, coords={"time": [0, 1, 2]}) + + +def open_sample(cache_folder: Path) -> xr.Dataset: + cacher = engine_ecmwf.DatasetCacher( + DummyClient({}), open_dummy, cache_folder=str(cache_folder) + ) + with cacher.cached_empty_dataset(REQUEST) as ds: + return ds + + +def test_cached_empty_dataset_keeps_the_sample_saved_first( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + to_zarr = xr.Dataset.to_zarr + + def to_zarr_after_another_process( + self: xr.Dataset, store: Any, *args: Any, **kw: Any + ) -> Any: + sample = os.path.splitext(store)[0] + # another process saves its complete sample while this one is writing + to_zarr(self.assign_attrs(saved_by="another process"), sample, *args, **kw) + return to_zarr(self, store, *args, **kw) + + monkeypatch.setattr(xr.Dataset, "to_zarr", to_zarr_after_another_process) + + ds = open_sample(tmp_path) + + assert ds.attrs["saved_by"] == "another process" + assert len([name for name in os.listdir(tmp_path) if ".zarr" in name]) == 1 + + +def test_robust_save_to_file_removes_a_folder_it_failed_to_write( + tmp_path: Path, +) -> None: + def failing_saver(target: str) -> None: + os.makedirs(os.path.join(target, "half-written")) + raise ValueError("write failed") + + with pytest.raises(ValueError): + engine_ecmwf.robust_save_to_file(failing_saver, (), str(tmp_path / "a.zarr")) + + assert os.listdir(tmp_path) == [] + + +def test_robust_save_to_file_raises_when_a_file_can_not_be_renamed( + tmp_path: Path, +) -> None: + os.makedirs(tmp_path / "a.grib" / "not-a-grib") + + with pytest.raises(OSError): + engine_ecmwf.robust_save_to_file( + lambda target: Path(target).touch(), (), str(tmp_path / "a.grib") + ) + + assert os.listdir(tmp_path) == ["a.grib"] diff --git a/xarray_ecmwf/engine_ecmwf.py b/xarray_ecmwf/engine_ecmwf.py index a270abc..38504e7 100644 --- a/xarray_ecmwf/engine_ecmwf.py +++ b/xarray_ecmwf/engine_ecmwf.py @@ -1,8 +1,10 @@ import contextlib +import errno import functools import hashlib import logging import os +import shutil import socket import uuid from typing import Any, Callable, Iterable, Iterator, Sequence @@ -51,6 +53,16 @@ def _raw_indexing_method(self, key: tuple[int | slice, ...]) -> np.typing.ArrayL return out +def remove_file_or_folder(path: str) -> None: + if os.path.isdir(path): + shutil.rmtree(path, ignore_errors=True) + else: + try: + os.remove(path) + except Exception: + pass + + def robust_save_to_file( saver: Callable[..., Any], args: Sequence[Any], path: str ) -> None: @@ -59,13 +71,16 @@ def robust_save_to_file( try: saver(*args, tmp_path) except Exception as ex: - try: - os.remove(tmp_path) - except Exception: - pass + remove_file_or_folder(tmp_path) raise ex - os.rename(tmp_path, path) + try: + os.rename(tmp_path, path) + except OSError as ex: + remove_file_or_folder(tmp_path) + # another process saved it first and it can't be replaced: keep that one + if ex.errno not in (errno.EEXIST, errno.ENOTEMPTY): + raise ex @attrs.define(slots=False) @@ -147,7 +162,11 @@ def cached_empty_dataset(self, request: dict[str, Any]) -> Iterator[xr.Dataset]: # NOTE: be sure that read_ds is chunked so compute=False only # writes the metadata. Some open_dataset read_ds = read_ds.chunk() - read_ds.to_zarr(path, compute=False) + # every process writes its own copy, and the first one + # renamed into place is the one everybody opens + robust_save_to_file( + functools.partial(read_ds.to_zarr, compute=False), (), path + ) yield xr.open_dataset(path, engine="zarr")