From 497a3adf04d2e1577c63e6de299895599ba3b870 Mon Sep 17 00:00:00 2001
From: Ben Perry <bhperry94@gmail.com>
Date: Mon, 26 Feb 2024 17:59:59 -0600
Subject: [PATCH] Add callback on rsync CLI command

---
 saturnfs/cli/commands.py    | 17 +++++++++++++++--
 saturnfs/client/saturnfs.py | 12 +++++++++---
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/saturnfs/cli/commands.py b/saturnfs/cli/commands.py
index 79e9abf..95c7296 100644
--- a/saturnfs/cli/commands.py
+++ b/saturnfs/cli/commands.py
@@ -177,6 +177,7 @@ def delete(path: str, recursive: bool):
 @cli.command("rsync")
 @click.argument("source_path", type=str)
 @click.argument("destination_path", type=str)
+@click.option("--quiet", "-q", is_flag=True, default=False, help="Do not print file operations")
 @click.option(
     "-d",
     "--delete-missing",
@@ -184,12 +185,24 @@ def delete(path: str, recursive: bool):
     default=False,
     help="Delete paths from the destination that are missing in the source",
 )
-def rsync(source_path: str, destination_path: str, delete_missing: bool):
+def rsync(source_path: str, destination_path: str, delete_missing: bool, quiet: bool):
     """
     Recursively sync files between two directory trees
     """
     sfs = SaturnFS()
-    sfs.rsync(source_path, destination_path, delete_missing=delete_missing)
+
+    src_is_local = not source_path.startswith(settings.SATURNFS_FILE_PREFIX)
+    dst_is_local = not destination_path.startswith(settings.SATURNFS_FILE_PREFIX)
+    if src_is_local and dst_is_local:
+        raise SaturnError(PathErrors.AT_LEAST_ONE_REMOTE_PATH)
+
+    if quiet:
+        callback = NoOpCallback()
+    else:
+        operation = file_op(src_is_local, dst_is_local)
+        callback = FileOpCallback(operation=operation)
+
+    sfs.rsync(source_path, destination_path, delete_missing=delete_missing, callback=callback)
 
 
 @cli.command("ls")
diff --git a/saturnfs/client/saturnfs.py b/saturnfs/client/saturnfs.py
index 36675d5..ba4991e 100644
--- a/saturnfs/client/saturnfs.py
+++ b/saturnfs/client/saturnfs.py
@@ -19,6 +19,7 @@
 from fsspec.spec import AbstractBufferedFile, AbstractFileSystem, _Cached
 from fsspec.utils import other_paths
 from saturnfs import settings
+from saturnfs.cli.callback import FileOpCallback
 from saturnfs.client.file_transfer import (
     DownloadPart,
     FileTransferClient,
@@ -1248,18 +1249,23 @@ async def _cp_file(
         # by put/get instead of opening as a buffered file
         proto1, path1 = split_protocol(url)
         proto2, path2 = split_protocol(url2)
+        if isinstance(callback, FileOpCallback) and not callback.inner:
+            callback.branch(path1, path2, kwargs)
+        else:
+            kwargs["callback"] = callback
+
         if self._is_local(proto1) and self._is_saturnfs(proto2):
             if blocksize < settings.S3_MIN_PART_SIZE:
                 blocksize = settings.S3_MIN_PART_SIZE
             return self.sfs.put_file(
-                path1, path2, callback=callback, block_size=blocksize, **kwargs
+                path1, path2, block_size=blocksize, **kwargs
             )
         elif self._is_saturnfs(proto1) and self._is_local(proto2):
             return self.sfs.get_file(
-                path1, path2, callback=callback, block_size=blocksize, **kwargs
+                path1, path2, block_size=blocksize, **kwargs
             )
 
-        return await super()._cp_file(url, url2, blocksize, callback, **kwargs)
+        return await super()._cp_file(url, url2, blocksize, **kwargs)
 
     def _is_local(self, protocol: str) -> bool:
         if isinstance(LocalFileSystem.protocol, tuple):