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()
|
||||
|
||||
@@ -1,49 +1,58 @@
|
||||
from dataclasses import dataclass
|
||||
from playbook.models import StepLog, ActionFn
|
||||
from typing import Callable, Dict, Any
|
||||
import os
|
||||
import sys
|
||||
import importlib.util
|
||||
from playbook.logging_models import StepLogModel
|
||||
from typing import Callable
|
||||
from pydantic import BaseModel, PrivateAttr
|
||||
|
||||
|
||||
@dataclass
|
||||
class Action(object):
|
||||
ActionFn = Callable[[dict[str, object], str], StepLogModel]
|
||||
|
||||
|
||||
class Action(BaseModel):
|
||||
name: str
|
||||
fn: ActionFn
|
||||
ver: str
|
||||
|
||||
|
||||
class ActionRegistry:
|
||||
def __init__(self, name: str):
|
||||
self.name: str = name
|
||||
self.__actions = {}
|
||||
class ActionRegistry(BaseModel):
|
||||
name: str
|
||||
_actions: dict[str, Action] = PrivateAttr(default_factory=dict)
|
||||
|
||||
def register(self, name: str, version: str):
|
||||
"""Decorator to register python functions with a name and version."""
|
||||
def decorator(function: ActionFn):
|
||||
if name not in self.__actions:
|
||||
self.__actions[name] = {}
|
||||
|
||||
self.__actions[name][version] = function
|
||||
self._actions[name] = Action(name=name, fn=function, ver=version)
|
||||
return function
|
||||
|
||||
return decorator
|
||||
|
||||
def get(self, name: str, ver: str) -> ActionFn:
|
||||
if name not in self.__actions:
|
||||
raise ValueError(f"Action '{name}' is not registered in registry '{self.name}'")
|
||||
|
||||
versions_dict = self.__actions[name]
|
||||
if not versions_dict:
|
||||
raise ValueError(f"No registered versions found for function '{name}'")
|
||||
|
||||
if ver == "*":
|
||||
return next(iter(versions_dict.values()))
|
||||
def get(self, name: str, ver: str = '*') -> ActionFn:
|
||||
if name not in self._actions:
|
||||
raise ValueError(
|
||||
f"Action '{name}' is not registered in registry '{self.name}'")
|
||||
|
||||
if ver not in self.__actions[name]:
|
||||
raise ValueError(f"Version '{ver}' is not registered for function '{name}'")
|
||||
return self._actions[name].fn
|
||||
|
||||
return versions_dict[ver]
|
||||
@staticmethod
|
||||
def load_registries_from_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 manifest(self):
|
||||
functions = {}
|
||||
for fn_name, versions_dict in self.__actions.items():
|
||||
functions[fn_name] = list(versions_dict.keys())
|
||||
return {"registry_name": self.name, "functions": functions}
|
||||
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 ValueError(f"No ActionRegistry instances were found in '{file_path}'")
|
||||
|
||||
return discovered_registries
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
from enum import StrEnum
|
||||
from dataclasses import field
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class Status(StrEnum):
|
||||
GOOD = "GOOD"
|
||||
BAD = "BAD"
|
||||
|
||||
|
||||
class LogTiming(BaseModel):
|
||||
start_date_timestamp: int = 0
|
||||
start_date: str = ""
|
||||
end_date: str = ""
|
||||
duration_sec: float = 0.0
|
||||
|
||||
|
||||
class StepLogModel(BaseModel):
|
||||
name: str
|
||||
timing: LogTiming = field(default=LogTiming())
|
||||
status: Status
|
||||
error: str = ""
|
||||
msg: str = ""
|
||||
pipe_ctx: dict[str, object] = field(default_factory=dict)
|
||||
log_file_path: str = "" # If set we have logged this particular step to a file
|
||||
|
||||
@property
|
||||
def failed(self) -> bool:
|
||||
""" Checks if the step failed """
|
||||
return len(self.error) > 0 or self.status == Status.BAD
|
||||
|
||||
@classmethod
|
||||
def ok(cls, name: str, msg: str = "success", pipe_ctx: dict[str, object] = {}) -> StepLogModel:
|
||||
return cls(name=name, status=Status.GOOD, msg=msg, pipe_ctx=pipe_ctx)
|
||||
|
||||
@classmethod
|
||||
def fail(cls, name: str, err: str) -> StepLogModel:
|
||||
return cls(name=name, status=Status.BAD, error=err)
|
||||
+112
-46
@@ -1,59 +1,128 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import yaml
|
||||
import time
|
||||
import functools
|
||||
from enum import StrEnum
|
||||
from typing import List, Dict, Any, Callable
|
||||
from typing import Callable
|
||||
from datetime import datetime, timezone
|
||||
from dataclasses import dataclass, field
|
||||
from pydantic import BaseModel, field_validator
|
||||
|
||||
from playbook.action_registry import ActionRegistry, ActionFn
|
||||
from playbook.logging_models import StepLogModel
|
||||
|
||||
|
||||
class Status(StrEnum):
|
||||
GOOD = "GOOD"
|
||||
BAD = "BAD"
|
||||
class StepModel(BaseModel):
|
||||
name: str
|
||||
_actions: list[ActionFn]
|
||||
_context: dict[str, object]
|
||||
|
||||
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) -> StepLogModel:
|
||||
""" Run the step. """
|
||||
return StepLogModel.ok(self.name, msg="success")
|
||||
# substeps: list[StepLogModel] = []
|
||||
# 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, object] = {}
|
||||
#
|
||||
# for op in self.__actions:
|
||||
# try:
|
||||
# log: StepLogModel = timed_run(
|
||||
# op, ctx | prev_step_ctx | self.__context, self.name)
|
||||
# except Exception as e:
|
||||
# log: StepLogModel = StepLogModel.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 StepLogModel(
|
||||
# step_name=self.name, status=status,
|
||||
# msg=msg, error=errors, substeps=substeps, pipe_ctx=prev_step_ctx
|
||||
# )
|
||||
|
||||
|
||||
@dataclass
|
||||
class StepLog(object):
|
||||
step_name: str
|
||||
status: Status = Status.BAD # By default set it to bad so the it fails by default
|
||||
# The user needs to be explicit
|
||||
cmd: str = ""
|
||||
start_date_timestamp: int = 0
|
||||
start_date: str = ""
|
||||
end_date: str = ""
|
||||
duration_sec: float = 0.0
|
||||
human_readable_duration: str = ""
|
||||
msg: str = ""
|
||||
error: List[Dict[Any, Any]] = field(default_factory=list)
|
||||
substeps: List[StepLog] = field(default_factory=list)
|
||||
log_file_path: str = ""
|
||||
pipe_ctx: Dict[str, Any] = field(default_factory=dict)
|
||||
class ActModel(BaseModel):
|
||||
name: str
|
||||
steps: list[StepModel]
|
||||
|
||||
@property
|
||||
def failed(self) -> bool:
|
||||
""" Checks if the step failed """
|
||||
return len(self.error) > 0 or self.status == Status.BAD
|
||||
# NOTE: this shouldnt return a steplogmodel but an actlogmodel or something like that
|
||||
def run(self, ctx) -> StepLogModel:
|
||||
return StepLogModel.ok(self.name, msg="success")
|
||||
|
||||
|
||||
class PlaybookModel(BaseModel):
|
||||
playbook_name: str
|
||||
log_dir: str
|
||||
registries: list[ActionRegistry]
|
||||
global_context: dict[str, str]
|
||||
acts: list[ActModel]
|
||||
|
||||
@field_validator("registries", mode="before")
|
||||
@classmethod
|
||||
def parse_registries(cls, v: object) -> list[ActionRegistry]:
|
||||
result = []
|
||||
if isinstance(v, list):
|
||||
for item in v:
|
||||
if isinstance(item, str):
|
||||
try:
|
||||
registries: list[ActionRegistry] = \
|
||||
ActionRegistry.load_registries_from_file(item)
|
||||
result.extend(registries)
|
||||
except Exception as e:
|
||||
raise e
|
||||
else:
|
||||
raise ValueError(f"Invalid registry type: {type(item)}")
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def ok(cls, name: str, msg: str = "success", pipe_ctx: Dict[str, Any] = {}) -> StepLog:
|
||||
return cls(step_name=name, status=Status.GOOD, msg=msg, pipe_ctx=pipe_ctx)
|
||||
def from_yaml_file(cls, fp: str):
|
||||
with open(fp, "rb") as f:
|
||||
data = yaml.safe_load(f)
|
||||
return cls(**data)
|
||||
|
||||
@classmethod
|
||||
def fail(cls, name: str, errors: List[Any]) -> StepLog:
|
||||
return cls(step_name=name, status=Status.BAD, error=errors)
|
||||
def _run(self) -> StepLogModel:
|
||||
for act in self.acts:
|
||||
act.run(self.global_context)
|
||||
|
||||
return StepLogModel.ok(name="somasjd", msg="ajshdsajhd")
|
||||
|
||||
def run(self) -> StepLogModel:
|
||||
return timed_run(self._run)
|
||||
|
||||
|
||||
class ContextChecker:
|
||||
"""Safely extract values with error messages."""
|
||||
def __init__(self, ctx: Dict[str, Any]):
|
||||
|
||||
def __init__(self, ctx: dict[str, object]):
|
||||
self._ctx = ctx
|
||||
|
||||
@classmethod
|
||||
def requires(cls, *args):
|
||||
def decorator(func: ActionFn):
|
||||
@functools.wraps(func)
|
||||
def wrapper(ctx: Dict[str, Any], name: str) -> StepLog:
|
||||
def wrapper(ctx: dict[str, object], name: str) -> StepLogModel:
|
||||
if all(key in ctx for key in args):
|
||||
return func(ctx, name)
|
||||
else:
|
||||
@@ -73,24 +142,21 @@ def human_readable_date(time: int) -> str:
|
||||
return dt.strftime(f"%Y-%m-%dT%H:%M:%S.{nanos:09d}%:z")
|
||||
|
||||
|
||||
def timed_run(op: Callable[..., StepLog], *args: Any, **kwargs: Any) -> StepLog:
|
||||
def timed_run(op: Callable[..., StepLogModel], *args: object, **kwargs: object) -> StepLogModel:
|
||||
start_date: int = time.time_ns()
|
||||
start_time: int = time.perf_counter_ns()
|
||||
|
||||
|
||||
try:
|
||||
log: StepLog = op(*args, **kwargs)
|
||||
log: StepLogModel = op(*args, **kwargs)
|
||||
finally:
|
||||
delta = time.perf_counter_ns() - start_time
|
||||
|
||||
|
||||
# NOTE: we could simplyfy this by having a marshalling method
|
||||
if isinstance(log, StepLog):
|
||||
if isinstance(log, StepLogModel):
|
||||
end_date: int = time.time_ns()
|
||||
log.duration_sec = delta / 1_000_000_000.0
|
||||
log.start_date = human_readable_date(start_date)
|
||||
log.start_date_timestamp = start_date
|
||||
log.end_date = human_readable_date(end_date)
|
||||
|
||||
log.timing.duration_sec = delta / 1_000_000_000.0
|
||||
log.timing.start_date = human_readable_date(start_date)
|
||||
log.timing.start_date_timestamp = start_date
|
||||
log.timing.end_date = human_readable_date(end_date)
|
||||
|
||||
return log
|
||||
|
||||
|
||||
ActionFn = Callable[[Dict[str, Any], str], StepLog]
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
# In this file we will only put premade steps
|
||||
|
||||
from playbook.models import StepLog, Status
|
||||
from playbook.action_registry import ActionRegistry
|
||||
|
||||
|
||||
# Module wide registry with built in functions
|
||||
registry: ActionRegistry = ActionRegistry("testrunner_registry")
|
||||
registry: ActionRegistry = ActionRegistry(name="testrunner_registry")
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
from playbook import ActionRegistrar, ActionRegistrarYAMLReader
|
||||
from playbook.action_registry import ActionRegistry
|
||||
|
||||
|
||||
CDIR = os.path.dirname(__file__)
|
||||
TEST_DATA = os.path.join(CDIR, "test_data")
|
||||
|
||||
|
||||
class TestActionRegistrar(unittest.TestCase):
|
||||
@patch.object(ActionRegistrarYAMLReader, '_load_registries_from_file')
|
||||
def test_reading_action_registry(self, mock_load):
|
||||
mock_reg1 = ActionRegistry("docker_reg")
|
||||
mock_reg2 = ActionRegistry("restic_reg")
|
||||
mock_load.side_effect = [[mock_reg1], [mock_reg2]]
|
||||
|
||||
yaml_data = {
|
||||
"registries": ["docker.py", "restic.py"]
|
||||
}
|
||||
|
||||
expected_size = 2
|
||||
expected_manifest = {
|
||||
"registrar_manifest": [
|
||||
{"registry_name": "docker_reg", "functions": {}},
|
||||
{"registry_name": "restic_reg", "functions": {}}
|
||||
]
|
||||
}
|
||||
|
||||
reg: ActionRegistrar = ActionRegistrarYAMLReader.read(yaml_data)
|
||||
|
||||
self.assertEqual(expected_manifest, reg.manifest())
|
||||
self.assertEqual(expected_size, reg.size())
|
||||
|
||||
pass
|
||||
|
||||
# def test_creating_action_registry(self):
|
||||
# yaml_data = {
|
||||
# "registries": ["docker.py", "restic.py"]
|
||||
# }
|
||||
#
|
||||
# expected_size = 2
|
||||
# expected_manifest = {
|
||||
# "registrar_manifest": []
|
||||
# }
|
||||
#
|
||||
# reg: ActionRegistrar = ActionRegistrarYAMLReader.read(
|
||||
# yaml_data, base_dir=os.path.join(TEST_DATA, "registries")
|
||||
# )
|
||||
#
|
||||
# self.assertEqual(expected_manifest, reg.manifest())
|
||||
# self.assertEqual(expected_size, reg.size())
|
||||
@@ -0,0 +1,65 @@
|
||||
from os import error
|
||||
|
||||
from yaml import dump
|
||||
|
||||
from playbook import models
|
||||
from playbook.action_registry import ActionRegistry
|
||||
from playbook.models import StepLog, ContextChecker
|
||||
|
||||
import sys
|
||||
import json
|
||||
import docker
|
||||
import subprocess
|
||||
from typing import List, Dict, Any, Tuple
|
||||
# from docker.models.containers import Container
|
||||
|
||||
|
||||
dockeractions = ActionRegistry("docker_actions_reg")
|
||||
|
||||
|
||||
@dockeractions.register(name="get_service_container_name", version="")
|
||||
@ContextChecker.requires("service_name")
|
||||
def get_service_container_name(ctx, name) -> StepLog:
|
||||
client = docker.from_env()
|
||||
|
||||
filters: Dict[str, Any]= {
|
||||
"label": [f"com.docker.compose.service={ctx["service_name"]}"]
|
||||
}
|
||||
containers = client.containers.list(filters=filters, all=True)
|
||||
|
||||
if not containers:
|
||||
return StepLog.fail(name, [{
|
||||
"status": "failed",
|
||||
"output": f"no container found for service {ctx["service_name"]}"}
|
||||
])
|
||||
|
||||
new_data = {"container": containers[0].id}
|
||||
return StepLog.ok(name, f"found container {containers[0].id}", pipe_ctx=new_data)
|
||||
|
||||
|
||||
@dockeractions.register(name="dump_container_pg_db", version="")
|
||||
@ContextChecker.requires("container", "dump_path", "db_user", "database")
|
||||
def dump_container_pg_db(ctx, name) -> StepLog:
|
||||
client = docker.from_env()
|
||||
container = client.containers.get(ctx["container"])
|
||||
|
||||
res = container.exec_run(
|
||||
cmd=f"pg_dump -U {ctx["db_user"]} {ctx["database"]}",
|
||||
stream=False, # change to True to stream directly to host file (need to check)
|
||||
demux=True,
|
||||
)
|
||||
stdout, stderr = res.output
|
||||
|
||||
if stderr:
|
||||
if not isinstance(stderr, bytes):
|
||||
return StepLog.fail(name, [{"status": "failed", "output": "stderr is not bytes"}])
|
||||
if res.exit_code != 0:
|
||||
error_msg = stderr.decode("utf-8") if stderr else "dump failed with no output"
|
||||
return StepLog.fail(name, [{"status": "failed", "output": error_msg}])
|
||||
|
||||
if not isinstance(stdout, bytes):
|
||||
return StepLog.fail(name, [{"status": "failed", "output": "stdout is not bytes"}])
|
||||
with open(ctx["dump_path"], "wb") as f:
|
||||
f.write(stdout)
|
||||
|
||||
return StepLog.ok(name, "database dump successful")
|
||||
@@ -0,0 +1,167 @@
|
||||
# restic actions registry
|
||||
from playbook.action_registry import ActionRegistry
|
||||
from playbook.models import ContextChecker, StepLog
|
||||
|
||||
import sys
|
||||
import json
|
||||
import subprocess
|
||||
from typing import List, Dict, Any, Tuple
|
||||
|
||||
|
||||
# NOTE:
|
||||
# most restic run commands seem to be the same pattern of:
|
||||
# restic command -> parse output -> grab errors
|
||||
# we can create a class for this I reckon
|
||||
|
||||
|
||||
restic = ActionRegistry("restic_actions_reg")
|
||||
|
||||
|
||||
def parse_output(line: str) -> List[Dict[str, Any]]:
|
||||
errors = []
|
||||
for line in line.split('\n'):
|
||||
if line.startswith("{"): # } treesitter is borked lmao
|
||||
data = json.loads(line)
|
||||
errors.append(data)
|
||||
return errors
|
||||
|
||||
|
||||
def find_restic_errors(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
errors = []
|
||||
for s in data:
|
||||
if s["message_type"] == "exit_error":
|
||||
errors.append(s)
|
||||
return errors
|
||||
|
||||
|
||||
def run_restic_command(cmd: List[str]) -> Tuple[int, List[Dict[str, Any]]]:
|
||||
"""
|
||||
Executes a restic command, streams output live to terminal,
|
||||
and returns (returncode, parsed_json_objects).
|
||||
"""
|
||||
parsed_json = []
|
||||
|
||||
# Start process with combined stdout/stderr pipe
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1
|
||||
)
|
||||
|
||||
# NOTE: is this what we want?
|
||||
assert process.stdout is not None
|
||||
|
||||
# Read line-by-line in real-time
|
||||
for line in iter(process.stdout.readline, ''):
|
||||
# Write live output directly to terminal screen
|
||||
sys.stdout.write(line)
|
||||
sys.stdout.flush()
|
||||
|
||||
# Parse and capture valid JSON objects on the fly
|
||||
clean_line = line.strip()
|
||||
if clean_line.startswith('{'): # }
|
||||
try:
|
||||
data = json.loads(clean_line)
|
||||
parsed_json.append(data)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
process.stdout.close()
|
||||
returncode = process.wait()
|
||||
return returncode, parsed_json
|
||||
|
||||
|
||||
@restic.register(name="check_restic_environment", version="")
|
||||
def check_restic_environment(ctx, name: str = "") -> StepLog:
|
||||
for v in ["passwordFile", "repoPath", "sourcePath", "serviceTag"]:
|
||||
if not ctx.get(v):
|
||||
return StepLog.fail(name, [f"{v} variable is empty"])
|
||||
|
||||
return StepLog.ok(name, "environment configuration seems to be valid")
|
||||
|
||||
|
||||
@restic.register(name="create_restic_repo", version="")
|
||||
def create_restic_repo(ctx, name: str = "") -> StepLog:
|
||||
cmd = ["restic", "init", "--json", "-r", ctx["repoPath"],
|
||||
"--password-file", ctx["passwordFile"]]
|
||||
|
||||
returncode, json_output = run_restic_command(cmd)
|
||||
|
||||
errors: List[Dict[str, Any]] = find_restic_errors(json_output)
|
||||
|
||||
# Handle error messages in JSON output
|
||||
if len(errors) > 0:
|
||||
return StepLog.fail(name, errors)
|
||||
|
||||
# Make sure the return code is good too
|
||||
if returncode != 0:
|
||||
return StepLog.fail(name, [{"status": "failed", "output": json_output}])
|
||||
|
||||
return StepLog.ok(name, "repo create succesfully")
|
||||
|
||||
|
||||
@restic.register(name="create_repo_if_not_exists", version="")
|
||||
def create_repo_if_not_exists(ctx, name: str = "") -> StepLog:
|
||||
log: StepLog = check_if_repo_exists(ctx, "")
|
||||
|
||||
msg: str = ""
|
||||
if log.failed:
|
||||
# Repo doesn't exist is error code 10
|
||||
if log.error[0]["code"] != 10:
|
||||
return StepLog.fail(name, [{"status": "failed", "output": log.error}])
|
||||
|
||||
# Create repo
|
||||
repoCreateLog: StepLog = create_restic_repo(ctx, name + f".{create_restic_repo.__name__}")
|
||||
if repoCreateLog.failed:
|
||||
return StepLog.fail(name, [{"status": "failed", "output": repoCreateLog.error}])
|
||||
msg = "repo created"
|
||||
else:
|
||||
msg = "repo already exists"
|
||||
|
||||
|
||||
return StepLog.ok(name, msg)
|
||||
|
||||
|
||||
@restic.register(name="check_if_repo_exists", version="")
|
||||
def check_if_repo_exists(ctx, name: str = "") -> StepLog:
|
||||
cmd = ["restic", "-r", ctx["repoPath"], "cat", "config", "--json",
|
||||
"--password-file", ctx["passwordFile"]]
|
||||
returncode, json_output = run_restic_command(cmd)
|
||||
|
||||
errors: List[Dict[str, Any]] = find_restic_errors(json_output)
|
||||
|
||||
# Handle error messages in JSON output
|
||||
if len(errors) > 0:
|
||||
return StepLog.fail(name, errors)
|
||||
|
||||
# Make sure the return code is good too
|
||||
if returncode != 0:
|
||||
return StepLog.fail(name, [{"status": "failed", "output": json_output}])
|
||||
|
||||
return StepLog.ok(name, "repo seems to exist")
|
||||
|
||||
|
||||
@restic.register(name="backup_data_to_restic_repo", version="")
|
||||
@ContextChecker.requires("source_path", "passwordFile")
|
||||
def backup_data_to_restic_repo(ctx, name) -> StepLog:
|
||||
cmd = ["restic", "backup", ctx["source_path"], "--json", "--quiet",
|
||||
"-r", ctx["repoPath"], "--password-file", ctx["passwordFile"]]
|
||||
if ctx["tags"]:
|
||||
cmd.extend(["--tag", ','.join(ctx["tags"])])
|
||||
|
||||
returncode, json_output = run_restic_command(cmd)
|
||||
print(json_output, file=sys.stderr)
|
||||
errors: List[Dict[str, Any]] = find_restic_errors(json_output)
|
||||
|
||||
# Handle error messages in JSON output
|
||||
if len(errors) > 0:
|
||||
return StepLog.fail(name, errors)
|
||||
|
||||
# Make sure the return code is good too
|
||||
if returncode != 0:
|
||||
return StepLog.fail(name, [{"status": "failed", "output": errors}])
|
||||
|
||||
return StepLog.ok(name, "backup successful")
|
||||
Reference in New Issue
Block a user