From 07be5c79e197520ec3cc0f50d694e0b10f93f96d Mon Sep 17 00:00:00 2001 From: Tobias Wolf Date: Thu, 9 May 2024 17:15:01 +0200 Subject: [PATCH] Add "transitions" based execution machine This adapts the logic how modules are loaded and sorted for execution. Signed-off-by: Tobias Wolf --- requirements.txt | 1 + src/rookify/__main__.py | 62 ++------------ src/rookify/modules/__init__.py | 140 +++++++++----------------------- src/rookify/modules/machine.py | 28 +++++++ src/rookify/modules/module.py | 84 +++++++++++++------ 5 files changed, 132 insertions(+), 183 deletions(-) create mode 100644 src/rookify/modules/machine.py diff --git a/requirements.txt b/requirements.txt index 9246059..0464477 100644 --- a/requirements.txt +++ b/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 diff --git a/src/rookify/__main__.py b/src/rookify/__main__.py index 89c8853..6474da3 100644 --- a/src/rookify/__main__.py +++ b/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() diff --git a/src/rookify/modules/__init__.py b/src/rookify/modules/__init__.py index 3167e96..40993f7 100644 --- a/src/rookify/modules/__init__.py +++ b/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,9 +21,39 @@ 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. @@ -34,97 +61,6 @@ def load_modules( :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) diff --git a/src/rookify/modules/machine.py b/src/rookify/modules/machine.py new file mode 100644 index 0000000..5b7340c --- /dev/null +++ b/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 diff --git a/src/rookify/modules/module.py b/src/rookify/modules/module.py index 8642717..5b4c514 100644 --- a/src/rookify/modules/module.py +++ b/src/rookify/modules/module.py @@ -9,9 +9,11 @@ import fabric import jinja2 import structlog -from rookify.logger import get_logger from typing import Any, Dict, List, Optional +from ..logger import get_logger +from .machine import Machine + class ModuleException(Exception): pass @@ -152,7 +154,7 @@ def yaml(self) -> Any: self.__result_yaml = yaml.safe_load(self.raw) return self.__result_yaml - def __init__(self, config: Dict[str, Any], data: Dict[str, Any], module_path: str): + def __init__(self, machine: Machine, config: Dict[str, Any]): """ Construct a new 'ModuleHandler' object. @@ -162,31 +164,13 @@ def __init__(self, config: Dict[str, Any], data: Dict[str, Any], module_path: st :return: returns nothing """ self._config = config - self._data = data - self.__module_path = module_path + self._machine = machine + self.__ceph: Optional[ModuleHandler.__Ceph] = None self.__k8s: Optional[ModuleHandler.__K8s] = None self.__ssh: Optional[ModuleHandler.__SSH] = None self.__logger = get_logger() - self.__logger.debug("Executing {0}".format(self.__class__.__name__)) - - @abc.abstractmethod - def preflight(self) -> None: - """ - Run the modules preflight check - """ - pass - - @abc.abstractmethod - def run(self) -> Dict[str, Any]: - """ - Run the modules tasks - - :return: returns result - """ - pass - @property def ceph(self) -> __Ceph: if self.__ceph is None: @@ -194,8 +178,8 @@ def ceph(self) -> __Ceph: return self.__ceph @property - def logger(self) -> structlog.getLogger: - return self.__logger + def machine(self) -> Machine: + return self._machine @property def k8s(self) -> __K8s: @@ -203,14 +187,64 @@ def k8s(self) -> __K8s: self.__k8s = ModuleHandler.__K8s(self._config["kubernetes"]) return self.__k8s + @property + def logger(self) -> structlog.getLogger: + return self.__logger + @property def ssh(self) -> __SSH: if self.__ssh is None: self.__ssh = ModuleHandler.__SSH(self._config["ssh"]) return self.__ssh + @abc.abstractmethod + def preflight(self) -> None: + """ + Run the modules preflight check + """ + pass + + @abc.abstractmethod + def run(self) -> Dict[str, Any]: + """ + Run the modules tasks + + :return: returns result + """ + pass + def load_template(self, filename: str, **variables: Any) -> __Template: - template_path = os.path.join(self.__module_path, "templates", filename) + template_path = os.path.join(os.path.dirname(__file__), "templates", filename) template = ModuleHandler.__Template(template_path) template.render(**variables) return template + + @classmethod + def register_state( + cls, machine: Machine, config: Dict[str, Any], **kwargs: Any + ) -> None: + """ + Register state for transitions + """ + + state_name = cls.STATE_NAME if hasattr(cls, "STATE_NAME") else cls.__name__ + + handler = cls(machine, config) + + if hasattr(cls, "preflight") and not getattr( + cls.preflight, "__isabstractmethod__", False + ): + kwargs["on_enter"] = handler.preflight + + if hasattr(cls, "run") and not getattr(cls.run, "__isabstractmethod__", False): + kwargs["on_exit"] = handler.run + + if len(kwargs) > 0: + get_logger().debug("Registering state {0}".format(state_name)) + machine.add_migrating_state(state_name, **kwargs) + else: + get_logger().warn( + "Not registering state {0} because ModuleHandler has no expected callables".format( + state_name + ) + )