diff --git a/antarest/study/storage/study_upgrader/__init__.py b/antarest/study/storage/study_upgrader/__init__.py index 6fbfcfa29b..f5e64a87e6 100644 --- a/antarest/study/storage/study_upgrader/__init__.py +++ b/antarest/study/storage/study_upgrader/__init__.py @@ -6,7 +6,7 @@ from http import HTTPStatus from http.client import HTTPException from pathlib import Path -from typing import Callable, NamedTuple +from typing import Callable, List, NamedTuple from antarest.core.exceptions import StudyValidationError @@ -29,18 +29,34 @@ class UpgradeMethod(NamedTuple): old: str new: str method: Callable[[Path], None] + files: List[Path] UPGRADE_METHODS = [ - UpgradeMethod("700", "710", upgrade_710), - UpgradeMethod("710", "720", upgrade_720), - UpgradeMethod("720", "800", upgrade_800), - UpgradeMethod("800", "810", upgrade_810), - UpgradeMethod("810", "820", upgrade_820), - UpgradeMethod("820", "830", upgrade_830), - UpgradeMethod("830", "840", upgrade_840), - UpgradeMethod("840", "850", upgrade_850), - UpgradeMethod("850", "860", upgrade_860), + UpgradeMethod("700", "710", upgrade_710, [Path("settings/generaldata.ini")]), + UpgradeMethod("710", "720", upgrade_720, []), + UpgradeMethod("720", "800", upgrade_800, [Path("settings/generaldata.ini")]), + UpgradeMethod( + "800", + "810", + upgrade_810, + [Path("settings/generaldata.ini"), Path("input")], + ), + UpgradeMethod("810", "820", upgrade_820, [Path("input/links")]), + UpgradeMethod( + "820", + "830", + upgrade_830, + [Path("settings/generaldata.ini"), Path("input/areas")], + ), + UpgradeMethod("830", "840", upgrade_840, [Path("settings/generaldata.ini")]), + UpgradeMethod("840", "850", upgrade_850, [Path("settings/generaldata.ini")]), + UpgradeMethod( + "850", + "860", + upgrade_860, + [Path("input"), Path("settings/generaldata.ini")], + ), ] @@ -70,8 +86,9 @@ def upgrade_study(study_path: Path, target_version: str) -> None: tmp_dir = Path(tempfile.mkdtemp(suffix=".upgrade.tmp", prefix="~", dir=study_path.parent)) shutil.copytree(study_path, tmp_dir, dirs_exist_ok=True) try: - src_version = get_current_version(tmp_dir) - can_upgrade_version(src_version, target_version) + src_version = get_current_version(study_path) + files_to_upgrade = can_upgrade_version(src_version, target_version) + files_to_retrieve = _copies_only_necessary_files(files_to_upgrade, study_path, tmp_dir) _do_upgrade(tmp_dir, src_version, target_version) except (StudyValidationError, InvalidUpgrade) as e: shutil.rmtree(tmp_dir) @@ -82,11 +99,7 @@ def upgrade_study(study_path: Path, target_version: str) -> None: logger.error(f"Unhandled exception : {e}", exc_info=True) raise else: - backup_dir = Path(tempfile.mkdtemp(suffix=".backup.tmp", prefix="~", dir=study_path.parent)) - backup_dir.rmdir() - study_path.rename(backup_dir) - tmp_dir.rename(study_path) - shutil.rmtree(backup_dir, ignore_errors=True) + _replace_safely_original_files(files_to_retrieve, study_path, tmp_dir) def get_current_version(study_path: Path) -> str: @@ -116,7 +129,7 @@ def get_current_version(study_path: Path) -> str: ) -def can_upgrade_version(from_version: str, to_version: str) -> None: +def can_upgrade_version(from_version: str, to_version: str) -> List[Path]: """ Checks if upgrading from one version to another is possible. @@ -124,9 +137,13 @@ def can_upgrade_version(from_version: str, to_version: str) -> None: from_version: The current version of the study. to_version: The target version of the study. + Returns: + If the upgrade is possible, the list of concerned folders and files + Raises: InvalidUpgrade: If the upgrade is not possible. """ + list_versions = [] if from_version == to_version: raise InvalidUpgrade(f"Your study is already in version '{to_version}'") @@ -138,12 +155,16 @@ def can_upgrade_version(from_version: str, to_version: str) -> None: if to_version not in targets: raise InvalidUpgrade(f"Version '{to_version}' unknown: possible versions are {', '.join(targets)}") + files = [u.files for u in UPGRADE_METHODS] curr_version = from_version - for src, dst in zip(sources, targets): + for src, dst, file in zip(sources, targets, files): if curr_version == src: + for path in file: + if path not in list_versions: + list_versions.append(path) curr_version = dst if curr_version == to_version: - return + return list_versions # This code must be unreachable! raise InvalidUpgrade( @@ -171,10 +192,66 @@ def _update_study_antares_file(target_version: str, study_path: Path) -> None: file.write_text(content, encoding="utf-8") +def _copies_only_necessary_files(files_to_upgrade: List[Path], study_path: Path, tmp_path: Path) -> List[Path]: + """ + Copies files concerned by the version upgrader into a temporary directory. + Args: + study_path: Path to the study. + tmp_path: Path to the temporary directory where the file modification will be performed. + files_to_upgrade: List[Path]: List of the files and folders concerned by the upgrade. + Returns: + The list of files and folders that were really copied. It's the same as files_to_upgrade but + without any children that has parents already in the list. + """ + files_to_upgrade.append(Path("study.antares")) + files_to_retrieve = [] + for path in files_to_upgrade: + entire_path = study_path / path + if entire_path.is_dir(): + if not (tmp_path / path).exists(): + shutil.copytree(entire_path, tmp_path / path, dirs_exist_ok=True) + files_to_retrieve.append(path) + elif len(path.parts) == 1: + shutil.copy(entire_path, tmp_path / path) + files_to_retrieve.append(path) + else: + parent_path = path.parent + (tmp_path / parent_path).mkdir(parents=True) + shutil.copy(entire_path, tmp_path / parent_path) + files_to_retrieve.append(path) + return files_to_retrieve + + +def _replace_safely_original_files(files_to_replace: List[Path], study_path: Path, tmp_path: Path) -> None: + """ + Replace files/folders of the study that should be upgraded by their copy already upgraded in the tmp directory. + It uses Path.rename() and an intermediary tmp directory to swap the folders safely. + In the end, all tmp directories are removed. + Args: + study_path: Path to the study. + tmp_path: Path to the temporary directory where the file modification will be performed. + files_to_replace: List[Path]: List of files and folders that were really copied + (cf. _copies_only_necessary_files's doc just above) + """ + for k, path in enumerate(files_to_replace): + backup_dir = Path( + tempfile.mkdtemp( + suffix=f".backup_{k}.tmp", + prefix="~", + dir=study_path.parent, + ) + ) + backup_dir.rmdir() + original_path = study_path / path + original_path.rename(backup_dir) + (tmp_path / path).rename(original_path) + shutil.rmtree(backup_dir, ignore_errors=True) + + def _do_upgrade(study_path: Path, src_version: str, target_version: str) -> None: _update_study_antares_file(target_version, study_path) curr_version = src_version - for old, new, method in UPGRADE_METHODS: + for old, new, method, _ in UPGRADE_METHODS: if curr_version == old and curr_version != target_version: method(study_path) curr_version = new