diff --git a/antarest/study/common/studystorage.py b/antarest/study/common/studystorage.py index 7f7c1193c0..c804c864e9 100644 --- a/antarest/study/common/studystorage.py +++ b/antarest/study/common/studystorage.py @@ -244,27 +244,6 @@ def export_output(self, metadata: T, output_id: str, target: Path) -> None: """ raise NotImplementedError() - @abstractmethod - def export_study_flat( - self, - metadata: T, - dst_path: Path, - outputs: bool = True, - output_list_filter: Optional[List[str]] = None, - denormalize: bool = True, - ) -> None: - """ - Export study to destination - - Args: - metadata: study. - dst_path: destination path. - outputs: list of outputs to keep. - output_list_filter: list of outputs to keep (None indicate all outputs). - denormalize: denormalize the study (replace matrix links by real matrices). - """ - raise NotImplementedError() - @abstractmethod def get_synthesis( self, metadata: T, params: Optional[RequestParameters] = None diff --git a/antarest/study/service.py b/antarest/study/service.py index 200a32123b..02de8d2d51 100644 --- a/antarest/study/service.py +++ b/antarest/study/service.py @@ -2,6 +2,7 @@ import io import json import logging +import shutil import os from datetime import datetime, timedelta from http import HTTPStatus @@ -1089,9 +1090,36 @@ def export_study( def export_task(notifier: TaskUpdateNotifier) -> TaskResult: try: target_study = self.get_study(uuid) - self.storage_service.get_storage(target_study).export_study( - target_study, export_path, outputs - ) + # if study.type == 'rawstudy': + # path_study = Path(study.path) + # if study.archived: + # self.storage_service.get_storage(study).unarchive(study) + # try: + # return self.storage_service.get_storage(study).export_study_flat(path_study, + # dest, + # len(output_list or []) > 0, + # output_list + # ) + # finally: + # if study.archived: + # shutil.rmtree(study.path) + if target_study.type == "ramstudy": + if target_study.archived: + self.storage_service.get_storage( + target_study + ).unarchive(target_study) + try: + self.storage_service.get_storage( + target_study + ).export_study(target_study, export_path, outputs) + finally: + if target_study.archived: + shutil.rmtree(target_study.path) + else: + self.storage_service.get_storage( + target_study + ).export_study(target_study, export_path, outputs) + self.file_transfer_manager.set_ready(export_id) return TaskResult( success=True, message=f"Study {uuid} successfully exported" @@ -1201,9 +1229,30 @@ def export_study_flat( study = self.get_study(uuid) assert_permission(params.user, study, StudyPermissionType.READ) self._assert_study_unarchived(study) - - return self.storage_service.get_storage(study).export_study_flat( - study, dest, len(output_list or []) > 0, output_list + path_study = Path(study.path) + if study.type == "rawstudy": + if study.archived: + self.storage_service.get_storage(study).unarchive(study) + try: + return self.storage_service.get_storage( + study + ).export_study_flat( + path_study=path_study, + dest=dest, + outputs=len(output_list or []) > 0, + output_list_filter=output_list, + ) + finally: + if study.archived: + shutil.rmtree(study.path) + snapshot_path = path_study / "snapshot" + output_src_path = path_study / "output" + self.storage_service.get_storage(study).export_study_flat( + path_study=snapshot_path, + dest=dest, + outputs=len(output_list or []) > 0, + output_list_filter=output_list, + output_src_path=output_src_path, ) def delete_study( diff --git a/antarest/study/storage/abstract_storage_service.py b/antarest/study/storage/abstract_storage_service.py index 404019a289..917cdb5106 100644 --- a/antarest/study/storage/abstract_storage_service.py +++ b/antarest/study/storage/abstract_storage_service.py @@ -5,6 +5,10 @@ from pathlib import Path from typing import List, Union, Optional, IO from uuid import uuid4 +import time +from zipfile import ZipFile +import os + from antarest.core.config import Config from antarest.core.exceptions import BadOutputError, StudyOutputNotFoundError @@ -272,7 +276,21 @@ def export_study( logger.info(f"Exporting study {metadata.id} to tmp path {tmpdir}") assert_this(target.name.endswith(".zip")) tmp_study_path = Path(tmpdir) / "tmp_copy" - self.export_study_flat(metadata, tmp_study_path, outputs) + if metadata.type == "": + snapshot_path = path_study / "" + output_src_path = path_study / "output" + self.export_study_flat( + snapshot_path, + tmp_study_path, + outputs, + ) + self.export_study_flat( + path_study=snapshot_path, + dest=tmp_study_path, + outputs=outputs, + output_src_path=output_src_path, + ) + self.export_study_flat(path_study, tmp_study_path, outputs) stopwatch = StopWatch() zip_dir(tmp_study_path, target) stopwatch.log_elapsed( @@ -369,3 +387,59 @@ def unarchive_study_output( exc_info=e, ) return False + + def export_study_flat( + self, + path_study: Path, + dest: Path, + outputs: bool = True, + output_list_filter: Optional[List[str]] = None, + output_src_path: Optional[Path] = None, + ) -> None: + """ + Export study to destination + Args: + path_study: Study source path + dest: Destination path. + outputs: List of outputs to keep. + output_list_filter: List of outputs to keep (None indicate all outputs). + output_src_path: Denormalize the study (replace matrix links by real matrices). + + Returns: None + + """ + start_time = time.time() + output_src_path = output_src_path or path_study / "output" + output_dest_path = dest / "output" + ignore_patterns = ( + lambda directory, contents: ["output"] + if str(directory) == str(path_study) + else [] + ) + + shutil.copytree(src=path_study, dst=dest, ignore=ignore_patterns) + + if outputs and output_src_path.is_dir(): + if output_dest_path.exists(): + shutil.rmtree(output_dest_path) + if output_list_filter is not None: + os.mkdir(output_dest_path) + for output in output_list_filter: + zip_path = output_src_path / f"{output}.zip" + if zip_path.exists(): + with ZipFile(zip_path) as zf: + zf.extractall(output_dest_path / output) + else: + shutil.copytree( + src=output_src_path / output, + dst=output_dest_path / output, + ) + else: + shutil.copytree( + src=output_src_path, + dst=output_dest_path, + ) + + stop_time = time.time() + duration = "{:.3f}".format(stop_time - start_time) + logger.info(f"Study {path_study} exported (flat mode) in {duration}s") diff --git a/antarest/study/storage/rawstudy/raw_study_service.py b/antarest/study/storage/rawstudy/raw_study_service.py index ad79de032b..4981a63538 100644 --- a/antarest/study/storage/rawstudy/raw_study_service.py +++ b/antarest/study/storage/rawstudy/raw_study_service.py @@ -44,7 +44,6 @@ is_managed, remove_from_cache, create_new_empty_study, - export_study_flat, ) logger = logging.getLogger(__name__) @@ -362,31 +361,6 @@ def import_study(self, metadata: RawStudy, stream: IO[bytes]) -> Study: metadata.path = str(path_study) return metadata - def export_study_flat( - self, - metadata: RawStudy, - dst_path: Path, - outputs: bool = True, - output_list_filter: Optional[List[str]] = None, - denormalize: bool = True, - ) -> None: - path_study = Path(metadata.path) - - if metadata.archived: - self.unarchive(metadata) - try: - export_study_flat( - path_study, - dst_path, - self.study_factory, - outputs, - output_list_filter, - denormalize, - ) - finally: - if metadata.archived: - shutil.rmtree(metadata.path) - def check_errors( self, metadata: RawStudy, diff --git a/antarest/study/storage/utils.py b/antarest/study/storage/utils.py index 3c9ee5d901..e095917feb 100644 --- a/antarest/study/storage/utils.py +++ b/antarest/study/storage/utils.py @@ -367,55 +367,3 @@ def get_start_date( first_week_size=first_week_size, level=level, ) - - -def export_study_flat( - path_study: Path, - dest: Path, - study_factory: StudyFactory, - outputs: bool = True, - output_list_filter: Optional[List[str]] = None, - denormalize: bool = True, - output_src_path: Optional[Path] = None, -) -> None: - start_time = time.time() - - output_src_path = output_src_path or path_study / "output" - output_dest_path = dest / "output" - ignore_patterns = ( - lambda directory, contents: ["output"] - if str(directory) == str(path_study) - else [] - ) - - shutil.copytree(src=path_study, dst=dest, ignore=ignore_patterns) - - if outputs and output_src_path.is_dir(): - if output_dest_path.is_dir(): - shutil.rmtree(output_dest_path) - if output_list_filter is not None: - os.mkdir(output_dest_path) - for output in output_list_filter: - zip_path = output_src_path / f"{output}.zip" - if zip_path.exists(): - with ZipFile(zip_path) as zf: - zf.extractall(output_dest_path / output) - else: - shutil.copytree( - src=output_src_path / output, - dst=output_dest_path / output, - ) - else: - shutil.copytree( - src=output_src_path, - dst=output_dest_path, - ) - - stop_time = time.time() - duration = "{:.3f}".format(stop_time - start_time) - logger.info(f"Study {path_study} exported (flat mode) in {duration}s") - study = study_factory.create_from_fs(dest, "", use_cache=False) - if denormalize: - study.tree.denormalize() - duration = "{:.3f}".format(time.time() - stop_time) - logger.info(f"Study {path_study} denormalized in {duration}s") diff --git a/antarest/study/storage/variantstudy/variant_study_service.py b/antarest/study/storage/variantstudy/variant_study_service.py index 886932a740..b889f8aef0 100644 --- a/antarest/study/storage/variantstudy/variant_study_service.py +++ b/antarest/study/storage/variantstudy/variant_study_service.py @@ -70,7 +70,6 @@ from antarest.study.storage.rawstudy.raw_study_service import RawStudyService from antarest.study.storage.utils import ( assert_permission, - export_study_flat, get_default_workspace_path, is_managed, remove_from_cache, @@ -839,22 +838,30 @@ def _generate( last_executed_command_index = None if last_executed_command_index is None: - # Copy parent study to destination if isinstance(parent_study, VariantStudy): self._safe_generation(parent_study) + path_study = Path(parent_study.path) + snapshot_path = path_study / SNAPSHOT_RELATIVE_PATH + output_src_path = path_study / "output" self.export_study_flat( - metadata=parent_study, - dst_path=dst_path, + snapshot_path, + dst_path, outputs=False, - denormalize=False, + output_src_path=output_src_path, ) else: - self.raw_study_service.export_study_flat( - metadata=parent_study, - dst_path=dst_path, - outputs=False, - denormalize=False, - ) + path_study = Path(parent_study.path) + if parent_study.archived: + self.raw_study_service.unarchive(parent_study) + try: + self.raw_study_service.export_study_flat( + path_study=path_study, + dest=dst_path, + outputs=False, + ) + finally: + if parent_study.archived: + shutil.rmtree(parent_study.path) command_start_index = ( last_executed_command_index + 1 @@ -1234,29 +1241,6 @@ def get_study_path(self, metadata: Study) -> Path: """ return Path(metadata.path) / SNAPSHOT_RELATIVE_PATH - def export_study_flat( - self, - metadata: VariantStudy, - dst_path: Path, - outputs: bool = True, - output_list_filter: Optional[List[str]] = None, - denormalize: bool = True, - ) -> None: - self._safe_generation(metadata) - path_study = Path(metadata.path) - - snapshot_path = path_study / SNAPSHOT_RELATIVE_PATH - output_src_path = path_study / "output" - export_study_flat( - snapshot_path, - dst_path, - self.study_factory, - outputs, - output_list_filter, - denormalize, - output_src_path, - ) - def get_synthesis( self, metadata: VariantStudy,