Change all approach

Author: amritghimire

src/datachain/cli/__init__.py

@@ -5,7 +5,6 @@
 from multiprocessing import freeze_support
 from typing import TYPE_CHECKING, Optional
 
-from datachain.cli.commands.storages import cp_storage
 from datachain.cli.utils import get_logging_level
 
 from .commands import (
@@ -26,6 +25,8 @@
 logger = logging.getLogger("datachain")
 
 if TYPE_CHECKING:
+    from argparse import Namespace
+
     from datachain.catalog import Catalog
 
 
@@ -82,7 +83,6 @@ def main(argv: Optional[list[str]] = None) -> int:
 
 def handle_command(args, catalog, client_config) -> int:
     """Handle the different CLI commands."""
-    from datachain.cli.commands.storages import mv_storage, rm_storage
     from datachain.studio import process_auth_cli_args, process_jobs_args
 
     command_handlers = {
@@ -101,8 +101,8 @@ def handle_command(args, catalog, client_config) -> int:
         "gc": lambda: garbage_collect(catalog),
         "auth": lambda: process_auth_cli_args(args),
         "job": lambda: process_jobs_args(args),
-        "mv": lambda: mv_storage(args),
-        "rm": lambda: rm_storage(args),
+        "mv": lambda: handle_mv_command(args, catalog),
+        "rm": lambda: handle_rm_command(args, catalog),
     }
 
     handler = command_handlers.get(args.command)
@@ -115,23 +115,36 @@ def handle_command(args, catalog, client_config) -> int:
     return 1
 
 
-def handle_cp_command(args, catalog: "Catalog"):
+def _get_storage_implementation(args: "Namespace", catalog: "Catalog"):
+    from datachain.cli.commands.storage import (
+        LocalStorageImplementation,
+        StudioStorageImplementation,
+    )
     from datachain.config import Config
 
     config = Config().read().get("studio", {})
     token = config.get("token")
-    local = True if not token else args.local
-    if local:
-        return catalog.cp(
-            [args.source_path],
-            args.destination_path,
-            force=bool(args.force),
-            update=bool(args.update),
-            recursive=bool(args.recursive),
-            no_glob=args.no_glob,
-        )
+    studio = False if not token else args.studio_cloud_auth
+    return (
+        StudioStorageImplementation(args, catalog)
+        if studio
+        else LocalStorageImplementation(args, catalog)
+    )
+
+
+def handle_cp_command(args, catalog):
+    storage_implementation = _get_storage_implementation(args, catalog)
+    return storage_implementation.cp()
+
+
+def handle_mv_command(args, catalog):
+    storage_implementation = _get_storage_implementation(args, catalog)
+    return storage_implementation.mv()
+
 
-    return cp_storage(args)
+def handle_rm_command(args, catalog):
+    storage_implementation = _get_storage_implementation(args, catalog)
+    return storage_implementation.rm()
 
 
 def handle_clone_command(args, catalog):

src/datachain/cli/commands/storage/__init__.py

@@ -0,0 +1,10 @@
+from .local import LocalStorageImplementation
+from .studio import StudioStorageImplementation
+from .utils import build_file_paths, validate_upload_args
+
+__all__ = [
+    "LocalStorageImplementation",
+    "StudioStorageImplementation",
+    "build_file_paths",
+    "validate_upload_args",
+]

src/datachain/cli/commands/storage/base.py

@@ -0,0 +1,81 @@
+from typing import TYPE_CHECKING, Optional
+
+from datachain.error import DataChainError
+
+if TYPE_CHECKING:
+    from argparse import Namespace
+
+    from fsspec import AbstractFileSystem
+
+    from datachain.catalog import Catalog
+    from datachain.client.fsspec import Client
+    from datachain.remote.studio import StudioClient
+
+
+class StorageImplementation:
+    def __init__(self, args: "Namespace", catalog: "Catalog"):
+        self.args = args
+        self.catalog = catalog
+
+    def rm(self):
+        raise NotImplementedError("Remove is not implemented")
+
+    def mv(self):
+        raise NotImplementedError("Move is not implemented")
+
+    def cp(self):
+        from datachain.client.fsspec import Client
+
+        source_cls = Client.get_implementation(self.args.source_path)
+        destination_cls = Client.get_implementation(self.args.destination_path)
+
+        if source_cls.protocol == "file" and destination_cls.protocol == "file":
+            self.copy_local_to_local(source_cls)
+        elif source_cls.protocol == "file":
+            self.upload_to_remote(source_cls, destination_cls)
+        elif destination_cls.protocol == "file":
+            self.download_from_remote(destination_cls)
+        else:
+            self.copy_remote_to_remote(source_cls)
+
+    def copy_local_to_local(self, source_cls: "Client"):
+        source_fs = source_cls.create_fs()
+        source_fs.copy(
+            self.args.source_path,
+            self.args.destination_path,
+            recursive=self.args.recursive,
+        )
+        print(f"Copied {self.args.source_path} to {self.args.destination_path}")
+
+    def upload_to_remote(self, source_cls: "Client", destination_cls: "Client"):
+        raise NotImplementedError("Upload to remote is not implemented")
+
+    def download_from_remote(self, destination_cls: "Client"):
+        raise NotImplementedError("Download from remote is not implemented")
+
+    def copy_remote_to_remote(self, source_cls: "Client"):
+        raise NotImplementedError("Copy remote to remote is not implemented")
+
+    def save_upload_log(
+        self,
+        studio_client: Optional["StudioClient"],
+        destination_path: str,
+        file_paths: dict,
+        local_fs: "AbstractFileSystem",
+    ):
+        from datachain.remote.storages import get_studio_client
+
+        try:
+            if studio_client is None:
+                studio_client = get_studio_client(self.args)
+        except DataChainError:
+            return
+
+        uploads = [
+            {
+                "path": dst,
+                "size": local_fs.info(src).get("size", 0),
+            }
+            for dst, src in file_paths.items()
+        ]
+        studio_client.save_upload_log(destination_path, uploads)

src/datachain/cli/commands/storage/local.py

@@ -0,0 +1,70 @@
+from typing import TYPE_CHECKING
+
+from datachain.cli.commands.storage.base import StorageImplementation
+from datachain.cli.commands.storage.utils import build_file_paths, validate_upload_args
+
+if TYPE_CHECKING:
+    from datachain.client.fsspec import Client
+
+
+class LocalStorageImplementation(StorageImplementation):
+    def upload_to_remote(self, source_cls: "Client", destination_cls: "Client"):
+        from tqdm import tqdm
+
+        source_fs = source_cls.create_fs()
+        destination_fs = destination_cls.create_fs()
+        is_dir = validate_upload_args(self.args, source_fs)
+        file_paths = build_file_paths(self.args, source_fs, is_dir)
+        destination_bucket, _ = destination_cls.parse_url(self.args.destination_path)
+
+        for dest_path, source_path in file_paths.items():
+            file_size = source_fs.info(source_path)["size"]
+            print(f"Uploading {source_path} to {dest_path}")
+
+            with tqdm(total=file_size, unit="B", unit_scale=True) as pbar:
+                with source_fs.open(source_path, "rb") as source_file:
+                    with destination_fs.open(
+                        f"{destination_bucket}/{dest_path}", "wb"
+                    ) as dest_file:
+                        while True:
+                            chunk = source_file.read(8192)
+                            if not chunk:
+                                break
+                            dest_file.write(chunk)
+                            pbar.update(len(chunk))
+
+        self.save_upload_log(None, self.args.destination_path, file_paths, source_fs)
+
+    def download_from_remote(self, destination_cls: "Client"):
+        self.catalog.cp(
+            [self.args.source_path],
+            self.args.destination_path,
+            force=bool(self.args.force),
+            update=bool(self.args.update),
+            recursive=bool(self.args.recursive),
+            no_glob=self.args.no_glob,
+        )
+
+    def copy_remote_to_remote(self, source_cls: "Client"):
+        source_fs = source_cls.create_fs()
+        source_fs.copy(
+            self.args.source_path,
+            self.args.destination_path,
+            recursive=self.args.recursive,
+        )
+
+    def rm(self):
+        from datachain.client.fsspec import Client
+
+        client_cls = Client.get_implementation(self.args.path)
+        fs = client_cls.create_fs()
+        fs.rm(self.args.path, recursive=self.args.recursive)
+        # TODO: Add storage logging.
+
+    def mv(self):
+        from datachain.client.fsspec import Client
+
+        client_cls = Client.get_implementation(self.args.path)
+        fs = client_cls.create_fs()
+        fs.mv(self.args.path, self.args.new_path, recursive=self.args.recursive)
+        # TODO: Add storage logging.

src/datachain/cli/commands/storage/studio.py

@@ -0,0 +1,56 @@
+from typing import TYPE_CHECKING
+
+from datachain.cli.commands.storage.base import StorageImplementation
+from datachain.error import DataChainError
+
+if TYPE_CHECKING:
+    from datachain.client.fsspec import Client
+
+
+class StudioStorageImplementation(StorageImplementation):
+    def upload_to_remote(self, source_cls: "Client", destination_cls: "Client"):
+        from datachain.remote.storages import upload_to_storage
+
+        source_fs = source_cls.create_fs()
+        file_paths = upload_to_storage(self.args, source_fs)
+        self.save_upload_log(None, self.args.destination_path, file_paths, source_fs)
+
+    def download_from_remote(self, destination_cls: "Client"):
+        from datachain.remote.storages import download_from_storage
+
+        destination_fs = destination_cls.create_fs()
+        download_from_storage(self.args, destination_fs)
+
+    def copy_remote_to_remote(self, source_cls: "Client"):
+        from datachain.remote.storages import copy_inside_storage
+
+        copy_inside_storage(self.args)
+
+    def rm(self):
+        from datachain.remote.storages import get_studio_client
+
+        client = get_studio_client(self.args)
+
+        response = client.delete_storage_file(
+            self.args.path,
+            recursive=self.args.recursive,
+        )
+        if not response.ok:
+            raise DataChainError(response.message)
+
+        print(f"Deleted {self.args.path}")
+
+    def mv(self):
+        from datachain.remote.storages import get_studio_client
+
+        client = get_studio_client(self.args)
+
+        response = client.move_storage_file(
+            self.args.path,
+            self.args.new_path,
+            recursive=self.args.recursive,
+        )
+        if not response.ok:
+            raise DataChainError(response.message)
+
+        print(f"Moved {self.args.path} to {self.args.new_path}")

src/datachain/cli/commands/storage/utils.py

@@ -0,0 +1,38 @@
+import os.path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from argparse import Namespace
+
+    from fsspec.spec import AbstractFileSystem
+
+from datachain.error import DataChainError
+
+
+def validate_upload_args(args: "Namespace", local_fs: "AbstractFileSystem"):
+    """Validate upload arguments and raise appropriate errors."""
+    is_dir = local_fs.isdir(args.source_path)
+    if is_dir and not args.recursive:
+        raise DataChainError("Cannot copy directory without --recursive")
+    return is_dir
+
+
+def build_file_paths(args: "Namespace", local_fs: "AbstractFileSystem", is_dir: bool):
+    """Build mapping of destination paths to source paths."""
+    from datachain.client.fsspec import Client
+
+    client = Client.get_implementation(args.destination_path)
+    _, subpath = client.split_url(args.destination_path)
+
+    if is_dir:
+        return {
+            os.path.join(subpath, os.path.relpath(path, args.source_path)): path
+            for path in local_fs.find(args.source_path)
+        }
+
+    destination_path = (
+        os.path.join(subpath, os.path.basename(args.source_path))
+        if args.destination_path.endswith(("/", "\\")) or not subpath
+        else subpath
+    )
+    return {destination_path: args.source_path}