Merge pull request #10191 from RasaHQ/clean-up-rasa-data-migrate

`rasa data migrate` doesn't clean up if migration fails

rasa/core/migrate.py

@@ -1,4 +1,5 @@
 import copy
+import shutil
 from pathlib import Path
 from typing import List, Dict, Text, Any, Tuple, Optional, Union
 
@@ -180,75 +181,90 @@ def _write_final_domain(
         rasa.shared.utils.io.write_yaml(new_domain, out_file, True)
 
 
+def _migrate_domain_files(
+    domain_file: Path, backup_dir: Path, out_file: Path
+) -> Dict[Text, Any]:
+    slots = {}
+    forms = {}
+    entities = []
+
+    for file in domain_file.iterdir():
+        if not Domain.is_domain_file(file):
+            continue
+
+        backup = backup_dir / file.name
+        original_content = _create_back_up(file, backup)
+
+        if KEY_SLOTS not in original_content and KEY_FORMS not in original_content:
+            # this is done so that the other domain files can be moved
+            # in the migrated directory
+            rasa.shared.utils.io.write_yaml(
+                original_content, out_file / file.name, True
+            )
+        elif KEY_SLOTS in original_content and slots:
+            raise RasaException(
+                f"Domain files with multiple '{KEY_SLOTS}' "
+                f"sections were provided. Please group these sections "
+                f"in one file only to prevent content duplication across "
+                f"multiple files. "
+            )
+        elif KEY_FORMS in original_content and forms:
+            raise RasaException(
+                f"Domain files with multiple '{KEY_FORMS}' "
+                f"sections were provided. Please group these sections "
+                f"in one file only to prevent content duplication across "
+                f"multiple files. "
+            )
+
+        slots.update(original_content.get(KEY_SLOTS, {}))
+        forms.update(original_content.get(KEY_FORMS, {}))
+        entities.extend(original_content.get(KEY_ENTITIES, {}))
+
+    if not slots or not forms:
+        raise RasaException(
+            f"The files you have provided in '{domain_file}' are missing slots "
+            f"or forms. Please make sure to include these for a "
+            f"successful migration."
+        )
+
+    return {KEY_SLOTS: slots, KEY_FORMS: forms, KEY_ENTITIES: entities}
+
+
 def migrate_domain_format(
     domain_file: Union[Text, Path], out_file: Union[Text, Path]
 ) -> None:
     """Converts 2.0 domain to 3.0 format."""
     domain_file = Path(domain_file)
     out_file = Path(out_file)
+    created_out_dir = False
 
     current_dir = domain_file.parent
 
     if domain_file.is_dir():
         backup_dir = current_dir / "original_domain"
         backup_dir.mkdir()
 
-        slots = {}
-        forms = {}
-        entities = []
-        original_domain = {}
-
         if out_file.is_file() or not out_file.exists():
             out_file = current_dir / "new_domain"
             out_file.mkdir()
+            created_out_dir = True
             rasa.shared.utils.io.raise_warning(
                 f"The out path provided is not a directory, "
                 f"creating a new directory '{str(out_file)}' "
                 f"for migrated domain files."
             )
 
-        for file in domain_file.iterdir():
-            if not Domain.is_domain_file(file):
-                continue
-
-            backup = backup_dir / file.name
-            original_content = _create_back_up(file, backup)
-
-            if KEY_SLOTS not in original_content and KEY_FORMS not in original_content:
-                # this is done so that the other domain files can be moved
-                # in the migrated directory
-                rasa.shared.utils.io.write_yaml(
-                    original_content, out_file / file.name, True
-                )
-            elif KEY_SLOTS in original_content and slots:
-                raise RasaException(
-                    f"Domain files with multiple '{KEY_SLOTS}' "
-                    f"sections were provided. Please group these sections "
-                    f"in one file only to prevent content duplication across "
-                    f"multiple files. "
-                )
-            elif KEY_FORMS in original_content and forms:
-                raise RasaException(
-                    f"Domain files with multiple '{KEY_FORMS}' "
-                    f"sections were provided. Please group these sections "
-                    f"in one file only to prevent content duplication across "
-                    f"multiple files. "
-                )
-
-            slots.update(original_content.get(KEY_SLOTS, {}))
-            forms.update(original_content.get(KEY_FORMS, {}))
-            entities.extend(original_content.get(KEY_ENTITIES, {}))
-
-        if not slots or not forms:
-            raise RasaException(
-                f"The files you have provided in '{domain_file}' are missing slots "
-                f"or forms. Please make sure to include these for a "
-                f"successful migration."
-            )
-
-        original_domain.update(
-            {KEY_SLOTS: slots, KEY_FORMS: forms, KEY_ENTITIES: entities}
-        )
+        try:
+            original_domain = _migrate_domain_files(domain_file, backup_dir, out_file)
+        except Exception as e:
+            shutil.rmtree(backup_dir)
+            if out_file != domain_file:
+                shutil.rmtree(out_file)
+                if not created_out_dir:
+                    # we recreate the deleted directory
+                    # because it existed before
+                    out_file.mkdir()
+            raise e
     else:
         if not Domain.is_domain_file(domain_file):
             raise RasaException(

tests/core/test_migrate.py

@@ -1,3 +1,4 @@
+import os
 import re
 import textwrap
 from pathlib import Path
@@ -764,3 +765,88 @@ def test_migrate_domain_raises_for_non_domain_files(tmp_path: Path):
         f"Please make sure to include these for a successful migration.",
     ):
         rasa.core.migrate.migrate_domain_format(domain_dir, domain_dir)
+
+
+def test_migration_cleanup(tmp_path: Path,):
+    domain_dir = tmp_path / "domain"
+    domain_dir.mkdir()
+    migrated_domain_dir = tmp_path / "domain2"
+
+    prepare_domain_path(
+        domain_dir,
+        """
+        version: "2.0"
+        entities:
+            - outdoor
+        slots:
+          outdoor_seating:
+           type: bool
+           influence_conversation: false
+        """,
+        "slots_one.yml",
+    )
+
+    prepare_domain_path(
+        domain_dir,
+        """
+        version: "2.0"
+        entities:
+            - outdoor
+        slots:
+          cuisine:
+           type: text
+           influence_conversation: false
+        """,
+        "slots_two.yml",
+    )
+
+    prepare_domain_path(
+        domain_dir,
+        """
+        version: "2.0"
+        responses:
+          utter_greet:
+          - text: "Hi there!"
+        """,
+        "responses.yml",
+    )
+
+    with pytest.raises(
+        RasaException,
+        match="Domain files with multiple 'slots' sections were provided.",
+    ):
+        rasa.core.migrate.migrate_domain_format(domain_dir, domain_dir)
+
+    # here the migration fails but we check if the domain was
+    # not deleted in the process of cleanup
+    assert Path.exists(domain_dir)
+
+    with pytest.raises(
+        RasaException,
+        match="Domain files with multiple 'slots' sections were provided.",
+    ):
+        rasa.core.migrate.migrate_domain_format(domain_dir, migrated_domain_dir)
+
+    # the migration fails again and we check
+    # if the cleanup was done successfully
+    current_dir = domain_dir.parent
+    assert Path.exists(domain_dir)
+    assert not Path.exists(current_dir / "original_domain")
+    assert not Path.exists(current_dir / "new_domain")
+
+    # create the `migrated_domain_dir` to use it instead of the default one
+    migrated_domain_dir.mkdir()
+
+    with pytest.raises(
+        RasaException,
+        match="Domain files with multiple 'slots' sections were provided.",
+    ):
+        rasa.core.migrate.migrate_domain_format(domain_dir, migrated_domain_dir)
+
+    # check if the cleanup was done successfully after
+    # the `migrated_domain_dir` was created
+    assert (
+        Path.exists(migrated_domain_dir) and len(os.listdir(migrated_domain_dir)) == 0
+    )
+    assert not Path.exists(current_dir / "original_domain")
+    assert not Path.exists(current_dir / "new_domain")