Repository: ungarj/mapchete
Pull request #676 (2025-04-03 15:00:55)

CLI: mpath sync


Patch:
diff --git a/mapchete/cli/mpath.py b/mapchete/cli/mpath.py
index abbc86e..25c3e51 100644
--- a/mapchete/cli/mpath.py
+++ b/mapchete/cli/mpath.py
@@ -7,12 +7,17 @@ import click
 import tqdm
 
 from mapchete.cli import options
-from mapchete.io import copy
+from mapchete.cli.progress_bar import PBar
 from mapchete.path import MPath
 
 logger = logging.getLogger(__name__)
 
-opt_recursive = click.option("--recursive", "-r", is_flag=True)
+opt_recursive = click.option(
+    "--recursive",
+    "-r",
+    is_flag=True,
+    help="Recursively copy all files within a directory.",
+)
 
 
 @click.group()
@@ -33,7 +38,14 @@ def exists(path: MPath):
 @options.opt_dst_fs_opts
 @options.opt_overwrite
 @opt_recursive
-@click.option("--skip-existing", is_flag=True)
+@click.option("--skip-existing", is_flag=True, help="Skip file if already exists.")
+@click.option(
+    "--chunksize",
+    type=click.INT,
+    default=1024 * 1024,
+    show_default=True,
+    help="Read and write chunk size in bytes.",
+)
 @options.opt_debug
 def cp(
     path: MPath,
@@ -42,6 +54,7 @@ def cp(
     recursive: bool = False,
     skip_existing: bool = False,
     debug: bool = False,
+    chunksize: int = 1024 * 1024,
     **_,
 ):
     try:
@@ -50,9 +63,9 @@ def cp(
                 raise click.UsageError(
                     "source path is directory, --recursive flag required"
                 )
-            for _, __, files in tqdm.tqdm(path.walk(), disable=debug):
+            for contents in tqdm.tqdm(path.walk(), disable=debug):
                 for src_file in tqdm.tqdm(
-                    [file for file in files if not file.is_directory()],
+                    contents.files,
                     leave=False,
                     disable=debug,
                 ):
@@ -65,11 +78,108 @@ def cp(
                         logger.debug(message)
                     else:
                         tqdm.tqdm.write(message)
-                    copy(
-                        src_file, dst_file, overwrite=overwrite, exists_ok=skip_existing
-                    )
+                    with PBar(
+                        total=100,
+                        desc="chunks",
+                        disable=debug,
+                        leave=False,
+                        bar_format="{percentage:3.0f}%|{bar}|{elapsed}<{remaining}",
+                    ) as pbar:
+                        src_file.cp(
+                            dst_file,
+                            overwrite=overwrite,
+                            exists_ok=skip_existing,
+                            observers=[pbar],
+                            chunksize=chunksize,
+                        )
         else:
-            copy(path, out_path, overwrite=overwrite, exists_ok=skip_existing)
+            with PBar(
+                total=100,
+                disable=debug,
+                bar_format="{percentage:3.0f}%|{bar}|{elapsed}<{remaining}",
+            ) as pbar:
+                path.cp(
+                    out_path,
+                    overwrite=overwrite,
+                    exists_ok=skip_existing,
+                    observers=[pbar],
+                    chunksize=chunksize,
+                )
+    except Exception as exc:  # pragma: no cover
+        if debug:
+            raise
+        raise click.ClickException(str(exc))
+
+
+@mpath.command(help="Sync between paths.")
+@options.arg_path
+@options.arg_out_path
+@options.opt_src_fs_opts
+@options.opt_dst_fs_opts
+@click.option(
+    "--chunksize",
+    type=click.INT,
+    default=1024 * 1024,
+    show_default=True,
+    help="Read and write chunk size in bytes.",
+)
+@click.option(
+    "--compare-checksums",
+    is_flag=True,
+    help="Calculate checksums of objects. WARNING: this will effectively read all of the data (source and destination)!",
+)
+@options.opt_debug
+def sync(
+    path: MPath,
+    out_path: MPath,
+    debug: bool = False,
+    chunksize: int = 1024 * 1024,
+    compare_checksums: bool = False,
+    **_,
+):
+    try:
+        if path.is_directory():
+            tqdm.tqdm.write(f"sync {path} to {out_path} ...")
+            for contents in path.walk(absolute_paths=True):
+                dst_root = out_path / os.path.relpath(
+                    str(contents.root.without_protocol()),
+                    start=str(path.without_protocol()),
+                )
+                try:
+                    dst_files = set([file.name for file in dst_root.ls()])
+                except FileNotFoundError:
+                    dst_files = set()
+                for src_file in tqdm.tqdm(contents.files, desc="files", leave=False):
+                    dst_file = out_path / os.path.relpath(
+                        str(src_file.without_protocol()),
+                        start=str(path.without_protocol()),
+                    )
+                    if (
+                        src_file.name not in dst_files
+                        or (  # file does not exist on destination
+                            src_file.checksum()
+                            != dst_file.checksum()  # file contents are not identical
+                            if compare_checksums
+                            else src_file.size() != dst_file.size()  # file sizes differ
+                        )
+                    ):
+                        with PBar(
+                            total=100,
+                            disable=debug,
+                            unit="B",
+                            unit_scale=True,
+                            unit_divisor=1024,
+                            # bar_format="{percentage:3.0f}%|{bar}|{n_fmt}/{total_fmt}|{elapsed}<{remaining}",
+                            desc=src_file.name,
+                        ) as pbar:
+                            src_file.cp(
+                                dst_file,
+                                overwrite=True,
+                                chunksize=chunksize,
+                                observers=[pbar],
+                            )
+        else:  # pragma: no cover
+            raise NotImplementedError()
     except Exception as exc:  # pragma: no cover
         if debug:
             raise
diff --git a/mapchete/cli/progress_bar.py b/mapchete/cli/progress_bar.py
index 2c35c9b..bdd193c 100644
--- a/mapchete/cli/progress_bar.py
+++ b/mapchete/cli/progress_bar.py
@@ -36,7 +36,7 @@ class PBar(ObserverProtocol):
     ):
         if progress:
             if self._pbar.total is None or self._pbar.total != progress.total:
-                self._pbar.reset(progress.total)
+                self._pbar.reset(progress.total)  # type: ignore
 
             self._pbar.update(progress.current - self._pbar.n)
 
diff --git a/mapchete/commands/cp.py b/mapchete/commands/cp.py
index 7eb1e5b..9de8e6a 100644
--- a/mapchete/commands/cp.py
+++ b/mapchete/commands/cp.py
@@ -139,8 +139,6 @@ def cp(
                                 dst_mp,
                                 src_tiles_exist,
                                 dst_tiles_exist,
-                                src_tiledir.fs,
-                                dst_tiledir.fs,
                                 overwrite,
                             ),
                         ),
@@ -156,9 +154,7 @@ def cp(
                     all_observers.notify(message=f"{total_copied} tiles copied")
 
 
-def _copy_tile(
-    tile, src_mp, dst_mp, src_tiles_exist, dst_tiles_exist, src_fs, dst_fs, overwrite
-):
+def _copy_tile(tile, src_mp, dst_mp, src_tiles_exist, dst_tiles_exist, overwrite):
     src_path = src_mp.config.output_reader.get_path(tile)
     # only copy if source tile exists
     if src_tiles_exist[tile]:
@@ -170,7 +166,7 @@ def _copy_tile(
 
         # copy from source to target
         dst_path = dst_mp.config.output_reader.get_path(tile)
-        copy(src_path, dst_path, src_fs=src_fs, dst_fs=dst_fs, overwrite=overwrite)
+        copy(src_path, dst_path, overwrite=overwrite)
         msg = f"{tile}: copy {src_path} to {dst_path}"
         logger.debug(msg)
         return 1, msg
diff --git a/mapchete/io/_misc.py b/mapchete/io/_misc.py
index 1ef0d7f..b08390d 100644
--- a/mapchete/io/_misc.py
+++ b/mapchete/io/_misc.py
@@ -7,11 +7,13 @@ from rasterio.warp import calculate_default_transform
 from retry import retry
 from shapely.errors import TopologicalError
 from shapely.geometry import box
+from tilematrix import TilePyramid
 
 from mapchete.geometry import reproject_geometry, segmentize_geometry
 from mapchete.path import MPath
 from mapchete.settings import IORetrySettings
 from mapchete.tile import BufferedTile, BufferedTilePyramid
+from mapchete.types import MPathLike
 
 logger = logging.getLogger(__name__)
 
@@ -55,7 +57,7 @@ def get_best_zoom_level(input_file, tile_pyramid_type):
             return max([0, zoom - 1])
 
 
-def get_segmentize_value(input_file=None, tile_pyramid=None):
+def get_segmentize_value(input_file: MPathLike, tile_pyramid: TilePyramid):
     """
     Return the recommended segmentation value in input file units.
 
@@ -193,47 +195,18 @@ def get_boto3_bucket(bucket_name):  # pragma: no cover
     **dict(IORetrySettings(exceptions=(OSError, FSTimeoutError, TimeoutError))),
 )
 def copy(
-    src_path,
-    dst_path,
-    src_fs=None,
-    dst_fs=None,
-    overwrite=False,
+    src_path: MPathLike,
+    dst_path: MPathLike,
+    overwrite: bool = False,
     exists_ok: bool = False,
+    read_blocksize: int = 0,
+    read_chunksize: int = 1024 * 1024,
 ):
     """Copy path from one place to the other."""
-    src_path = MPath.from_inp(src_path, fs=src_fs)
-    dst_path = MPath.from_inp(dst_path, fs=dst_fs)
-
-    if dst_path.fs.exists(dst_path):
-        if overwrite:
-            pass
-        elif exists_ok:
-            logger.debug("%s already exists", str(dst_path))
-            return
-        else:
-            raise IOError(f"{dst_path} already exists")
-
-    # create parent directories on local filesystems
-    dst_path.parent.makedirs()
-
-    try:
-        # copy either within a filesystem or between filesystems
-        if src_path.fs == dst_path.fs:
-            src_path.fs.copy(str(src_path), str(dst_path))
-        else:
-            # read source data first
-            with src_path.open("rb") as src:
-                content = src.read()
-            # only write to destination if reading source data didn't raise errors,
-            # otherwise we can end up with empty objects on an object store
-            with dst_path.open("wb") as dst:
-                dst.write(content)
-    except Exception as exception:  # pragma: no cover
-        # This is a hack because some tool using aiohttp does not raise a
-        # ClientResponseError directly but masks it as a generic Exception and thus
-        # preventing our retry mechanism to kick in.
-        if repr(exception).startswith('Exception("ClientResponseError'):
-            raise ConnectionError(repr(exception)).with_traceback(
-                exception.__traceback__
-            ) from exception
-        raise
+    MPath.from_inp(src_path).cp(
+        dst_path,
+        overwrite=overwrite,
+        exists_ok=exists_ok,
+        read_block_size=read_blocksize,
+        chunksize=read_chunksize,
+    )
diff --git a/mapchete/path.py b/mapchete/path.py
index ad9ca0f..1e3682c 100644
--- a/mapchete/path.py
+++ b/mapchete/path.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import hashlib
 import json
 import logging
 import os
@@ -10,7 +11,7 @@ from collections import defaultdict
 from datetime import datetime
 from functools import cached_property
 from io import TextIOWrapper
-from typing import IO, Generator, List, Optional, Set, TextIO, Tuple, Union
+from typing import IO, Generator, List, Optional, Set, TextIO, NamedTuple, Union
 
 from aiohttp import BasicAuth
 import fiona
@@ -24,9 +25,11 @@ from retry import retry
 
 from mapchete.executor import Executor
 from mapchete.pretty import pretty_bytes
+from mapchete.protocols import ObserverProtocol
 from mapchete.settings import GDALHTTPOptions, IORetrySettings, mapchete_options
 from mapchete.tile import BatchBy, BufferedTile
-from mapchete.types import MPathLike
+from mapchete.timer import Timer
+from mapchete.types import MPathLike, Progress
 
 logger = logging.getLogger(__name__)
 
@@ -34,6 +37,12 @@ UNALLOWED_S3_KWARGS = ["timeout"]
 UNALLOWED_HTTP_KWARGS = ["username", "password"]
 
 
+class DirectoryContent(NamedTuple):
+    root: MPath
+    subdirs: List[MPath]
+    files: List[MPath]
+
+
 class MPath(os.PathLike):
     """
     Partially replicates pathlib.Path but with remote file support.
@@ -114,6 +123,16 @@ class MPath(os.PathLike):
             self._info = self.fs.info(self._path_str)
         return self._info
 
+    def checksum(self, algo: str = "sha256", block_size: int = 1024 * 1024) -> str:
+        """Stream a file and compute its checksum."""
+        hasher = hashlib.new(algo)
+
+        with self.open("rb") as src:
+            for chunk in iter(lambda: src.read(block_size), b""):
+                hasher.update(chunk)  # type: ignore
+
+        return hasher.hexdigest()
+
     def to_dict(self) -> dict:
         return dict(
             path=self._path_str,
@@ -241,9 +260,9 @@ class MPath(os.PathLike):
     def crop(self, elements: int) -> MPath:
         return self.new("/".join(self.elements[elements:]))
 
-    def new(self, path: MPathLike) -> "MPath":
+    def new(self, path: MPathLike, **kwargs) -> "MPath":
         """Create a new MPath instance with given path."""
-        return MPath(path, **self._kwargs)
+        return MPath(path, **self._kwargs, **kwargs)
 
     def exists(self) -> bool:
         """Check if path exists."""
@@ -313,10 +332,10 @@ class MPath(os.PathLike):
         )
 
     def open(
-        self, mode: str = "r"
+        self, mode: str = "r", **kwargs
     ) -> Union[IO, TextIO, TextIOWrapper, AbstractBufferedFile]:
         """Open file."""
-        return self.fs.open(self._path_str, mode)
+        return self.fs.open(self._path_str, mode, **kwargs)
 
     @retry(logger=logger, **dict(IORetrySettings()))
     def read_text(self) -> str:
@@ -363,37 +382,39 @@ class MPath(os.PathLike):
             logger.debug("create directory %s", str(self))
             self.fs.makedirs(self, exist_ok=exist_ok)
 
-    def _create_full_path(self, path: str, absolute_path: bool = True):
+    def _create_full_path(
+        self, path: Union[MPath, str, dict], absolute_path: bool = True
+    ):
+        if isinstance(path, str):
+            path_info = None
+            path_str = path
+        elif isinstance(path, dict):
+            path_info = path
+            path_str: str = path["name"]
+        elif isinstance(path, MPath):
+            path_info = path._info
+            path_str = path._path_str
+        else:  # pragma: no cover
+            raise TypeError(f"invalid path type: {path}")
         if "s3" in self.protocols:
             # s3fs returns paths without protocol ("s3://"), therefore we have to append it again:
-            return self.new(path).with_protocol("s3")
+            out_path = self.new(path_str, _info=path_info).with_protocol("s3")
         elif absolute_path:
-            return self.new(path)
-        return self.new(os.path.relpath(path, start=self))
+            out_path = self.new(path_str, _info=path_info)
+        else:
+            out_path = self.new(os.path.relpath(path_str, start=self), _info=path_info)
+        out_path._info = path_info
+        return out_path
 
     def ls(
-        self,
-        detail: bool = False,
-        absolute_paths: bool = True,
-    ) -> List[Union[MPath, dict]]:
-        if detail:
-            # return as is but convert "name" from string to MPath instead
-            return [
-                {
-                    k: (
-                        self._create_full_path(v, absolute_path=absolute_paths)
-                        if k == "name"
-                        else v
-                    )
-                    for k, v in path_info.items()
-                }
-                for path_info in self.fs.ls(self._path_str, detail=detail)
-            ]
-        else:
-            return [
-                self._create_full_path(path, absolute_path=absolute_paths)
-                for path in self.fs.ls(self._path_str, detail=detail)
-            ]
+        self, absolute_paths: bool = True, detail: Optional[bool] = None
+    ) -> List[MPath]:
+        if detail is not None:  # pragma: no cover
+            warnings.warn(DeprecationWarning("'detail' kwarg is deprecated."))
+        return [
+            self._create_full_path(path_info, absolute_path=absolute_paths)
+            for path_info in self.fs.ls(self._path_str, detail=True)
+        ]
 
     def walk(
         self,
@@ -402,7 +423,7 @@ class MPath(os.PathLike):
         absolute_paths: bool = True,
         **kwargs,
     ) -> Generator[
-        Tuple[Union[MPath, dict], List[Union[MPath, dict]], List[Union[MPath, dict]]],
+        DirectoryContent,
         None,
         None,
     ]:
@@ -410,17 +431,17 @@ class MPath(os.PathLike):
             str(self), maxdepth=maxdepth, topdown=topdown, **kwargs
         ):
             if isinstance(root, str):
-                yield (
-                    self._create_full_path(root, absolute_path=absolute_paths),
-                    [
+                yield DirectoryContent(
+                    root=self._create_full_path(root, absolute_path=absolute_paths),
+                    subdirs=[
                         self._create_full_path(
-                            MPath(root) / subpath, absolute_path=absolute_paths
+                            self.new(root) / subpath, absolute_path=absolute_paths
                         )
                         for subpath in subpaths
                     ],
-                    [
+                    files=[
                         self._create_full_path(
-                            MPath(root) / file, absolute_path=absolute_paths
+                            self.new(root) / file, absolute_path=absolute_paths
                         )
                         for file in files
                     ],
@@ -439,8 +460,84 @@ class MPath(os.PathLike):
             else:  # pragma: no cover
                 raise
 
+    def cp(
+        self,
+        destination: MPathLike,
+        overwrite: bool = False,
+        exists_ok: bool = False,
+        read_block_size: int = 0,
+        chunksize: int = 1024 * 1024,  # 1MB
+        observers: Optional[List[ObserverProtocol]] = None,
+    ) -> None:
+        """
+        Copy file contents to destination.
+        """
+        from mapchete.commands.observer import Observers
+
+        all_observers = Observers(observers)
+        dst_path = MPath.from_inp(destination)
+        try:
+            if dst_path.endswith("/") or dst_path.is_directory():
+                dst_path = dst_path / self.name
+        except FileNotFoundError:  # pragma: no cover
+            pass
+
+        if dst_path.exists():
+            if overwrite:
+                pass
+            elif exists_ok:
+                msg = f"{str(dst_path)} already exists"
+                all_observers.notify(message=msg)
+                logger.debug(msg)
+                return
+            else:
+                raise IOError(f"{dst_path} already exists")
+
+        # create parent directories on local filesystems
+        dst_path.parent.makedirs()
+
+        try:
+            # copy either within a filesystem or between filesystems
+            msg = f"copy {self} to {dst_path} {'(overwrite)' if overwrite else ''}..."
+            all_observers.notify(message=msg)
+            logger.debug(msg)
+            with Timer() as duration:
+                if self.fs == dst_path.fs:
+                    self.fs.copy(str(self), str(dst_path))
+                else:
+                    with self.open("rb", block_size=read_block_size) as src:
+                        with dst_path.open("wb") as dst:
+                            total_size = self.size()
+                            transferred = 0
+                            all_observers.notify(
+                                progress=Progress(current=transferred, total=total_size)
+                            )
+                            for chunk in iter(lambda: src.read(chunksize), b""):
+                                transferred += chunksize
+                                dst.write(chunk)  # type: ignore
+                                all_observers.notify(
+                                    progress=Progress(
+                                        current=min([transferred, total_size]),
+                                        total=total_size,
+                                    )
+                                )
+            all_observers.notify(message=f"copied in {duration}")
+        except Exception as exception:  # pragma: no cover
+            # delete file if something failed
+            # dst_path should either not even exist and if, the overwrite flag is active anyways
+            dst_path.rm(ignore_errors=True)
+
+            # This is a hack because some tool using aiohttp does not raise a
+            # ClientResponseError directly but masks it as a generic Exception and thus
+            # preventing our retry mechanism to kick in.
+            if repr(exception).startswith('Exception("ClientResponseError'):
+                raise ConnectionError(repr(exception)).with_traceback(
+                    exception.__traceback__
+                ) from exception
+            raise
+
     def size(self) -> int:
-        return self.info().get("size", self.info().get("Size"))
+        return self.info().get("size", self.info().get("Size"))  # type: ignore
 
     def pretty_size(self) -> str:
         return pretty_bytes(self.size())
@@ -458,11 +555,13 @@ class MPath(os.PathLike):
             raise ValueError("Object timestamp could not be determined.")
 
     def is_directory(self) -> bool:
-        # for S3 objects use the possible cached info directory
-        if "StorageClass" in self.info():  # pragma: no cover
-            return self.info().get("StorageClass") == "DIRECTORY"
-        else:
-            return self.fs.isdir(self._path_str)
+        try:
+            # for S3 objects use the possible cached info directory
+            if "StorageClass" in self.info():  # pragma: no cover
+                return self.info().get("StorageClass") == "DIRECTORY"
+        except FileNotFoundError:
+            pass
+        return self.fs.isdir(self._path_str)
 
     def joinpath(self, *other: Union[MPathLike, List[MPathLike]]) -> MPath:
         """Join path with other."""


Test patch:
diff --git a/test/cli/test_mpath.py b/test/cli/test_mpath.py
index 3ac0a7d..39fdbf1 100644
--- a/test/cli/test_mpath.py
+++ b/test/cli/test_mpath.py
@@ -1,6 +1,7 @@
 from test.cli.default import run_cli
 
 import pytest
+from pytest_lazyfixture import lazy_fixture
 
 from mapchete.cli.mpath import mpath
 
@@ -53,3 +54,18 @@ def test_read_json(metadata_json):
 
 def test_read_text(metadata_json):
     assert run_cli(["read-text", str(metadata_json)], cli=mpath)
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize(
+    "tiledir",
+    [
+        lazy_fixture("local_tiledir"),
+        # lazy_fixture("http_tiledir"),
+        # lazy_fixture("secure_http_tiledir"),
+    ],
+)
+def test_sync_dir(tiledir, mp_tmpdir):
+    assert not mp_tmpdir.ls()
+    assert run_cli(["sync", str(tiledir), str(mp_tmpdir), "--debug"], cli=mpath)
+    assert mp_tmpdir.ls()
diff --git a/test/test_path.py b/test/test_path.py
index 7dbe122..b5ad515 100644
--- a/test/test_path.py
+++ b/test/test_path.py
@@ -5,6 +5,7 @@ import pytest
 from pytest_lazyfixture import lazy_fixture
 
 from mapchete.config import get_hash
+from mapchete.io.raster.referenced_raster import ReferencedRaster
 from mapchete.path import MPath, batch_sort_property
 
 
@@ -160,8 +161,7 @@ def test_ls(path):
     for p in path.ls():
         assert isinstance(p, MPath)
         assert p.is_remote() == dir_is_remote
-    for p in path.ls(detail=True):
-        assert isinstance(p.get("name"), MPath)
+        assert p._info
 
 
 @pytest.mark.integration
@@ -535,3 +535,50 @@ def test_s3_region_name():
     path.fs
     # not an ideal way to test this, but the storage_options are passed on to S3FS that way
     assert path.storage_options.get("client_kwargs", {}).get("region_name") == "bar"
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize(
+    "src_path",
+    [
+        lazy_fixture("raster_4band"),
+        lazy_fixture("raster_4band_s3"),
+        lazy_fixture("raster_4band_http"),
+        lazy_fixture("raster_4band_secure_http"),
+        lazy_fixture("raster_4band_aws_s3"),
+    ],
+)
+@pytest.mark.parametrize(
+    "dst_dir",
+    [
+        lazy_fixture("mp_s3_tmpdir"),
+        lazy_fixture("mp_tmpdir"),
+    ],
+)
+def test_cp(src_path: MPath, dst_dir: MPath):
+    tempdir = dst_dir / "temp/"
+    dst_path = tempdir / src_path.name
+    try:
+        src_path.cp(tempdir)
+        assert dst_path.exists()
+        assert not ReferencedRaster.from_file(dst_path).data.mask.all()
+    finally:
+        dst_path.rm(ignore_errors=True)
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize(
+    "src_path",
+    [
+        lazy_fixture("raster_4band"),
+        lazy_fixture("raster_4band_s3"),
+        lazy_fixture("raster_4band_http"),
+        lazy_fixture("raster_4band_secure_http"),
+        lazy_fixture("raster_4band_aws_s3"),
+    ],
+)
+def test_checksum(src_path: MPath):
+    assert (
+        src_path.checksum()
+        == "fff06260d08965a898021b9513dc9226f6c3b964d755734357603c21fb2359ad"
+    )