232 lines
11 KiB
Python
232 lines
11 KiB
Python
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}")
|