Add "transitions" based execution machine

This adapts the logic how modules are loaded and sorted for execution.

Signed-off-by: Tobias Wolf <[email protected]>

requirements.txt

@@ -26,6 +26,7 @@ requests-oauthlib==1.3.1
 rsa==4.9
 six==1.16.0
 structlog==24.1.0
+transitions==0.9.0
 urllib3==2.2.1
 yamale==5.1.0
 websocket-client==1.7.0

src/rookify/__main__.py

@@ -1,10 +1,9 @@
 # -*- coding: utf-8 -*-
 
-import os
-import rookify.modules
-from types import MappingProxyType
+from .modules import load_modules
+from .modules.machine import Machine
 from .logger import configure_logging, get_logger
-from .yaml import load_config, load_module_data, save_module_data
+from .yaml import load_config
 
 
 def main() -> None:
@@ -23,56 +22,7 @@ def main() -> None:
     log = get_logger()
     log.debug("Executing Rookify")
 
-    preflight_modules, migration_modules = rookify.modules.load_modules(
-        config["migration_modules"]
-    )
+    machine = Machine()
+    load_modules(machine, config)
 
-    module_data = dict()
-
-    try:
-        module_data.update(load_module_data(config["general"]["module_data_file"]))
-    except FileNotFoundError:
-        pass
-
-    # Get absolute path of the rookify instance
-    rookify_path = os.path.dirname(__file__)
-
-    # Run preflight requirement modules
-    for preflight_module in preflight_modules:
-        module_path = os.path.join(
-            rookify_path, "modules", preflight_module.MODULE_NAME
-        )
-        handler = preflight_module.HANDLER_CLASS(
-            config=MappingProxyType(config),
-            data=MappingProxyType(module_data),
-            module_path=module_path,
-        )
-        result = handler.run()
-        module_data[preflight_module.MODULE_NAME] = result
-
-    # Run preflight and append handlers to list
-    handlers = list()
-    for migration_module in migration_modules:
-        module_path = os.path.join(
-            rookify_path, "modules", migration_module.MODULE_NAME
-        )
-        handler = migration_module.HANDLER_CLASS(
-            config=MappingProxyType(config),
-            data=MappingProxyType(module_data),
-            module_path=module_path,
-        )
-        handler.preflight()
-        handlers.append((migration_module, handler))
-
-    # Run migration modules
-    for migration_module, handler in handlers:
-        result = handler.run()
-        module_data[migration_module.MODULE_NAME] = result
-
-    save_module_data(config["general"]["module_data_file"], module_data)
-
-    log.info("Data was updated to module_data_file.")
-
-
-if __name__ == "__main__":
-    main()
+    machine.execute()

src/rookify/modules/__init__.py

@@ -1,11 +1,8 @@
 # -*- coding: utf-8 -*-
 
 import importlib
-import types
-from typing import List
-
-from collections import OrderedDict
-from .module import ModuleHandler
+from typing import Any, Dict
+from .machine import Machine
 
 
 class ModuleLoadException(Exception):
@@ -24,107 +21,46 @@ def __init__(self, module_name: str, message: str):
         self.message = message
 
 
-def load_modules(
-    module_names: List[str],
-) -> tuple[List[types.ModuleType], List[types.ModuleType]]:
+def _load_module(machine: Machine, config: Dict[str, Any], module_name: str) -> None:
+    """
+    Dynamically loads a module from the 'rookify.modules' package.
+
+    :param module_names: The module names to load
+    :return: returns tuple of preflight_modules, modules
+    """
+
+    module = importlib.import_module("rookify.modules.{0}".format(module_name))
+    additional_modules = []
+
+    if not hasattr(module, "ModuleHandler") or not callable(
+        getattr(module.ModuleHandler, "register_state")
+    ):
+        raise ModuleLoadException(module_name, "Module structure is invalid")
+
+    if hasattr(module.ModuleHandler, "PREFLIGHT_REQUIRES"):
+        assert isinstance(module.ModuleHandler.PREFLIGHT_REQUIRES, list)
+        additional_modules += module.ModuleHandler.PREFLIGHT_REQUIRES
+
+    if hasattr(module.ModuleHandler, "REQUIRES"):
+        assert isinstance(module.ModuleHandler.REQUIRES, list)
+        for module_name in module.ModuleHandler.REQUIRES:
+            if module_name not in additional_modules:
+                additional_modules.append(module_name)
+
+    for module_name in additional_modules:
+        _load_module(machine, config, module_name)
+
+    module.ModuleHandler.register_state(machine, config)
+
+
+def load_modules(machine: Machine, config: Dict[str, Any]) -> None:
     """
     Dynamically loads modules from the 'modules' package.
 
     :param module_names: The module names to load
     :return: returns tuple of preflight_modules, modules
     """
 
-    # Sanity checks for modules
-    def check_module_sanity(module_name: str, module: types.ModuleType) -> None:
-        for attr_type, attr_name in (
-            (ModuleHandler, "HANDLER_CLASS"),
-            (str, "MODULE_NAME"),
-            (list, "REQUIRES"),
-            (list, "AFTER"),
-            (list, "PREFLIGHT_REQUIRES"),
-        ):
-            if not hasattr(module, attr_name):
-                raise ModuleLoadException(
-                    module_name, f"Module has no attribute {attr_name}"
-                )
-
-            attr = getattr(module, attr_name)
-            if not isinstance(attr, attr_type) and not issubclass(attr, attr_type):
-                raise ModuleLoadException(
-                    module_name, f"Attribute {attr_name} is not type {attr_type}"
-                )
-
-    # Load the modules in the given list and recursivley load required modules
-    required_modules: OrderedDict[str, types.ModuleType] = OrderedDict()
-
-    def load_required_modules(
-        modules_out: OrderedDict[str, types.ModuleType], module_names: List[str]
-    ) -> None:
-        for module_name in module_names:
-            if module_name in modules_out:
-                continue
-
-            module = importlib.import_module(f".{module_name}", "rookify.modules")
-            check_module_sanity(module_name, module)
-
-            load_required_modules(modules_out, module.REQUIRES)
-            module.AFTER.extend(module.REQUIRES)
-
-            modules_out[module_name] = module
-
-    load_required_modules(required_modules, module_names)
-
-    # Recursively load the modules in the PREFLIGHT_REQUIRES attribute of the given modules
-    preflight_modules: OrderedDict[str, types.ModuleType] = OrderedDict()
-
-    def load_preflight_modules(
-        modules_in: OrderedDict[str, types.ModuleType],
-        modules_out: OrderedDict[str, types.ModuleType],
-        module_names: List[str],
-    ) -> None:
-        for module_name in module_names:
-            if module_name in modules_out:
-                continue
-
-            module = importlib.import_module(f".{module_name}", "rookify.modules")
-            check_module_sanity(module_name, module)
-
-            # We have to check, if the preflight_requires list is already loaded as migration requirement
-            for preflight_requirement in module.PREFLIGHT_REQUIRES:
-                if preflight_requirement in modules_in:
-                    raise ModuleLoadException(
-                        module_name,
-                        f"Module {preflight_requirement} is already loaded as migration requirement",
-                    )
-
-            load_preflight_modules(modules_in, modules_out, module.PREFLIGHT_REQUIRES)
-            if module_name not in modules_in:
-                modules_out[module_name] = module
-
-    load_preflight_modules(
-        required_modules, preflight_modules, list(required_modules.keys())
-    )
-
-    # Sort the modules by the AFTER keyword
-    modules: OrderedDict[str, types.ModuleType] = OrderedDict()
-
-    def sort_modules(
-        modules_in: OrderedDict[str, types.ModuleType],
-        modules_out: OrderedDict[str, types.ModuleType],
-        module_names: List[str],
-    ) -> None:
-        for module_name in module_names:
-            if module_name not in modules_in:
-                continue
-
-            if module_name in modules_out:
-                continue
-
-            after_modules_name = modules_in[module_name].AFTER
-            sort_modules(modules_in, modules_out, after_modules_name)
-
-            modules_out[module_name] = modules_in[module_name]
-
-    sort_modules(required_modules, modules, list(required_modules.keys()))
-
-    return list(preflight_modules.values()), list(modules.values())
+    for entry in importlib.resources.files("rookify.modules").iterdir():
+        if entry.is_dir() and entry.name in config["migration_modules"]:
+            _load_module(machine, config, entry.name)

src/rookify/modules/machine.py

@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+
+from transitions import MachineError
+from transitions import Machine as _Machine
+from transitions.extensions.states import add_state_features, Tags, Timeout
+from typing import Any, Dict
+
+
+@add_state_features(Tags, Timeout)
+class Machine(_Machine):  # type: ignore
+    def __init__(self) -> None:
+        _Machine.__init__(self, states=["uninitialized"], initial="uninitialized")
+
+    def add_migrating_state(self, name: Any, **kwargs: Dict[str, Any]) -> None:
+        if not isinstance(name, str):
+            raise MachineError("Migration state name must be string")
+        self.add_state(name, **kwargs)
+
+    def execute(self) -> None:
+        self.add_state("migrated")
+        self.add_ordered_transitions(loop=False)
+
+        try:
+            while True:
+                self.next_state()
+        except MachineError:
+            if self.state != "migrated":
+                raise