mdrs-client-python/mdrsclient/commands/upload.py

import os
from argparse import Namespace
from concurrent.futures import ThreadPoolExecutor
from typing import Any
from unicodedata import normalize

from pydantic.dataclasses import dataclass

from mdrsclient.api import FilesApi, FoldersApi
from mdrsclient.commands.base import BaseCommand
from mdrsclient.connection import MDRSConnection
from mdrsclient.exceptions import IllegalArgumentException, MDRSException
from mdrsclient.models import File, Folder
from mdrsclient.models.file import find_file
from mdrsclient.settings import CONCURRENT


@dataclass(frozen=True)
class UploadFileInfo:
    folder: Folder
    files: list[File]
    path: str


class UploadCommand(BaseCommand):
    @classmethod
    def register(cls, parsers: Any) -> None:
        upload_parser = parsers.add_parser("upload", help="upload the file or directory")
        upload_parser.add_argument(
            "-r", "--recursive", help="upload directories and their contents recursive", action="store_true"
        )
        upload_parser.add_argument(
            "-s",
            "--skip-if-exists",
            help="skip the upload if file is already uploaded and file size is the same",
            action="store_true",
        )
        upload_parser.add_argument("local_path", help="local file path (/foo/bar/data.txt)")
        upload_parser.add_argument("remote_path", help="remote folder path (remote:/lab/path/)")
        upload_parser.set_defaults(func=cls.func)

    @classmethod
    def func(cls, args: Namespace) -> None:
        local_path = str(args.local_path)
        remote_path = str(args.remote_path)
        is_recursive = bool(args.recursive)
        is_skip_if_exists = bool(args.skip_if_exists)
        cls.upload(local_path, remote_path, is_recursive, is_skip_if_exists)

    @classmethod
    def upload(cls, local_path: str, remote_path: str, is_recursive: bool, is_skip_if_exists: bool) -> None:
        (remote, laboratory_name, r_path) = cls._parse_remote_host_with_path(remote_path)
        l_path = os.path.abspath(local_path)
        if not os.path.exists(l_path):
            raise IllegalArgumentException(f"File or directory `{local_path}` not found.")
        connection = cls._create_connection(remote)
        laboratory = cls._find_laboratory(connection, laboratory_name)
        folder = cls._find_folder(connection, laboratory, r_path)
        files = cls._find_files(connection, folder.id)
        infos: list[UploadFileInfo] = []
        if os.path.isdir(l_path):
            if not is_recursive:
                raise IllegalArgumentException(f"Cannot upload `{local_path}`: Is a directory.")
            folder_api = FoldersApi(connection)
            folder_map: dict[str, Folder] = {}
            folder_map[r_path] = folder
            files_map: dict[str, list[File]] = {}
            files_map[r_path] = files
            l_basename = os.path.basename(l_path)
            for dirpath, _, filenames in os.walk(l_path, followlinks=True):
                sub = l_basename if dirpath == l_path else os.path.join(l_basename, os.path.relpath(dirpath, l_path))
                d_dirname = os.path.join(r_path, sub)
                d_basename = os.path.basename(d_dirname)
                # prepare destination parent path
                d_parent_dirname = os.path.dirname(d_dirname)
                if folder_map.get(d_parent_dirname) is None:
                    parent_folder = cls._find_folder(connection, laboratory, d_parent_dirname)
                    folder_map[d_parent_dirname] = parent_folder
                    parent_files = cls._find_files(connection, parent_folder.id)
                    files_map[d_parent_dirname] = parent_files
                # prepare destination path
                if folder_map.get(d_dirname) is None:
                    d_folder = folder_map[d_parent_dirname].find_sub_folder(d_basename)
                    if d_folder is None:
                        d_folder_id = folder_api.create(normalize("NFC", d_basename), folder_map[d_parent_dirname].id)
                    else:
                        d_folder_id = d_folder.id
                    print(d_dirname)
                    folder_map[d_dirname] = folder_api.retrieve(d_folder_id)
                    files_map[d_dirname] = cls._find_files(connection, d_folder_id)
                    if d_folder is None:
                        folder_map[d_parent_dirname].sub_folders.append(folder_map[d_dirname])
                # register upload file list
                for filename in filenames:
                    infos.append(
                        UploadFileInfo(folder_map[d_dirname], files_map[d_dirname], os.path.join(dirpath, filename))
                    )
        else:
            infos.append(UploadFileInfo(folder, files, l_path))
        cls.__multiple_upload(connection, infos, is_skip_if_exists)

    @classmethod
    def __multiple_upload(
        cls, connection: MDRSConnection, infos: list[UploadFileInfo], is_skip_if_exists: bool
    ) -> None:
        file_api = FilesApi(connection)
        with ThreadPoolExecutor(max_workers=CONCURRENT) as pool:
            pool.map(lambda x: cls.__multiple_upload_worker(file_api, x, is_skip_if_exists), infos)

    @classmethod
    def __multiple_upload_worker(cls, file_api: FilesApi, info: UploadFileInfo, is_skip_if_exists: bool) -> None:
        basename = os.path.basename(info.path)
        file = find_file(info.files, basename)
        try:
            if file is None:
                file_api.create(info.folder.id, info.path)
            elif not is_skip_if_exists or file.size != os.path.getsize(info.path):
                file_api.update(file, info.path)
            print(os.path.join(info.folder.path, basename))
        except MDRSException as e:
            print(f"Error: {e}")