niu/mdrs-client-python
7
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use concurrent for the multiple file uploads instead of multiprocess.

Browse Source
This commit is contained in:
Yoshihiro OKUMURA
2023-05-09 13:08:50 +09:00
parent 6e065e7e25
commit c724af538b
11 changed files with 122 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
mdrsclient/commands/file.py
View File

@ -1,7 +1,7 @@
import dataclasses
import os
from argparse import Namespace, _SubParsersAction
from multiprocessing import Process
from concurrent.futures import ThreadPoolExecutor
from pydantic import parse_obj_as
from pydantic.dataclasses import dataclass
@ -9,18 +9,18 @@ from pydantic.dataclasses import dataclass
from mdrsclient.api import FileApi, FolderApi
from mdrsclient.commands.base import BaseCommand
from mdrsclient.commands.utils import (
create_session,
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
from mdrsclient.session import MDRSSession
from mdrsclient.settings import NUMBER_OF_PROCESS
@ -66,14 +66,14 @@ class FileCommand(BaseCommand):
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)
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(session)
folder_api = FolderApi(connection)
folders: dict[str, Folder] = {}
folders[path] = folder
local_basename = os.path.basename(local_path)
@ -107,7 +107,7 @@ class FileCommand(BaseCommand):
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)
FileCommand._multiple_upload(connection, upload_files)
@staticmethod
def download(args: Namespace) -> None:
@ -115,16 +115,16 @@ class FileCommand(BaseCommand):
path = path.rstrip("/")
parent_path = os.path.dirname(path)
file_name = os.path.basename(path)
session = create_session(remote)
connection = create_connection(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)
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.")
r = session.get(session.build_url("v2/" + file.download_url), stream=True)
r = connection.session.get(connection.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):
@ -151,17 +151,17 @@ class FileCommand(BaseCommand):
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)
connection = create_connection(src_remote)
laboratory = find_laboratory(connection, src_laboratory_name)
src_folder = find_folder(connection, laboratory, src_dirpath)
dest_folder = find_folder(connection, 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)
file_api = FileApi(connection)
if src_folder.id != dest_folder.id:
file_api.move(src_file, dest_folder.id)
if dest_filename != src_filename:
@ -175,13 +175,13 @@ class FileCommand(BaseCommand):
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)
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(session)
file_api = FileApi(connection)
file_api.destroy(file)
@staticmethod
@ -190,42 +190,31 @@ class FileCommand(BaseCommand):
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)
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(session)
file_api = FileApi(connection)
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()
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(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}")
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}")

36
mdrsclient/commands/folder.py
View File

@ -4,7 +4,7 @@ from argparse import Namespace, _SubParsersAction
from mdrsclient.api import FolderApi
from mdrsclient.commands.base import BaseCommand
from mdrsclient.commands.utils import (
create_session,
create_connection,
find_folder,
find_laboratory,
parse_remote_host_with_path,
@ -34,9 +34,9 @@ class FolderCommand(BaseCommand):
@staticmethod
def list(args: Namespace) -> None:
(remote, laboratory_name, path) = parse_remote_host_with_path(args.remote_path)
session = create_session(remote)
laboratory = find_laboratory(session, laboratory_name)
folder = find_folder(session, laboratory, path)
connection = create_connection(remote)
laboratory = find_laboratory(connection, laboratory_name)
folder = find_folder(connection, laboratory, path)
label = {
"type": "Type",
"acl": "Access",
@ -49,7 +49,7 @@ class FolderCommand(BaseCommand):
for key in label.keys():
length[key] = len(label[key])
for sub_folder in folder.sub_folders:
sub_laboratory = session.laboratories.find_by_id(sub_folder.lab_id)
sub_laboratory = connection.laboratories.find_by_id(sub_folder.lab_id)
sub_laboratory_name = sub_laboratory.name if sub_laboratory is not None else "(invalid)"
length["acl"] = max(length["acl"], len(sub_folder.access_level_name))
length["laboratory"] = max(length["laboratory"], len(sub_laboratory_name))
@ -71,7 +71,7 @@ class FolderCommand(BaseCommand):
print("-" * len(header.expandtabs()))
for sub_folder in sorted(folder.sub_folders, key=lambda x: x.name):
sub_laboratory = session.laboratories.find_by_id(sub_folder.lab_id)
sub_laboratory = connection.laboratories.find_by_id(sub_folder.lab_id)
sub_laboratory_name = sub_laboratory.name if sub_laboratory is not None else "(invalid)"
print(
f"{'[d]':{length['type']}}\t{sub_folder.access_level_name:{length['acl']}}\t"
@ -91,27 +91,27 @@ class FolderCommand(BaseCommand):
path = path.rstrip("/")
parent_path = os.path.dirname(path)
folder_name = os.path.basename(path)
session = create_session(remote)
laboratory = find_laboratory(session, laboratory_name)
folder = find_folder(session, laboratory, parent_path)
folder_api = FolderApi(session)
connection = create_connection(remote)
laboratory = find_laboratory(connection, laboratory_name)
folder = find_folder(connection, laboratory, parent_path)
folder_api = FolderApi(connection)
folder_api.create(folder_name, folder.id)
@staticmethod
def rmdir(args: Namespace) -> None:
(remote, laboratory_name, path) = parse_remote_host_with_path(args.remote_path)
session = create_session(remote)
laboratory = find_laboratory(session, laboratory_name)
folder = find_folder(session, laboratory, path)
folder_api = FolderApi(session)
connection = create_connection(remote)
laboratory = find_laboratory(connection, laboratory_name)
folder = find_folder(connection, laboratory, path)
folder_api = FolderApi(connection)
folder_api.destroy(folder.id)
@staticmethod
def metadata(args: Namespace) -> None:
(remote, laboratory_name, path) = parse_remote_host_with_path(args.remote_path)
session = create_session(remote)
laboratory = find_laboratory(session, laboratory_name)
folder = find_folder(session, laboratory, path)
folder_api = FolderApi(session)
connection = create_connection(remote)
laboratory = find_laboratory(connection, laboratory_name)
folder = find_folder(connection, laboratory, path)
folder_api = FolderApi(connection)
metadata = folder_api.metadata(folder.id)
print(metadata)

8
mdrsclient/commands/laboratory.py
View File

@ -2,7 +2,7 @@ from argparse import Namespace, _SubParsersAction
from mdrsclient.api import LaboratoryApi
from mdrsclient.commands.base import BaseCommand
from mdrsclient.commands.utils import create_session, parse_remote_host
from mdrsclient.commands.utils import create_connection, parse_remote_host
class LaboratoryCommand(BaseCommand):
@ -16,10 +16,10 @@ class LaboratoryCommand(BaseCommand):
@staticmethod
def list(args: Namespace) -> None:
remote = parse_remote_host(args.remote)
session = create_session(remote)
laboratory_api = LaboratoryApi(session)
connection = create_connection(remote)
laboratory_api = LaboratoryApi(connection)
laboratories = laboratory_api.list()
session.laboratories = laboratories
connection.laboratories = laboratories
label = {"id": "ID", "name": "Name", "pi_name": "PI", "full_name": "Laboratory"}
length: dict[str, int] = {}
for key in label.keys():

22
mdrsclient/commands/user.py
View File

@ -5,8 +5,8 @@ from mdrsclient.api import UserApi
from mdrsclient.commands.base import BaseCommand
from mdrsclient.commands.utils import parse_remote_host
from mdrsclient.config import ConfigFile
from mdrsclient.connection import MDRSConnection
from mdrsclient.exceptions import MissingConfigurationException
from mdrsclient.session import MDRSSession
class UserCommand(BaseCommand):
@ -31,13 +31,13 @@ class UserCommand(BaseCommand):
config = ConfigFile(remote)
if config.url is None:
raise MissingConfigurationException(f"Remote host `{remote}` is not found.")
session = MDRSSession(config.remote, config.url)
connection = MDRSConnection(config.remote, config.url)
username = input("Username: ").strip()
password = getpass.getpass("Password: ").strip()
user_api = UserApi(session)
user_api = UserApi(connection)
(user, token) = user_api.auth(username, password)
session.user = user
session.token = token
connection.user = user
connection.token = token
@staticmethod
def logout(args: Namespace) -> None:
@ -45,8 +45,8 @@ class UserCommand(BaseCommand):
config = ConfigFile(remote)
if config.url is None:
raise MissingConfigurationException(f"Remote host `{remote}` is not found.")
session = MDRSSession(config.remote, config.url)
session.logout()
connection = MDRSConnection(config.remote, config.url)
connection.logout()
@staticmethod
def whoami(args: Namespace) -> None:
@ -54,8 +54,8 @@ class UserCommand(BaseCommand):
config = ConfigFile(remote)
if config.url is None:
raise MissingConfigurationException(f"Remote host `{remote}` is not found.")
session = MDRSSession(config.remote, config.url)
if session.token is not None and session.token.is_expired:
session.logout()
username = session.user.username if session.user is not None else "(Anonymous)"
connection = MDRSConnection(config.remote, config.url)
if connection.token is not None and connection.token.is_expired:
connection.logout()
username = connection.user.username if connection.user is not None else "(Anonymous)"
print(username)

20
mdrsclient/commands/utils.py
View File

@ -2,6 +2,7 @@ import re
from mdrsclient.api import FolderApi, LaboratoryApi
from mdrsclient.config import ConfigFile
from mdrsclient.connection import MDRSConnection
from mdrsclient.exceptions import (
IllegalArgumentException,
MissingConfigurationException,
@ -9,28 +10,27 @@ from mdrsclient.exceptions import (
UnexpectedException,
)
from mdrsclient.models import Folder, Laboratory
from mdrsclient.session import MDRSSession
def create_session(remote: str) -> MDRSSession:
def create_connection(remote: str) -> MDRSConnection:
config = ConfigFile(remote)
if config.url is None:
raise MissingConfigurationException(f"Remote host `{remote}` is not found.")
return MDRSSession(config.remote, config.url)
return MDRSConnection(config.remote, config.url)
def find_laboratory(session: MDRSSession, laboratory_name: str) -> Laboratory:
if session.laboratories.empty() or session.token is not None and session.token.is_expired:
laboratory_api = LaboratoryApi(session)
session.laboratories = laboratory_api.list()
laboratory = session.laboratories.find_by_name(laboratory_name)
def find_laboratory(connection: MDRSConnection, laboratory_name: str) -> Laboratory:
if connection.laboratories.empty() or connection.token is not None and connection.token.is_expired:
laboratory_api = LaboratoryApi(connection)
connection.laboratories = laboratory_api.list()
laboratory = connection.laboratories.find_by_name(laboratory_name)
if laboratory is None:
raise IllegalArgumentException(f"Laboratory `{laboratory_name}` not found.")
return laboratory
def find_folder(session: MDRSSession, laboratory: Laboratory, path: str) -> Folder:
folder_api = FolderApi(session)
def find_folder(connection: MDRSConnection, laboratory: Laboratory, path: str) -> Folder:
folder_api = FolderApi(connection)
folders = folder_api.list(laboratory.id, path)
if len(folders) != 1:
raise UnexpectedException(f"Folder `{path}` not found.")