Fix nested root zipped import (#285)

* Fix nested root zipped import

Signed-off-by: Paul Bui-Quang <[email protected]>

* Fix mypy

Signed-off-by: Paul Bui-Quang <[email protected]>

antarest/storage/business/importer_service.py

@@ -1,19 +1,23 @@
 import json
+import logging
 import os
 import shutil
 from datetime import datetime
 from pathlib import Path
 from typing import IO
+from uuid import uuid4
 
 from antarest.common.custom_types import JSON
 from antarest.storage.business.storage_service_utils import StorageServiceUtils
 from antarest.storage.business.study_service import StudyService
 from antarest.storage.repository.antares_io.reader import IniReader
 from antarest.storage.repository.filesystem.factory import StudyFactory
-from antarest.common.requests import (
-    RequestParameters,
+from antarest.storage.web.exceptions import (
+    BadOutputError,
+    StudyValidationError,
 )
-from antarest.storage.web.exceptions import BadOutputError
+
+logger = logging.getLogger(__name__)
 
 
 class ImporterService:
@@ -43,26 +47,33 @@ def import_study(self, stream: IO[bytes]) -> str:
         uuid = StorageServiceUtils.generate_uuid()
         path_study = Path(self.path_to_studies) / uuid
         path_study.mkdir()
-        StorageServiceUtils.extract_zip(stream, path_study)
-
-        data_file = path_study / "data.json"
-
-        # If compact study generate tree and launch save with data.json
-        if data_file.is_file() and (path_study / "res").is_dir():
-            with open(data_file) as file:
-                data = json.load(file)
-                _, study = self.study_factory.create_from_json(
-                    path_study, data
-                )
-                study.save(data)
-            del study
-            shutil.rmtree(path_study / "res")
-            os.remove(str(data_file.absolute()))
-
-        data = self.study_service.get(uuid, -1)
-        if data is None:
-            self.study_service.delete_study(uuid)
-            return ""  # TODO return exception
+
+        try:
+            StorageServiceUtils.extract_zip(stream, path_study)
+
+            data_file = path_study / "data.json"
+
+            # If compact study generate tree and launch save with data.json
+            if data_file.is_file() and (path_study / "res").is_dir():
+                with open(data_file) as file:
+                    data = json.load(file)
+                    _, study = self.study_factory.create_from_json(
+                        path_study, data
+                    )
+                    study.save(data)
+                del study
+                shutil.rmtree(path_study / "res")
+                os.remove(str(data_file.absolute()))
+            else:
+                fix_study_root(path_study)
+
+            data = self.study_service.get(uuid, -1)
+            if data is None:
+                self.study_service.delete_study(uuid)
+                raise StudyValidationError("Fail to import study")
+        except Exception as e:
+            shutil.rmtree(path_study)
+            raise e
 
         return uuid
 
@@ -106,3 +117,34 @@ def import_output(self, uuid: str, stream: IO[bytes]) -> JSON:
             raise BadOutputError("The output provided is not conform.")
 
         return data
+
+
+def fix_study_root(study_path: Path) -> None:
+    """
+    Fix possibly the wrong study root on zipped archive (when the study root is nested)
+
+    @param study_path the study initial root path
+    """
+    if not study_path.is_dir():
+        raise StudyValidationError("Not a directory")
+
+    root_path = study_path
+    contents = os.listdir(root_path)
+    sub_root_path = None
+    while len(contents) == 1 and (root_path / contents[0]).is_dir():
+        new_root = root_path / contents[0]
+        if sub_root_path is None:
+            sub_root_path = root_path / str(uuid4())
+            shutil.move(str(new_root), str(sub_root_path))
+            new_root = sub_root_path
+
+        logger.debug(f"Searching study root in {new_root}")
+        root_path = new_root
+        if not new_root.is_dir():
+            raise StudyValidationError("Not a directory")
+        contents = os.listdir(new_root)
+
+    if sub_root_path is not None:
+        for item in os.listdir(root_path):
+            shutil.move(str(root_path / item), str(study_path))
+        shutil.rmtree(sub_root_path)

tests/storage/business/test_importer_service.py

@@ -1,14 +1,19 @@
+import os
 import shutil
 from pathlib import Path
 import io
 from unittest.mock import Mock
 
 import pytest
 
-from antarest.storage.business.importer_service import ImporterService
+from antarest.storage.business.importer_service import (
+    ImporterService,
+    fix_study_root,
+)
 from antarest.storage.web.exceptions import (
     IncorrectPathError,
     BadZipBinary,
+    StudyValidationError,
 )
 
 
@@ -67,3 +72,29 @@ def test_import_study(tmp_path: Path, storage_service_builder) -> None:
 
     with pytest.raises(BadZipBinary):
         importer_service.import_study(io.BytesIO(b""))
+
+
+@pytest.mark.unit_test
+def test_fix_root(tmp_path: Path):
+    name = "my-study"
+    study_path = tmp_path / name
+    study_nested_root = study_path / "nested" / "real_root"
+    os.makedirs(study_nested_root)
+    (study_nested_root / "antares.study").touch()
+    # when the study path is a single file
+    with pytest.raises(StudyValidationError):
+        fix_study_root(study_nested_root / "antares.study")
+
+    shutil.rmtree(study_path)
+    study_path = tmp_path / name
+    study_nested_root = study_path / "nested" / "real_root"
+    os.makedirs(study_nested_root)
+    (study_nested_root / "antares.study").touch()
+    os.mkdir(study_nested_root / "input")
+
+    fix_study_root(study_path)
+    study_files = os.listdir(study_path)
+    assert len(study_files) == 2
+    assert "antares.study" in study_files and "input" in study_files
+
+    shutil.rmtree(study_path)