import dataclasses import os from argparse import Namespace, _SubParsersAction from multiprocessing import Process from pydantic import parse_obj_as from pydantic.dataclasses import dataclass from mdrsclient.api import FileApi, FolderApi from mdrsclient.commands.base import BaseCommand from mdrsclient.commands.utils import ( create_session, find_folder, find_laboratory, parse_remote_host_with_path, ) from mdrsclient.exceptions import ( IllegalArgumentException, MDRSException, UnexpectedException, ) from mdrsclient.models import File, Folder from mdrsclient.session import MDRSSession from mdrsclient.settings import NUMBER_OF_PROCESS @dataclass(frozen=True) class UploadFile: folder: Folder path: str class FileCommand(BaseCommand): @staticmethod def register(top_level_subparsers: _SubParsersAction) -> None: # upload upload_parser = top_level_subparsers.add_parser("upload", help="upload the file or directories") upload_parser.add_argument( "-r", "--recursive", help="Upload directories and their contents recursive", 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=FileCommand.upload) # download download_parser = top_level_subparsers.add_parser("download", help="download a file") download_parser.add_argument("remote_path", help="Remote file path (remote:/lab/path/file)") download_parser.add_argument("local_path", help="Local folder path (/foo/bar/)") download_parser.set_defaults(func=FileCommand.download) # move move_parser = top_level_subparsers.add_parser("move", help="move a file") move_parser.add_argument("src_path", help="Source remote file path (remote:/lab/path/file)") move_parser.add_argument("dest_path", help="Destination remote file path (remote:/lab/path/file)") move_parser.set_defaults(func=FileCommand.move) # remove remove_parser = top_level_subparsers.add_parser("remove", help="remove a file") remove_parser.add_argument("remote_path", help="Remote file path (remote:/lab/path/file)") remove_parser.set_defaults(func=FileCommand.remove) # file-metadata metadata_parser = top_level_subparsers.add_parser("file-metadata", help="get the file metadata") metadata_parser.add_argument("remote_path", help="Remote file path (remote:/lab/path/file)") metadata_parser.set_defaults(func=FileCommand.metadata) @staticmethod def upload(args: Namespace) -> None: (remote, laboratory_name, path) = parse_remote_host_with_path(args.remote_path) local_path = os.path.realpath(args.local_path) if not os.path.exists(local_path): raise IllegalArgumentException(f"File or directory `{args.local_path}` not found.") session = create_session(remote) laboratory = find_laboratory(session, laboratory_name) folder = find_folder(session, laboratory, path) upload_files: list[UploadFile] = [] if os.path.isdir(local_path): if not args.recursive: raise IllegalArgumentException(f"Cannot upload `{args.local_path}`: Is a directory.") folder_api = FolderApi(session) folders: dict[str, Folder] = {} folders[path] = folder local_basename = os.path.basename(local_path) for dirpath, dirnames, filenames in os.walk(local_path): sub = ( local_basename if dirpath == local_path else os.path.join(local_basename, os.path.relpath(dirpath, local_path)) ) dest_folder_path = os.path.join(path, sub) dest_folder_name = os.path.basename(dest_folder_path) # prepare destination parent path dest_parent_folder_path = os.path.dirname(dest_folder_path) if folders.get(dest_parent_folder_path) is None: res = folder_api.list(laboratory.id, dest_parent_folder_path) if len(res) != 1: raise UnexpectedException(f"Remote folder `{dest_parent_folder_path}` not found.") folders[dest_parent_folder_path] = folder_api.retrieve(res[0].id) # prepare destination path if folders.get(dest_folder_path) is None: dest_folder_simple = folders[dest_parent_folder_path].find_sub_folder(dest_folder_name) if dest_folder_simple is None: dest_folder_id = folder_api.create(dest_folder_name, folders[dest_parent_folder_path].id) else: dest_folder_id = dest_folder_simple.id folders[dest_folder_path] = folder_api.retrieve(dest_folder_id) if dest_folder_simple is None: folders[dest_parent_folder_path].sub_folders.append(folders[dest_folder_path]) # register upload file list for filename in filenames: upload_files.append(UploadFile(folders[dest_folder_path], os.path.join(dirpath, filename))) else: upload_files.append(UploadFile(folder, local_path)) FileCommand._multiple_upload(session, upload_files) @staticmethod def download(args: Namespace) -> None: (remote, laboratory_name, path) = parse_remote_host_with_path(args.remote_path) path = path.rstrip("/") parent_path = os.path.dirname(path) file_name = os.path.basename(path) session = create_session(remote) if not os.path.isdir(args.local_path): raise IllegalArgumentException(f"Local directory `{args.local_path}` not found.") local_file = os.path.join(args.local_path, file_name) laboratory = find_laboratory(session, laboratory_name) folder = find_folder(session, laboratory, parent_path) file = folder.find_file(file_name) if file is None: raise IllegalArgumentException(f"File `{file_name}` not found.") r = session.get(session.build_url("v2/" + file.download_url), stream=True) try: with open(local_file, "wb") as f: for chunk in r.iter_content(chunk_size=4096): if chunk: f.write(chunk) f.flush() except PermissionError: raise IllegalArgumentException(f"Cannot create file `{local_file}`: Permission denied.") @staticmethod def move(args: Namespace) -> None: (src_remote, src_laboratory_name, src_path) = parse_remote_host_with_path(args.src_path) (dest_remote, dest_laboratory_name, dest_path) = parse_remote_host_with_path(args.dest_path) if src_remote != dest_remote: raise IllegalArgumentException("Remote host mismatched.") if src_laboratory_name != dest_laboratory_name: raise IllegalArgumentException("Laboratory mismatched.") src_path = src_path.rstrip("/") src_dirpath = os.path.dirname(src_path) src_filename = os.path.basename(src_path) if dest_path.endswith("/"): dest_dirpath = dest_path dest_filename = src_filename else: dest_dirpath = os.path.dirname(dest_path) dest_filename = os.path.basename(dest_path) session = create_session(src_remote) laboratory = find_laboratory(session, src_laboratory_name) src_folder = find_folder(session, laboratory, src_dirpath) dest_folder = find_folder(session, laboratory, dest_dirpath) src_file = src_folder.find_file(src_filename) if src_file is None: raise IllegalArgumentException(f"File `{src_filename}` not found.") dest_file = dest_folder.find_file(dest_filename) if dest_file is not None: raise IllegalArgumentException(f"File `{dest_filename}` already exists.") file_api = FileApi(session) if src_folder.id != dest_folder.id: file_api.move(src_file, dest_folder.id) if dest_filename != src_filename: dest_file_dict = dataclasses.asdict(src_file) | {"name": dest_filename} dest_file = parse_obj_as(File, dest_file_dict) file_api.update(dest_file, None) @staticmethod def remove(args: Namespace) -> None: (remote, laboratory_name, path) = parse_remote_host_with_path(args.remote_path) path = path.rstrip("/") parent_path = os.path.dirname(path) file_name = os.path.basename(path) session = create_session(remote) laboratory = find_laboratory(session, laboratory_name) folder = find_folder(session, laboratory, parent_path) file = folder.find_file(file_name) if file is None: raise IllegalArgumentException(f"File `{file_name}` not found.") file_api = FileApi(session) file_api.destroy(file) @staticmethod def metadata(args: Namespace) -> None: (remote, laboratory_name, path) = parse_remote_host_with_path(args.remote_path) path = path.rstrip("/") parent_path = os.path.dirname(path) file_name = os.path.basename(path) session = create_session(remote) laboratory = find_laboratory(session, laboratory_name) folder = find_folder(session, laboratory, parent_path) file = folder.find_file(file_name) if file is None: raise IllegalArgumentException(f"File `{file_name}` not found.") file_api = FileApi(session) metadata = file_api.metadata(file) print(metadata) @staticmethod def _multiple_upload(session: MDRSSession, upload_files: list[UploadFile]) -> None: processes: list[Process] = [] for idx in range(NUMBER_OF_PROCESS): processes.append( Process( target=FileCommand._multiple_upload_worker, args=(session, upload_files, idx, NUMBER_OF_PROCESS), ) ) for process in processes: process.start() for process in processes: process.join() @staticmethod def _multiple_upload_worker(session: MDRSSession, upload_files: list[UploadFile], idx: int, num_proc: int) -> None: file_api = FileApi(session) for upload_file in upload_files[idx::num_proc]: file_name = os.path.basename(upload_file.path) file = next((x for x in upload_file.folder.files if x.name == file_name), None) try: if file is None: file_api.create(upload_file.folder.id, upload_file.path) else: file_api.update(file, upload_file.path) pass except MDRSException as e: print(f"API Error: {e}")