This commit is contained in:
+181
-328
@@ -1,355 +1,208 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import yaml
|
||||
import json
|
||||
import importlib.util
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Dict, List, Any, Optional, Tuple
|
||||
from dataclasses import MISSING, asdict, dataclass
|
||||
from dataclasses import dataclass
|
||||
from abc import ABC, abstractmethod
|
||||
from pydantic import BaseModel
|
||||
|
||||
from playbook.models import ActModel, PlaybookModel, timed_run
|
||||
from playbook.logging_models import Status, StepLogModel
|
||||
from playbook.action_registry import ActionFn
|
||||
|
||||
from playbook.models import Status, StepLog, timed_run
|
||||
from playbook.action_registry import ActionFn, ActionRegistry
|
||||
from playbook.premade_steps_registry import registry
|
||||
|
||||
|
||||
class StepIF(ABC):
|
||||
@abstractmethod
|
||||
def run(self, ctx: Dict[str, Any]) -> StepLog:
|
||||
def run(self, ctx: dict[str, object]) -> StepLogModel:
|
||||
"""Perform the actions."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_action_names(self) -> List[str]:
|
||||
def get_action_names(self) -> list[str]:
|
||||
""" Return a mapping of action roles (pre, play, post) to function names. """
|
||||
pass
|
||||
|
||||
|
||||
class CustomStep(StepIF):
|
||||
""" Setup custom steps. """
|
||||
|
||||
def __init__(self, name: str, actions: List[ActionFn], context: Dict[str, Any]):
|
||||
self.name: str = name
|
||||
self.__actions = actions
|
||||
self.__context = context
|
||||
|
||||
def get_action_names(self) -> List[str]:
|
||||
return [action.__name__ for action in self.__actions]
|
||||
|
||||
# NOTE: refactor this method to be more readable
|
||||
def run(self, ctx) -> StepLog:
|
||||
""" Run the step. """
|
||||
substeps: List[StepLog] = []
|
||||
status: Status = Status.GOOD
|
||||
|
||||
# NOTE:
|
||||
# Use this to move context from the previous step to the next step
|
||||
# This is wrong, we are moving away from having a list of steps here to a single
|
||||
# step for each operation
|
||||
prev_step_ctx: Dict[str, Any] = {}
|
||||
|
||||
for op in self.__actions:
|
||||
try:
|
||||
log: StepLog = timed_run(op, ctx | prev_step_ctx | self.__context, self.name)
|
||||
except Exception as e:
|
||||
log: StepLog = StepLog.fail(self.name, [{"status": "failed", "output": str(e)}])
|
||||
|
||||
status = Status.BAD if log.failed else Status.GOOD
|
||||
prev_step_ctx = log.pipe_ctx
|
||||
|
||||
substeps.append(log)
|
||||
|
||||
if status == Status.BAD:
|
||||
break
|
||||
|
||||
errors = []
|
||||
msg: str = ""
|
||||
if status == Status.BAD:
|
||||
# Note, use this to set the standar error output/formatting
|
||||
# errors = [{"status": "failed", "output": ""}]
|
||||
errors = []
|
||||
else:
|
||||
msg = "success"
|
||||
|
||||
return StepLog(
|
||||
step_name=self.name, status=status,
|
||||
msg=msg, error=errors, substeps=substeps, pipe_ctx=prev_step_ctx
|
||||
)
|
||||
# class PlaybookError(Exception):
|
||||
# """ Base exception for Playbook parsing and execution failures. """
|
||||
# pass
|
||||
|
||||
|
||||
class PlaybookError(Exception):
|
||||
""" Base exception for Playbook parsing and execution failures. """
|
||||
pass
|
||||
# @dataclass
|
||||
# class StepEntry(object):
|
||||
# name: str
|
||||
# step: StepIF
|
||||
|
||||
|
||||
@dataclass
|
||||
class StepEntry(object):
|
||||
name: str
|
||||
step: StepIF
|
||||
# class Play(object):
|
||||
# def __init__(self, name: str, registries: dict[str, ActionRegistry]):
|
||||
# self.name: str = name
|
||||
# self.log_dir: str = ""
|
||||
# self._context: dict[str, object] = {}
|
||||
# self._acts: dict[str, list[CustomStep]] = {}
|
||||
# self._steps: list[StepEntry] = []
|
||||
# self._registries: dict[str, ActionRegistry] = {
|
||||
# registry.name: registry} | registries
|
||||
#
|
||||
#
|
||||
# def add_step(self, name: str, stepRunner: StepIF):
|
||||
# self._steps.append(StepEntry(name=name, step=stepRunner))
|
||||
#
|
||||
# def view_playbook(self) -> str:
|
||||
# acts_info = {}
|
||||
# for act_name, step_list in self._acts.items():
|
||||
# steps_info = []
|
||||
# for s in step_list:
|
||||
# steps_info.append({
|
||||
# "name": s.name,
|
||||
# "actions": s.get_action_names()
|
||||
# })
|
||||
#
|
||||
# new_act = {
|
||||
# "name": act_name,
|
||||
# "steps": steps_info
|
||||
# }
|
||||
#
|
||||
# acts_info[act_name] = new_act
|
||||
#
|
||||
# registry_data = []
|
||||
# for reg in self._registries.values():
|
||||
# registry_data.append(reg.manifest())
|
||||
#
|
||||
# data = {
|
||||
# "name": self.name,
|
||||
# "registries": registry_data,
|
||||
# "number_of_acts": len(acts_info),
|
||||
# "acts": acts_info,
|
||||
# }
|
||||
#
|
||||
# return json.dumps(data)
|
||||
#
|
||||
# def play(self) -> StepLogModel:
|
||||
# log: StepLogModel = timed_run(self._play_act)
|
||||
#
|
||||
# if self.log_dir != "":
|
||||
# log.log_file_path = os.path.join(
|
||||
# self.log_dir, self.log_file_name(log.start_date_timestamp)
|
||||
# )
|
||||
# self._write_log(log)
|
||||
#
|
||||
# return log
|
||||
#
|
||||
# def log_file_name(self, start_date_ns: int) -> str:
|
||||
# """ Format to ISO 8601. """
|
||||
# seconds = start_date_ns // 1_000_000_000
|
||||
# ms = (start_date_ns % 1_000_000_000) // 1_000_000
|
||||
#
|
||||
# dt = datetime.fromtimestamp(seconds, tz=timezone.utc).astimezone()
|
||||
#
|
||||
# return dt.strftime(f"{self.name}_%Y-%m-%d_%H-%M-%S.{ms:03d}.log")
|
||||
#
|
||||
# def _write_log(self, log: StepLogModel):
|
||||
# try:
|
||||
# with open(log.log_file_path, "w") as file:
|
||||
# json.dump(asdict(log), file)
|
||||
# except Exception as e:
|
||||
# raise PlaybookError(f"failed to create log file {e}") from e
|
||||
# return ""
|
||||
#
|
||||
# def _play_act(self) -> StepLogModel:
|
||||
# playLog: StepLogModel = StepLogModel(
|
||||
# step_name=self.name,
|
||||
# status=Status.GOOD,
|
||||
# duration_sec=0,
|
||||
# msg="",
|
||||
# error=[],
|
||||
# substeps=[],
|
||||
# )
|
||||
#
|
||||
# for a in self._acts:
|
||||
# log: StepLogModel = self._play_steps(self._acts[a])
|
||||
#
|
||||
# if log.failed:
|
||||
# playLog.error.append({
|
||||
# "status": "failed",
|
||||
# "output": f"failed in step: {log.step_name}"
|
||||
# })
|
||||
# playLog.status = Status.BAD
|
||||
# break
|
||||
#
|
||||
# if playLog.status == Status.GOOD:
|
||||
# playLog.msg = "success"
|
||||
#
|
||||
# return playLog
|
||||
#
|
||||
# def _play_steps(self, steps: list[CustomStep]) -> StepLogModel:
|
||||
# playLog: StepLogModel = StepLogModel(
|
||||
# step_name=self.name,
|
||||
# status=Status.GOOD,
|
||||
# duration_sec=0,
|
||||
# msg="",
|
||||
# error=[],
|
||||
# substeps=[],
|
||||
# )
|
||||
#
|
||||
# prev_ctx: dict[str, object] = {}
|
||||
# for step in steps:
|
||||
# log: StepLogModel = timed_run(step.run, self._context | prev_ctx)
|
||||
# prev_ctx = log.pipe_ctx
|
||||
#
|
||||
# playLog.substeps.append(log)
|
||||
#
|
||||
# if log.failed:
|
||||
# playLog.error.append({
|
||||
# "status": "failed",
|
||||
# "output": f"failed in step: {step.name}"
|
||||
# })
|
||||
# playLog.status = Status.BAD
|
||||
# break
|
||||
#
|
||||
# if playLog.status == Status.GOOD:
|
||||
# playLog.msg = "success"
|
||||
#
|
||||
# return playLog
|
||||
#
|
||||
# def _get_function_from_registry(self, fn_name) -> ActionFn:
|
||||
# for reg in self._registries.values():
|
||||
# try:
|
||||
# return reg.get(fn_name, "*")
|
||||
# except ValueError:
|
||||
# continue
|
||||
#
|
||||
# raise ValueError(f"no function with name '{fn_name}' was found")
|
||||
#
|
||||
# def _build_custom_step(self, name, actions, ctx) -> CustomStep:
|
||||
# functionlist: list[ActionFn] = []
|
||||
# for f in actions:
|
||||
# try:
|
||||
# fn = self._get_function_from_registry(actions[f])
|
||||
# functionlist.append(fn)
|
||||
# except KeyError as e:
|
||||
# raise PlaybookError(
|
||||
# f"Step '{name}' is missing required action field: {e}") from e
|
||||
#
|
||||
# return CustomStep(name, functionlist, ctx)
|
||||
|
||||
|
||||
class Play(object):
|
||||
def __init__(self, name : str, registries: Dict[str, ActionRegistry]):
|
||||
self.name: str = name
|
||||
self.log_dir: str = ""
|
||||
self._context: Dict[str, Any] = {}
|
||||
self._acts: Dict[str, List[CustomStep]] = {}
|
||||
self._steps: List[StepEntry] = []
|
||||
self._registries: Dict[str, ActionRegistry] = {registry.name: registry} | registries
|
||||
class Playbook():
|
||||
model: PlaybookModel
|
||||
|
||||
# NOTE: this should move to the ActionRegistry thing probably
|
||||
@staticmethod
|
||||
def _load_registry_file(file_path: str) -> List[ActionRegistry]:
|
||||
abs_path = os.path.abspath(file_path)
|
||||
module_name = os.path.splitext(os.path.basename(abs_path))[0]
|
||||
def __init__(self, model: PlaybookModel):
|
||||
self.model = model
|
||||
|
||||
@property
|
||||
def acts(self) -> list[ActModel]:
|
||||
return self.model.acts
|
||||
|
||||
spec = importlib.util.spec_from_file_location(module_name, abs_path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError(f"could not create spec for registry file: {file_path}")
|
||||
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
discovered_registries = [
|
||||
obj for obj in vars(module).values()
|
||||
if isinstance(obj, ActionRegistry)
|
||||
]
|
||||
|
||||
if not discovered_registries:
|
||||
raise PlaybookError(f"No ActionRegistry instances were found in '{file_path}'")
|
||||
|
||||
return discovered_registries
|
||||
|
||||
def add_step(self, name: str, stepRunner : StepIF):
|
||||
self._steps.append(StepEntry(name=name, step=stepRunner))
|
||||
@classmethod
|
||||
def from_yaml(cls, fp: str) -> Playbook:
|
||||
try:
|
||||
new_playboook = PlaybookModel.from_yaml_file(fp)
|
||||
return Playbook(new_playboook)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def view_playbook(self) -> str:
|
||||
acts_info = {}
|
||||
for act_name, step_list in self._acts.items():
|
||||
steps_info = []
|
||||
for s in step_list:
|
||||
steps_info.append({
|
||||
"name": s.name,
|
||||
"actions": s.get_action_names()
|
||||
})
|
||||
return self.model.model_dump_json()
|
||||
|
||||
new_act = {
|
||||
"name": act_name,
|
||||
"steps": steps_info
|
||||
}
|
||||
|
||||
acts_info[act_name] = new_act
|
||||
|
||||
|
||||
|
||||
registry_data = []
|
||||
for reg in self._registries.values():
|
||||
registry_data.append(reg.manifest())
|
||||
|
||||
data = {
|
||||
"name": self.name,
|
||||
"registries": registry_data,
|
||||
"number_of_acts": len(acts_info),
|
||||
"acts": acts_info,
|
||||
}
|
||||
|
||||
return json.dumps(data)
|
||||
|
||||
def play(self) -> StepLog:
|
||||
log: StepLog = timed_run(self._play_act)
|
||||
|
||||
if self.log_dir != "":
|
||||
log.log_file_path = os.path.join(
|
||||
self.log_dir, self.log_file_name(log.start_date_timestamp)
|
||||
)
|
||||
self._write_log(log)
|
||||
|
||||
return log
|
||||
|
||||
def log_file_name(self, start_date_ns: int) -> str:
|
||||
""" Format to ISO 8601. """
|
||||
seconds = start_date_ns // 1_000_000_000
|
||||
ms = (start_date_ns % 1_000_000_000) // 1_000_000
|
||||
|
||||
dt = datetime.fromtimestamp(seconds, tz=timezone.utc).astimezone()
|
||||
|
||||
return dt.strftime(f"{self.name}_%Y-%m-%d_%H-%M-%S.{ms:03d}.log")
|
||||
|
||||
def _write_log(self, log: StepLog):
|
||||
try:
|
||||
with open(log.log_file_path, "w") as file:
|
||||
json.dump(asdict(log), file)
|
||||
except Exception as e:
|
||||
raise PlaybookError(f"failed to create log file {e}") from e
|
||||
return ""
|
||||
|
||||
def _play_act(self) -> StepLog:
|
||||
playLog: StepLog = StepLog(
|
||||
step_name=self.name,
|
||||
status=Status.GOOD,
|
||||
duration_sec=0,
|
||||
msg="",
|
||||
error=[],
|
||||
substeps=[],
|
||||
)
|
||||
|
||||
for a in self._acts:
|
||||
log: StepLog = self._play_steps(self._acts[a])
|
||||
|
||||
if log.failed:
|
||||
playLog.error.append({
|
||||
"status": "failed",
|
||||
"output": f"failed in step: {log.step_name}"
|
||||
})
|
||||
playLog.status = Status.BAD
|
||||
break
|
||||
|
||||
if playLog.status == Status.GOOD:
|
||||
playLog.msg = "success"
|
||||
|
||||
return playLog
|
||||
|
||||
def _play_steps(self, steps: List[CustomStep]) -> StepLog:
|
||||
playLog: StepLog = StepLog(
|
||||
step_name=self.name,
|
||||
status=Status.GOOD,
|
||||
duration_sec=0,
|
||||
msg="",
|
||||
error=[],
|
||||
substeps=[],
|
||||
)
|
||||
|
||||
prev_ctx: Dict[str, Any] = {}
|
||||
for step in steps:
|
||||
log: StepLog = timed_run(step.run, self._context | prev_ctx)
|
||||
prev_ctx = log.pipe_ctx
|
||||
|
||||
playLog.substeps.append(log)
|
||||
|
||||
if log.failed:
|
||||
playLog.error.append({
|
||||
"status": "failed",
|
||||
"output": f"failed in step: {step.name}"
|
||||
})
|
||||
playLog.status = Status.BAD
|
||||
break
|
||||
|
||||
if playLog.status == Status.GOOD:
|
||||
playLog.msg = "success"
|
||||
|
||||
return playLog
|
||||
|
||||
def _get_function_from_registry(self, fn_name) -> ActionFn:
|
||||
for reg in self._registries.values():
|
||||
try:
|
||||
return reg.get(fn_name, "*")
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
raise ValueError(f"no function with name '{fn_name}' was found")
|
||||
|
||||
def _build_custom_step(self, name, actions, ctx) -> CustomStep:
|
||||
functionList: List[ActionFn] = []
|
||||
for f in actions:
|
||||
try:
|
||||
fn = self._get_function_from_registry(actions[f])
|
||||
functionList.append(fn)
|
||||
except KeyError as e:
|
||||
raise PlaybookError(f"Step '{name}' is missing required action field: {e}") from e
|
||||
|
||||
return CustomStep(name, functionList, ctx)
|
||||
|
||||
# NOTE: Move to factory
|
||||
@classmethod
|
||||
def from_yaml(cls, fp: str) -> Play:
|
||||
""" Loads the playbook from a given YAML file. """
|
||||
try:
|
||||
with open(fp, "r") as f:
|
||||
file_contents = f.read()
|
||||
except FileNotFoundError as e:
|
||||
raise PlaybookError(f"Playbook file not found: {fp}") from e
|
||||
except Exception as e:
|
||||
raise PlaybookError(f"Failed to read file {fp}: {e}") from e
|
||||
|
||||
yaml_dir = os.path.dirname(os.path.abspath(fp))
|
||||
return cls.__from_yaml_str(file_contents, base_dir=yaml_dir)
|
||||
|
||||
# NOTE: Move to factory
|
||||
@classmethod
|
||||
def __from_yaml_str_load_registries(cls, dir, data) -> Dict[str, ActionRegistry]:
|
||||
custom_registries: Dict[str, ActionRegistry] = {}
|
||||
registry_paths = data.get("registries", {})
|
||||
|
||||
custom_registries: Dict[str, ActionRegistry] = {}
|
||||
for reg_path in registry_paths:
|
||||
full_path = os.path.join(dir, reg_path) if not os.path.isabs(reg_path) else reg_path
|
||||
try:
|
||||
regs = cls._load_registry_file(full_path)
|
||||
for new_reg in regs:
|
||||
custom_registries[new_reg.name] = new_reg
|
||||
except Exception as e:
|
||||
raise PlaybookError(f"Error loading registry '{reg_path}': {e}") from e
|
||||
|
||||
return custom_registries
|
||||
|
||||
# NOTE: yes, I know, I have to refactor this
|
||||
def __from_yaml_str_load_acts(self, data):
|
||||
for idx, act in enumerate(data.get("acts", [])):
|
||||
act_name = act.get("name", f"step_{idx}")
|
||||
steps = act.get("steps", [])
|
||||
|
||||
self._acts[act_name] = []
|
||||
try:
|
||||
self.__from_yaml_str_load_steps(act_name, steps)
|
||||
except Exception as e:
|
||||
raise PlaybookError(f"error loading steps for act {act_name}") from e
|
||||
|
||||
|
||||
# NOTE: Move to factory
|
||||
def __from_yaml_str_load_steps(self, key, data):
|
||||
for step_cfg in data:
|
||||
step_name = step_cfg.get("name", f"")
|
||||
actions = step_cfg.get("actions", {})
|
||||
context = step_cfg.get("context", {})
|
||||
|
||||
if not isinstance(actions, dict) or len(actions) <= 0:
|
||||
raise PlaybookError(
|
||||
f"Step '{step_name}' (index ) must define at least one action "
|
||||
)
|
||||
|
||||
try:
|
||||
customStep: CustomStep = self._build_custom_step(step_name, actions, context)
|
||||
self._acts[key].append(customStep)
|
||||
except Exception as e:
|
||||
raise PlaybookError(f"Error configuring step '{step_name}': {e}") from e
|
||||
|
||||
# NOTE: Move to factory
|
||||
@classmethod
|
||||
def __from_yaml_str(cls, yaml_str: str, base_dir: str = ".") -> Play:
|
||||
data = None
|
||||
try:
|
||||
data = yaml.safe_load(yaml_str)
|
||||
except Exception as e:
|
||||
raise PlaybookError(f"YAML Syntax error: {e}") from e
|
||||
|
||||
if not isinstance(data, dict) or "playbook_name" not in data:
|
||||
raise PlaybookError("YAML must contain a top-level 'playbook_name' field.")
|
||||
|
||||
try:
|
||||
custom_registries = cls.__from_yaml_str_load_registries(base_dir, data)
|
||||
except Exception as e:
|
||||
raise PlaybookError(f"registries load error: {e}") from e
|
||||
|
||||
playbook = cls(data["playbook_name"], registries=custom_registries)
|
||||
try:
|
||||
playbook.__from_yaml_str_load_acts(data)
|
||||
except Exception as e:
|
||||
raise PlaybookError(f"steps load error: {e}") from e
|
||||
|
||||
playbook._context = data.get("global_context", {})
|
||||
playbook.log_dir = data.get("log_dir", "")
|
||||
|
||||
return playbook
|
||||
def run_play(self):
|
||||
return self.model.run()
|
||||
|
||||
Reference in New Issue
Block a user