From ef009c97c65df69e86e6aea90cbb02023e3fe4a8 Mon Sep 17 00:00:00 2001
From: TLAIDI <takfarinas.laidi@gmail.com>
Date: Wed, 26 Jul 2023 14:45:15 +0200
Subject: [PATCH] Features/1646 refactoring request for study function (#1669)

Co-authored-by: LAIDI Takfarinas (Externe) <laiditak@gm0winl638.bureau.si.interne>
Co-authored-by: TLAIDI <takfarinas.laidi_externe@rte-france.com>
---
 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)