From b85bfdc2e64092d2616f3ed3d760845bae43330e Mon Sep 17 00:00:00 2001 From: TLAIDI Date: Wed, 26 Jul 2023 14:45:15 +0200 Subject: [PATCH] refactor(export): request for Study export function (#1669) Co-authored-by: LAIDI Takfarinas (Externe) Co-authored-by: TLAIDI --- antarest/study/common/studystorage.py | 24 +------ antarest/study/service.py | 47 +++++++++++-- .../study/storage/abstract_storage_service.py | 70 ++++++++++++++++++- .../storage/rawstudy/raw_study_service.py | 26 ------- antarest/study/storage/utils.py | 52 -------------- .../variantstudy/variant_study_service.py | 60 ++++++---------- tests/storage/business/test_export.py | 11 +-- tests/storage/integration/test_exporter.py | 6 +- .../variantstudy/model/test_variant_model.py | 24 ++++--- 9 files changed, 155 insertions(+), 165 deletions(-) diff --git a/antarest/study/common/studystorage.py b/antarest/study/common/studystorage.py index 7f7c1193c0..85144d60af 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 @@ -292,3 +271,6 @@ def unarchive_study_output( self, study: T, output_id: str, keep_src_zip: bool ) -> bool: raise NotImplementedError() + + def unarchive(self, study: T) -> None: + raise NotImplementedError() diff --git a/antarest/study/service.py b/antarest/study/service.py index 200a32123b..d09a169c6a 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 @@ -165,6 +166,7 @@ remove_from_cache, study_matcher, ) +from antarest.study.storage.abstract_storage_service import export_study_flat from antarest.study.storage.variantstudy.model.command.icommand import ICommand from antarest.study.storage.variantstudy.model.command.replace_matrix import ( ReplaceMatrix, @@ -1089,9 +1091,23 @@ 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 isinstance(target_study, RawStudy): + 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 +1217,28 @@ 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 isinstance(study, RawStudy): + if study.archived: + self.storage_service.get_storage(study).unarchive(study) + try: + return 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" + 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..7fd85d52a5 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 @@ -39,6 +43,7 @@ StudyFactory, FileStudy, ) +from antarest.study.model import RawStudy from antarest.study.storage.rawstudy.model.helpers import FileStudyHelpers from antarest.study.storage.utils import ( fix_study_root, @@ -49,6 +54,60 @@ logger = logging.getLogger(__name__) +def export_study_flat( + 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). + + """ + 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") + + class AbstractStorageService(IStudyStorageService[T], ABC): def __init__( self, @@ -272,7 +331,16 @@ 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 not isinstance(metadata, RawStudy): + snapshot_path = path_study / "snapshot" + output_src_path = path_study / "output" + export_study_flat( + path_study=snapshot_path, + dest=tmp_study_path, + outputs=outputs, + output_src_path=output_src_path, + ) + export_study_flat(path_study, tmp_study_path, outputs) stopwatch = StopWatch() zip_dir(tmp_study_path, target) stopwatch.log_elapsed( 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..9d8ea14dfe 100644 --- a/antarest/study/storage/variantstudy/variant_study_service.py +++ b/antarest/study/storage/variantstudy/variant_study_service.py @@ -57,6 +57,7 @@ ) from antarest.study.storage.abstract_storage_service import ( AbstractStorageService, + export_study_flat, ) from antarest.study.storage.patch_service import PatchService from antarest.study.storage.rawstudy.model.filesystem.config.model import ( @@ -70,7 +71,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, @@ -80,9 +80,8 @@ ) from antarest.study.storage.variantstudy.command_factory import CommandFactory from antarest.study.storage.variantstudy.model.command.icommand import ICommand -from antarest.study.storage.variantstudy.model.command.update_config import ( - UpdateConfig, -) + + from antarest.study.storage.variantstudy.model.dbmodel import ( CommandBlock, VariantStudy, @@ -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) - self.export_study_flat( - metadata=parent_study, - dst_path=dst_path, + path_study = Path(parent_study.path) + snapshot_path = path_study / SNAPSHOT_RELATIVE_PATH + output_src_path = path_study / "output" + export_study_flat( + 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: + 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, diff --git a/tests/storage/business/test_export.py b/tests/storage/business/test_export.py index df4697a004..1764807418 100644 --- a/tests/storage/business/test_export.py +++ b/tests/storage/business/test_export.py @@ -8,6 +8,7 @@ from antarest.core.config import Config, StorageConfig from antarest.study.model import DEFAULT_WORKSPACE_NAME, RawStudy from antarest.study.storage.rawstudy.raw_study_service import RawStudyService +from antarest.study.storage.abstract_storage_service import export_study_flat @pytest.mark.unit_test @@ -105,17 +106,17 @@ def test_export_flat(tmp_path: Path): study_factory.create_from_fs.return_value = study_tree study = RawStudy(id="id", path=root) + path_study = Path(study.path) - study_service.export_study_flat( - study, tmp_path / "copy_with_output", outputs=True - ) + export_study_flat(path_study, tmp_path / "copy_with_output", outputs=True) copy_with_output_hash = dirhash(tmp_path / "copy_with_output", "md5") assert root_hash == copy_with_output_hash - study_service.export_study_flat( - study, tmp_path / "copy_without_output", outputs=False + path_study = Path(study.path) + export_study_flat( + path_study, tmp_path / "copy_without_output", outputs=False ) copy_without_output_hash = dirhash(tmp_path / "copy_without_output", "md5") diff --git a/tests/storage/integration/test_exporter.py b/tests/storage/integration/test_exporter.py index 355d5bf462..9a83fd212c 100644 --- a/tests/storage/integration/test_exporter.py +++ b/tests/storage/integration/test_exporter.py @@ -16,7 +16,7 @@ from antarest.matrixstore.service import MatrixService from antarest.study.main import build_study_service from antarest.study.model import DEFAULT_WORKSPACE_NAME, RawStudy -from antarest.study.storage.utils import export_study_flat +from antarest.study.storage.abstract_storage_service import export_study_flat from antarest.study.storage.variantstudy.business.matrix_constants_generator import ( GeneratorMatrixConstants, ) @@ -103,13 +103,11 @@ def test_exporter_file_no_output( @pytest.mark.parametrize( "output_list", [None, [], ["20201014-1427eco"], ["20201014-1430adq-2"]] ) -@pytest.mark.parametrize("denormalize", [True, False]) def test_export_flat( tmp_path: Path, sta_mini_zip_path: Path, outputs: bool, output_list: Optional[List[str]], - denormalize: bool, ) -> None: path_studies = tmp_path / "studies" path_studies.mkdir(exist_ok=True) @@ -123,10 +121,8 @@ def test_export_flat( export_study_flat( path_studies / "STA-mini", export_path / "STA-mini-export", - Mock(), outputs, output_list, - denormalize=denormalize, ) export_output_path = export_path / "STA-mini-export" / "output" diff --git a/tests/variantstudy/model/test_variant_model.py b/tests/variantstudy/model/test_variant_model.py index e254633ee0..801a86a497 100644 --- a/tests/variantstudy/model/test_variant_model.py +++ b/tests/variantstudy/model/test_variant_model.py @@ -1,6 +1,6 @@ import datetime from pathlib import Path -from unittest.mock import ANY, Mock +from unittest.mock import ANY, Mock, patch from antarest.core.cache.business.local_chache import LocalCache from antarest.core.config import Config, StorageConfig, WorkspaceConfig @@ -15,7 +15,6 @@ StudyAdditionalData, ) from antarest.study.storage.variantstudy.command_factory import CommandFactory -from antarest.study.storage.variantstudy.model.dbmodel import VariantStudy from antarest.study.storage.variantstudy.model.model import ( CommandDTO, GenerationResultInfoDTO, @@ -81,6 +80,7 @@ def test_commands_service(tmp_path: Path, command_factory: CommandFactory): id=origin_id, name="my-study", additional_data=StudyAdditionalData(), + path=str(tmp_path), ) repository.save(origin_study) @@ -166,8 +166,11 @@ def test_commands_service(tmp_path: Path, command_factory: CommandFactory): assert study.snapshot.id == study.id +@patch( + "antarest.study.storage.variantstudy.variant_study_service.export_study_flat" +) def test_smart_generation( - tmp_path: Path, command_factory: CommandFactory + mock_export, tmp_path: Path, command_factory: CommandFactory ) -> None: engine = create_engine( "sqlite:///:memory:", @@ -209,17 +212,16 @@ def test_smart_generation( # noinspection PyUnusedLocal def export_flat( - metadata: VariantStudy, - dst_path: Path, + path_study: Path, + dest: Path, outputs: bool = True, denormalize: bool = True, ) -> None: - dst_path.mkdir(parents=True) - (dst_path / "user").mkdir() - (dst_path / "user" / "some_unmanaged_config").touch() - - service.raw_study_service.export_study_flat.side_effect = export_flat + dest.mkdir(parents=True) + (dest / "user").mkdir() + (dest / "user" / "some_unmanaged_config").touch() + mock_export.side_effect = export_flat with db(): origin_id = "base-study" # noinspection PyArgumentList @@ -230,6 +232,7 @@ def export_flat( workspace=DEFAULT_WORKSPACE_NAME, additional_data=StudyAdditionalData(), updated_at=datetime.datetime(year=2000, month=1, day=1), + path=str(tmp_path), ) repository.save(origin_study) @@ -293,7 +296,6 @@ def export_flat( ], SADMIN, ) - assert unmanaged_user_config_path.exists() unmanaged_user_config_path.write_text("hello") service._generate(variant_id, SADMIN, False)