289 lines
14 KiB
Python
289 lines
14 KiB
Python
import dataclasses
|
|
import os
|
|
from argparse import Namespace, _SubParsersAction
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
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_connection,
|
|
find_folder,
|
|
find_laboratory,
|
|
parse_remote_host_with_path,
|
|
)
|
|
from mdrsclient.connection import MDRSConnection
|
|
from mdrsclient.exceptions import (
|
|
IllegalArgumentException,
|
|
MDRSException,
|
|
UnexpectedException,
|
|
)
|
|
from mdrsclient.models import File, Folder, FolderSimple
|
|
from mdrsclient.settings import NUMBER_OF_PROCESS
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UploadFile:
|
|
folder: Folder
|
|
path: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DownloadFile:
|
|
file: File
|
|
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 the file or folders")
|
|
download_parser.add_argument(
|
|
"-r", "--recursive", help="Download folders and their contents recursive", action="store_true"
|
|
)
|
|
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)
|
|
# mv
|
|
move_parser = top_level_subparsers.add_parser("mv", help="move or rename file or folder")
|
|
move_parser.add_argument("src_path", help="Source remote path (remote:/lab/path/src)")
|
|
move_parser.add_argument("dest_path", help="Destination remote path (remote:/lab/path/dest)")
|
|
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.")
|
|
connection = create_connection(remote)
|
|
laboratory = find_laboratory(connection, laboratory_name)
|
|
folder = find_folder(connection, 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(connection)
|
|
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(connection, 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)
|
|
connection = create_connection(remote)
|
|
local_path = os.path.abspath(args.local_path)
|
|
if not os.path.isdir(local_path):
|
|
raise IllegalArgumentException(f"Local directory `{args.local_path}` not found.")
|
|
laboratory = find_laboratory(connection, laboratory_name)
|
|
parent_folder = find_folder(connection, laboratory, parent_path)
|
|
file = parent_folder.find_file(file_name)
|
|
download_files: list[DownloadFile] = []
|
|
if file is not None:
|
|
file_path = os.path.join(local_path, file_name)
|
|
download_files.append(DownloadFile(file, file_path))
|
|
else:
|
|
if not args.recursive:
|
|
raise IllegalArgumentException(f"Cannot download `{args.remote_path}`: Is a folder.")
|
|
sub_folder = parent_folder.find_sub_folder(file_name)
|
|
if sub_folder is None:
|
|
raise IllegalArgumentException(f"File or Folder`{file_name}` not found.")
|
|
folder_api = FolderApi(connection)
|
|
sub_folder_dirname = os.path.join(local_path, sub_folder.name)
|
|
FileCommand._multiple_download_pickup_recursive_files(
|
|
folder_api, download_files, sub_folder.id, sub_folder_dirname
|
|
)
|
|
FileCommand._multiple_download(connection, download_files)
|
|
|
|
@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_dirname = os.path.dirname(src_path)
|
|
src_basename = os.path.basename(src_path)
|
|
if dest_path.endswith("/"):
|
|
dest_dirname = dest_path
|
|
dest_basename = src_basename
|
|
else:
|
|
dest_dirname = os.path.dirname(dest_path)
|
|
dest_basename = os.path.basename(dest_path)
|
|
connection = create_connection(src_remote)
|
|
laboratory = find_laboratory(connection, src_laboratory_name)
|
|
src_parent_folder = find_folder(connection, laboratory, src_dirname)
|
|
dest_parent_folder = find_folder(connection, laboratory, dest_dirname)
|
|
src_file = src_parent_folder.find_file(src_basename)
|
|
if src_file is not None:
|
|
# source is file
|
|
dest_file = dest_parent_folder.find_file(dest_basename)
|
|
if dest_file is not None:
|
|
raise IllegalArgumentException(f"File `{dest_basename}` already exists.")
|
|
dest_sub_folder = dest_parent_folder.find_sub_folder(dest_basename)
|
|
if dest_sub_folder is not None:
|
|
raise IllegalArgumentException(
|
|
f"Cannot overwrite non-folder `{dest_basename}` with folder `{dest_path}`."
|
|
)
|
|
file_api = FileApi(connection)
|
|
if src_parent_folder.id != dest_parent_folder.id:
|
|
file_api.move(src_file, dest_parent_folder.id)
|
|
if dest_basename != src_basename:
|
|
dest_file_dict = dataclasses.asdict(src_file) | {"name": dest_basename}
|
|
dest_file = parse_obj_as(File, dest_file_dict)
|
|
file_api.update(dest_file, None)
|
|
else:
|
|
src_folder = src_parent_folder.find_sub_folder(src_basename)
|
|
if src_folder is None:
|
|
raise IllegalArgumentException(f"File or Folder `{src_basename}` not found.")
|
|
# source is folder
|
|
dest_file = dest_parent_folder.find_file(dest_basename)
|
|
if dest_file is not None:
|
|
raise IllegalArgumentException(
|
|
f"Cannot overwrite non-folder `{dest_basename}` with folder `{src_path}`."
|
|
)
|
|
dest_folder = dest_parent_folder.find_sub_folder(dest_basename)
|
|
if dest_folder is not None:
|
|
if dest_folder.id == src_folder.id:
|
|
raise IllegalArgumentException(f"`{src_path}` and `{src_path}` are the same folder.")
|
|
raise IllegalArgumentException(f"Cannot move `{src_path}` to `{dest_path}`: Folder not empty.")
|
|
folder_api = FolderApi(connection)
|
|
if src_parent_folder.id != dest_parent_folder.id:
|
|
folder_api.move(src_folder, dest_parent_folder.id)
|
|
if src_basename != dest_basename:
|
|
dest_folder_dict = dataclasses.asdict(src_folder) | {"name": dest_basename}
|
|
dest_folder = parse_obj_as(FolderSimple, dest_folder_dict)
|
|
folder_api.update(dest_folder)
|
|
|
|
@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)
|
|
connection = create_connection(remote)
|
|
laboratory = find_laboratory(connection, laboratory_name)
|
|
folder = find_folder(connection, laboratory, parent_path)
|
|
file = folder.find_file(file_name)
|
|
if file is None:
|
|
raise IllegalArgumentException(f"File `{file_name}` not found.")
|
|
file_api = FileApi(connection)
|
|
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)
|
|
connection = create_connection(remote)
|
|
laboratory = find_laboratory(connection, laboratory_name)
|
|
folder = find_folder(connection, laboratory, parent_path)
|
|
file = folder.find_file(file_name)
|
|
if file is None:
|
|
raise IllegalArgumentException(f"File `{file_name}` not found.")
|
|
file_api = FileApi(connection)
|
|
metadata = file_api.metadata(file)
|
|
print(metadata)
|
|
|
|
@staticmethod
|
|
def _multiple_upload(connection: MDRSConnection, upload_files: list[UploadFile]) -> None:
|
|
file_api = FileApi(connection)
|
|
with ThreadPoolExecutor(max_workers=NUMBER_OF_PROCESS) as pool:
|
|
pool.map(lambda x: FileCommand._multiple_upload_worker(file_api, x), upload_files)
|
|
|
|
@staticmethod
|
|
def _multiple_upload_worker(file_api: FileApi, upload_file: UploadFile) -> None:
|
|
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}")
|
|
|
|
@staticmethod
|
|
def _multiple_download_pickup_recursive_files(
|
|
folder_api: FolderApi, download_files: list[DownloadFile], folder_id: str, local_dirname: str
|
|
) -> None:
|
|
folder = folder_api.retrieve(folder_id)
|
|
file_dirname = os.path.join(local_dirname, folder.name)
|
|
if not os.path.exists(file_dirname):
|
|
os.makedirs(file_dirname)
|
|
for file in folder.files:
|
|
file_path = os.path.join(file_dirname, file.name)
|
|
download_files.append(DownloadFile(file, file_path))
|
|
for sub_folder in folder.sub_folders:
|
|
sub_folder_dirname = os.path.join(local_dirname, sub_folder.name)
|
|
FileCommand._multiple_download_pickup_recursive_files(
|
|
folder_api, download_files, sub_folder.id, sub_folder_dirname
|
|
)
|
|
|
|
@staticmethod
|
|
def _multiple_download(connection: MDRSConnection, download_files: list[DownloadFile]) -> None:
|
|
file_api = FileApi(connection)
|
|
with ThreadPoolExecutor(max_workers=NUMBER_OF_PROCESS) as pool:
|
|
pool.map(lambda x: FileCommand._multiple_download_worker(file_api, x), download_files)
|
|
|
|
@staticmethod
|
|
def _multiple_download_worker(file_api: FileApi, download_file: DownloadFile) -> None:
|
|
file_api.download(download_file.file, download_file.path)
|