diff --git a/antarest/study/common/studystorage.py b/antarest/study/common/studystorage.py index 162266ebb2..f795b46d6c 100644 --- a/antarest/study/common/studystorage.py +++ b/antarest/study/common/studystorage.py @@ -230,27 +230,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) -> FileStudyTreeConfigDTO: """ @@ -274,3 +253,6 @@ def archive_study_output(self, study: T, output_id: str) -> bool: @abstractmethod 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 332fc384da..1523434822 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 @@ -106,6 +107,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 from antarest.study.storage.variantstudy.model.command.update_comments import UpdateComments @@ -918,7 +920,17 @@ 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) + storage = self.storage_service.get_storage(target_study) + if isinstance(target_study, RawStudy): + if target_study.archived: + storage.unarchive(target_study) + try: + storage.export_study(target_study, export_path, outputs) + finally: + if target_study.archived: + shutil.rmtree(target_study.path) + else: + storage.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") except Exception as e: @@ -1020,9 +1032,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(self, uuid: str, children: bool, params: RequestParameters) -> None: diff --git a/antarest/study/storage/abstract_storage_service.py b/antarest/study/storage/abstract_storage_service.py index f056839ac6..bc466580cd 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 IO, List, Optional, Union 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 @@ -32,6 +36,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, @@ -230,7 +288,16 @@ def export_study(self, metadata: T, target: Path, outputs: bool = True) -> Path: with tempfile.TemporaryDirectory(dir=self.config.storage.tmp_dir) as tmpdir: logger.info(f"Exporting study {metadata.id} to temporary path {tmpdir}") 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(lambda x: logger.info(f"Study {path_study} exported (zipped mode) in {x}s")) diff --git a/antarest/study/storage/rawstudy/raw_study_service.py b/antarest/study/storage/rawstudy/raw_study_service.py index c6c73e63b5..c9ce8fd0d1 100644 --- a/antarest/study/storage/rawstudy/raw_study_service.py +++ b/antarest/study/storage/rawstudy/raw_study_service.py @@ -23,7 +23,6 @@ from antarest.study.storage.rawstudy.model.filesystem.lazy_node import LazyNode from antarest.study.storage.utils import ( create_new_empty_study, - export_study_flat, fix_study_root, get_default_workspace_path, is_managed, @@ -323,31 +322,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 6197180f50..d0906e128c 100644 --- a/antarest/study/storage/utils.py +++ b/antarest/study/storage/utils.py @@ -326,50 +326,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.exists(): - if output_list_filter is None: - # Retrieve all directories or ZIP files without duplicates - output_list_filter = list( - {f.with_suffix("").name for f in output_src_path.iterdir() if f.is_dir() or f.suffix == ".zip"} - ) - # Copy each folder or uncompress each ZIP file to the destination dir. - shutil.rmtree(output_dest_path, ignore_errors=True) - output_dest_path.mkdir() - 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, - ) - - 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 f956564656..61ff09f025 100644 --- a/antarest/study/storage/variantstudy/variant_study_service.py +++ b/antarest/study/storage/variantstudy/variant_study_service.py @@ -34,14 +34,13 @@ from antarest.core.utils.utils import assert_this, suppress_exception from antarest.matrixstore.service import MatrixService from antarest.study.model import RawStudy, Study, StudyAdditionalData, StudyMetadataDTO, StudySimResultDTO -from antarest.study.storage.abstract_storage_service import AbstractStorageService +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 FileStudyTreeConfig, FileStudyTreeConfigDTO from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy, StudyFactory 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, @@ -719,22 +718,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 if last_executed_command_index is not None else 0 logger.info(f"Generating study snapshot from command index {command_start_index}") @@ -1068,29 +1075,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 759bc99ed3..e45f68dfcd 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 @@ -104,17 +105,14 @@ 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) - + export_study_flat(path_study, tmp_path / "copy_without_output", outputs=False) copy_without_output_hash = dirhash(tmp_path / "copy_without_output", "md5") - assert root_without_output_hash == copy_without_output_hash diff --git a/tests/storage/integration/test_exporter.py b/tests/storage/integration/test_exporter.py index 3e4e5666f3..be3d5c5c9f 100644 --- a/tests/storage/integration/test_exporter.py +++ b/tests/storage/integration/test_exporter.py @@ -100,13 +100,11 @@ def test_exporter_file_no_output(tmp_path: Path, sta_mini_zip_path: Path) -> Non @pytest.mark.parametrize("outputs", [True, False, "prout"]) @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) @@ -120,10 +118,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 d3a1760077..d1cebd2ce9 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 sqlalchemy import create_engine @@ -13,7 +13,6 @@ from antarest.core.utils.fastapi_sqlalchemy import DBSessionMiddleware, db from antarest.study.model import DEFAULT_WORKSPACE_NAME, RawStudy, 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 from antarest.study.storage.variantstudy.repository import VariantStudyRepository from antarest.study.storage.variantstudy.variant_study_service import SNAPSHOT_RELATIVE_PATH, VariantStudyService @@ -64,6 +63,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) @@ -138,7 +138,8 @@ def test_commands_service(tmp_path: Path, command_factory: CommandFactory): assert study.snapshot.id == study.id -def test_smart_generation(tmp_path: Path, command_factory: CommandFactory) -> None: +@patch("antarest.study.storage.variantstudy.variant_study_service.export_study_flat") +def test_smart_generation(mock_export, tmp_path: Path, command_factory: CommandFactory) -> None: engine = create_engine( "sqlite:///:memory:", echo=False, @@ -173,17 +174,16 @@ def test_smart_generation(tmp_path: Path, command_factory: CommandFactory) -> No # 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 @@ -194,6 +194,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) @@ -233,7 +234,6 @@ def export_flat( ], SADMIN, ) - assert unmanaged_user_config_path.exists() unmanaged_user_config_path.write_text("hello") service._generate(variant_id, SADMIN, False)