implemented to cancel to recursive download if some files failed to download.

This commit is contained in:
Yoshihiro OKUMURA 2024-07-08 20:35:48 +09:00
parent c8b16939d7
commit d392379235
Signed by: orrisroot
GPG Key ID: 470AA444C92904B2

View File

@ -8,7 +8,7 @@ from pydantic.dataclasses import dataclass
from mdrsclient.api import FilesApi, FoldersApi from mdrsclient.api import FilesApi, FoldersApi
from mdrsclient.commands.base import BaseCommand from mdrsclient.commands.base import BaseCommand
from mdrsclient.connection import MDRSConnection from mdrsclient.connection import MDRSConnection
from mdrsclient.exceptions import IllegalArgumentException from mdrsclient.exceptions import IllegalArgumentException, UnexpectedException
from mdrsclient.models import File from mdrsclient.models import File
from mdrsclient.settings import CONCURRENT from mdrsclient.settings import CONCURRENT
@ -19,6 +19,12 @@ class DownloadFileInfo:
path: str path: str
@dataclass
class DownloadContext:
hasError: bool
files: list[DownloadFileInfo]
class DownloadCommand(BaseCommand): class DownloadCommand(BaseCommand):
@classmethod @classmethod
def register(cls, parsers: Any) -> None: def register(cls, parsers: Any) -> None:
@ -52,10 +58,11 @@ class DownloadCommand(BaseCommand):
laboratory = cls._find_laboratory(connection, laboratory_name) laboratory = cls._find_laboratory(connection, laboratory_name)
r_parent_folder = cls._find_folder(connection, laboratory, r_dirname, password) r_parent_folder = cls._find_folder(connection, laboratory, r_dirname, password)
file = r_parent_folder.find_file(r_basename) file = r_parent_folder.find_file(r_basename)
download_files: list[DownloadFileInfo] = []
if file is not None: if file is not None:
context = DownloadContext(False, [])
l_path = os.path.join(l_dirname, r_basename) l_path = os.path.join(l_dirname, r_basename)
download_files.append(DownloadFileInfo(file, l_path)) context.files.append(DownloadFileInfo(file, l_path))
cls.__multiple_download(connection, context)
else: else:
folder = r_parent_folder.find_sub_folder(r_basename) folder = r_parent_folder.find_sub_folder(r_basename)
if folder is None: if folder is None:
@ -63,13 +70,13 @@ class DownloadCommand(BaseCommand):
if not is_recursive: if not is_recursive:
raise IllegalArgumentException(f"Cannot download `{r_path}`: Is a folder.") raise IllegalArgumentException(f"Cannot download `{r_path}`: Is a folder.")
folder_api = FoldersApi(connection) folder_api = FoldersApi(connection)
cls.__multiple_download_pickup_recursive_files(folder_api, download_files, folder.id, l_dirname) cls.__multiple_download_pickup_recursive_files(connection, folder_api, folder.id, l_dirname)
cls.__multiple_download(connection, download_files)
@classmethod @classmethod
def __multiple_download_pickup_recursive_files( def __multiple_download_pickup_recursive_files(
cls, folder_api: FoldersApi, infolist: list[DownloadFileInfo], folder_id: str, basedir: str cls, connection: MDRSConnection, folder_api: FoldersApi, folder_id: str, basedir: str
) -> None: ) -> None:
context = DownloadContext(False, [])
folder = folder_api.retrieve(folder_id) folder = folder_api.retrieve(folder_id)
dirname = os.path.join(basedir, folder.name) dirname = os.path.join(basedir, folder.name)
if not os.path.exists(dirname): if not os.path.exists(dirname):
@ -77,17 +84,27 @@ class DownloadCommand(BaseCommand):
print(dirname) print(dirname)
for file in folder.files: for file in folder.files:
path = os.path.join(dirname, file.name) path = os.path.join(dirname, file.name)
infolist.append(DownloadFileInfo(file, path)) context.files.append(DownloadFileInfo(file, path))
cls.__multiple_download(connection, context)
if context.hasError:
raise UnexpectedException("Some files failed to download.")
for sub_folder in folder.sub_folders: for sub_folder in folder.sub_folders:
cls.__multiple_download_pickup_recursive_files(folder_api, infolist, sub_folder.id, dirname) cls.__multiple_download_pickup_recursive_files(connection, folder_api, sub_folder.id, dirname)
@classmethod @classmethod
def __multiple_download(cls, connection: MDRSConnection, infolist: list[DownloadFileInfo]) -> None: def __multiple_download(cls, connection: MDRSConnection, context: DownloadContext) -> None:
file_api = FilesApi(connection) file_api = FilesApi(connection)
with ThreadPoolExecutor(max_workers=CONCURRENT) as pool: with ThreadPoolExecutor(max_workers=CONCURRENT) as pool:
pool.map(lambda x: cls.__multiple_download_worker(file_api, x), infolist) results = pool.map(lambda x: cls.__multiple_download_worker(file_api, x), context.files)
hasError = next(filter(lambda x: x is False, results), None)
if hasError is not None:
context.hasError = True
@classmethod @classmethod
def __multiple_download_worker(cls, file_api: FilesApi, info: DownloadFileInfo) -> None: def __multiple_download_worker(cls, file_api: FilesApi, info: DownloadFileInfo) -> bool:
try:
file_api.download(info.file, info.path) file_api.download(info.file, info.path)
except Exception:
return False
print(info.path) print(info.path)
return True