context checker decorator

playbook.yaml

@@ -17,12 +17,18 @@ steps:
     context:
       serviceName: "gitea-db"
       tmpPathPrefix: "/tmp/db_dump.sql"
-  - name: "backup db to restic repo"
+  - name: "test context"
     actions:
-      dump_database: "backup_data_to_restic_repo"
+      dump_database: "dump_pg_database"
     context:
-      sourcePath: "/tmp/db_dump.sql"
-      serviceTag: "gitea-db"
+      serviceName: "gitea-db"
+      tmpPathPrefix: "/tmp/db_dump.sql"
+  # - name: "backup db to restic repo"
+  #   actions:
+  #     dump_database: "backup_data_to_restic_repo"
+  #   context:
+  #     sourcePath: "/tmp/db_dump.sql"
+  #     serviceTag: "gitea-db"
   # - name: "backup rusty-packages"
   #   actions:
   #     create_repo_dir_if_not_exists: "create_dir_if_not_exists"

playbook/__init__.py

@@ -7,7 +7,7 @@ import json
 import importlib.util
 from datetime import datetime, timezone
 from typing import Dict, List, Any, Optional, Tuple
-from dataclasses import asdict, dataclass
+from dataclasses import MISSING, asdict, dataclass
 from abc import ABC, abstractmethod
 
 from playbook.models import Status, StepLog, timed_run
@@ -30,7 +30,7 @@ class StepIF(ABC):
 class CustomStep(StepIF):
     """ Setup custom steps. """
 
-    def __init__( self, name: str, actions: List[ActionFn], context: Dict[str, Any]):
+    def __init__(self, name: str, actions: List[ActionFn], context: Dict[str, Any]):
         self.name: str = name
         self.__actions = actions
         self.__context = context
@@ -43,9 +43,17 @@ class CustomStep(StepIF):
         """ 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:
-            log: StepLog = timed_run(op, ctx | self.__context, op.__name__)
+            log: StepLog = timed_run(op, ctx | prev_step_ctx | self.__context, op.__name__)
             status = Status.BAD if log.failed else Status.GOOD
+            prev_step_ctx = log.pipe_ctx
 
             substeps.append(log)
 
@@ -55,13 +63,15 @@ class CustomStep(StepIF):
         errors = []
         msg: str = ""
         if status == Status.BAD:
-            errors = ["substep failed"]
+            # 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,
+                msg=msg, error=errors, substeps=substeps, pipe_ctx=prev_step_ctx
                 )
 
 
@@ -70,20 +80,6 @@ class PlaybookError(Exception):
     pass
 
 
-class ContextWrapper:
-    """Safely extract values with error messages."""
-    def __init__(self, ctx: Dict[str, Any]):
-        self._ctx = ctx
-
-    def get(self, key: str, default: Any = None) -> Any:
-        return self._ctx.get(key, default)
-    
-    def require(self, *keys: str) -> Tuple[List[str], Dict[str, Any]]:
-        missing = [k for k in keys if not self._ctx.get(k)]
-        values = {k: self._ctx.get(k) for k in keys if k not in missing}
-        return missing, values
-
-
 @dataclass
 class StepEntry(object):
     name: str
@@ -184,8 +180,10 @@ class Play(object):
                 substeps=[],
                 )
 
+        prev_ctx: Dict[str, Any] = {}
         for s in self._steps:
-            log: StepLog = timed_run(s.step.run, self._context)
+            log: StepLog = timed_run(s.step.run, self._context | prev_ctx)
+            prev_ctx = log.pipe_ctx
 
             playLog.substeps.append(log)
 
@@ -217,7 +215,7 @@ class Play(object):
             except KeyError as e:
                 raise PlaybookError(f"Step '{name}' is missing required action field: {e}") from e
 
-        return CustomStep(name, functionList, self._context | ctx)
+        return CustomStep(name, functionList, ctx)
 
     # NOTE: Move to factory
     @classmethod

playbook/action_registry/__init__.py

@@ -1,11 +1,8 @@
 from dataclasses import dataclass
-from playbook.models import StepLog
+from playbook.models import StepLog, ActionFn
 from typing import Callable, Dict, Any
 
 
-ActionFn = Callable[[Dict[str, Any], str], StepLog]
-
-
 @dataclass
 class Action(object):
     name: str

playbook/models.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import time
 from enum import StrEnum
 from typing import List, Dict, Any, Callable
@@ -33,14 +35,33 @@ class StepLog(object):
         return len(self.error) > 0 or self.status == Status.BAD
 
     @classmethod
-    def ok(cls, name: str, msg: str = "success") -> StepLog:
-        return cls(step_name=name, status=Status.GOOD, msg=msg)
+    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)
 
     @classmethod
     def fail(cls, name: str, errors: List[Any]) -> StepLog:
         return cls(step_name=name, status=Status.BAD, error=errors)
 
 
+class ContextChecker:
+    """Safely extract values with error messages."""
+    def __init__(self, ctx: Dict[str, Any]):
+        self._ctx = ctx
+
+    @classmethod
+    def requires(cls, *args):
+        def decorator(func: ActionFn):
+            def wrapper(ctx: Dict[str, Any], name: str) -> StepLog:
+                if all(key in ctx for key in args):
+                    return func(ctx, name)
+                else:
+                    missing_keys = [key for key in args if key not in ctx]
+                    missing_keys_msg = f"missing context keys: {missing_keys}"
+                    return StepLog.fail(name, [{"status": "failed", "output": missing_keys_msg}])
+            return wrapper
+        return decorator
+
+
 def human_readable_date(time: int) -> str:
     seconds = time // 1_000_000_000
     nanos = time % 1_000_000_000
@@ -68,3 +89,6 @@ def timed_run(op: Callable[..., StepLog], *args: Any, **kwargs: Any) -> StepLog:
         log.end_date = human_readable_date(end_date)
     
     return log
+
+
+ActionFn = Callable[[Dict[str, Any], str], StepLog]

registry/dbpg_actions.py

@@ -1,7 +1,7 @@
 # restic actions registry
 from playbook.action_registry import ActionRegistry
 from playbook.models import StepLog
-from playbook import ContextWrapper
+from playbook.models import ContextChecker
 
 from registry.docker_actions import dockeractions
 
@@ -16,6 +16,7 @@ dbpgactions = ActionRegistry("dbpg_actions_reg")
  
  
 @dbpgactions.register(name="dump_pg_database", version="")
+@ContextChecker.requires("database")
 def dump_pg_database(ctx, name) -> StepLog:
-    c = ContextWrapper(ctx)
-    return StepLog.ok(name, "")
+    print(ctx, file=sys.stderr)
+    return StepLog.ok(name, "", pipe_ctx={"new_data": "coolData"})