ddeutil-workflow 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

ddeutil/workflow/stage.py

@@ -3,6 +3,18 @@
 # Licensed under the MIT License. See LICENSE in the project root for
 # license information.
 # ------------------------------------------------------------------------------
+"""Stage Model that use for getting stage data template from Job Model.
+The stage that handle the minimize task that run in some thread (same thread at
+its job owner) that mean it is the lowest executor of a pipeline workflow that
+can tracking logs.
+
+    The output of stage execution only return 0 status because I do not want to
+handle stage error on this stage model. I think stage model should have a lot of
+usecase and it does not worry when I want to create a new one.
+
+    Execution --> Ok    --> Result with 0
+              --> Error --> Raise StageException
+"""
 from __future__ import annotations
 
 import contextlib
@@ -15,13 +27,22 @@ import uuid
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
 from dataclasses import dataclass
+from functools import wraps
 from inspect import Parameter
 from pathlib import Path
 from subprocess import CompletedProcess
+from textwrap import dedent
 from typing import Callable, Optional, Union
 
+try:
+    from typing import ParamSpec
+except ImportError:
+    from typing_extensions import ParamSpec
+
 from ddeutil.core import str2bool
 from pydantic import BaseModel, Field
+from pydantic.functional_validators import model_validator
+from typing_extensions import Self
 
 from .__types import DictData, DictStr, Re, TupleStr
 from .exceptions import StageException
@@ -32,9 +53,68 @@ from .utils import (
     gen_id,
     make_exec,
     make_registry,
+    not_in_template,
     param2template,
 )
 
+P = ParamSpec("P")
+__all__: TupleStr = (
+    "Stage",
+    "EmptyStage",
+    "BashStage",
+    "PyStage",
+    "HookStage",
+    "TriggerStage",
+    "handler_result",
+)
+
+
+def handler_result(message: str | None = None) -> Callable[P, Result]:
+    """Decorator function for handler result from the stage execution. This
+    function should to use with execution method only.
+
+    :param message: A message that want to add at prefix of exception statement.
+    """
+    message: str = message or ""
+
+    def decorator(func: Callable[P, Result]) -> Callable[P, Result]:
+
+        @wraps(func)
+        def wrapped(self: Stage, *args, **kwargs):
+            try:
+                # NOTE: Start calling origin function with a passing args.
+                return func(self, *args, **kwargs).set_run_id(self.run_id)
+            except Exception as err:
+                # NOTE: Start catching error from the stage execution.
+                logging.error(
+                    f"({self.run_id}) [STAGE]: {err.__class__.__name__}: {err}"
+                )
+                if str2bool(
+                    os.getenv("WORKFLOW_CORE_STAGE_RAISE_ERROR", "true")
+                ):
+                    # NOTE: If error that raise from stage execution course by
+                    #   itself, it will return that error with previous
+                    #   dependency.
+                    if isinstance(err, StageException):
+                        raise StageException(
+                            f"{self.__class__.__name__}: {message}\n\t{err}"
+                        ) from err
+                    raise StageException(
+                        f"{self.__class__.__name__}: {message}\n\t"
+                        f"{err.__class__.__name__}: {err}"
+                    ) from None
+                rs: Result = Result(
+                    status=1,
+                    context={
+                        "error_message": f"{err.__class__.__name__}: {err}",
+                    },
+                )
+                return rs.set_run_id(self.run_id)
+
+        return wrapped
+
+    return decorator
+
 
 class BaseStage(BaseModel, ABC):
     """Base Stage Model that keep only id and name fields for the stage
@@ -50,12 +130,45 @@ class BaseStage(BaseModel, ABC):
         ),
     )
     name: str = Field(
-        description="A stage name that want to logging when start execution."
+        description="A stage name that want to logging when start execution.",
     )
     condition: Optional[str] = Field(
         default=None,
+        description="A stage condition statement to allow stage executable.",
         alias="if",
     )
+    run_id: Optional[str] = Field(
+        default=None,
+        description="A running stage ID.",
+        repr=False,
+    )
+
+    @model_validator(mode="after")
+    def __prepare_running_id(self):
+        """Prepare stage running ID that use default value of field and this
+        method will validate name and id fields should not contain any template
+        parameter (exclude matrix template).
+        """
+        if self.run_id is None:
+            self.run_id = gen_id(self.name + (self.id or ""), unique=True)
+
+        # VALIDATE: Validate stage id and name should not dynamic with params
+        #   template. (allow only matrix)
+        if not_in_template(self.id) or not_in_template(self.name):
+            raise ValueError(
+                "Stage name and ID should only template with matrix."
+            )
+
+        return self
+
+    def get_running_id(self, run_id: str) -> Self:
+        """Return Stage model object that changing stage running ID with an
+        input running ID.
+
+        :param run_id: A replace stage running ID.
+        :rtype: Self
+        """
+        return self.model_copy(update={"run_id": run_id})
 
     @abstractmethod
     def execute(self, params: DictData) -> Result:
@@ -67,31 +180,45 @@ class BaseStage(BaseModel, ABC):
         """
         raise NotImplementedError("Stage should implement ``execute`` method.")
 
-    def set_outputs(self, output: DictData, params: DictData) -> DictData:
+    def set_outputs(self, output: DictData, to: DictData) -> DictData:
         """Set an outputs from execution process to an input params.
 
         :param output: A output data that want to extract to an output key.
-        :param params: A context data that want to add output result.
+        :param to: A context data that want to add output result.
         :rtype: DictData
         """
-        if self.id:
-            _id: str = param2template(self.id, params)
-        elif str2bool(os.getenv("WORKFLOW_CORE_DEFAULT_STAGE_ID", "false")):
-            _id: str = gen_id(param2template(self.name, params))
-        else:
-            return params
+        if not (
+            self.id
+            or str2bool(os.getenv("WORKFLOW_CORE_STAGE_DEFAULT_ID", "false"))
+        ):
+            logging.debug(
+                f"({self.run_id}) [STAGE]: Output does not set because this "
+                f"stage does not set ID or default stage ID config flag not be "
+                f"True."
+            )
+            return to
 
         # NOTE: Create stages key to receive an output from the stage execution.
-        if "stages" not in params:
-            params["stages"] = {}
+        if "stages" not in to:
+            to["stages"] = {}
 
-        params["stages"][_id] = {"outputs": output}
-        return params
+        if self.id:
+            _id: str = param2template(self.id, params=to)
+        else:
+            _id: str = gen_id(param2template(self.name, params=to))
+
+        # NOTE: Set the output to that stage generated ID.
+        logging.debug(
+            f"({self.run_id}) [STAGE]: Set output complete with stage ID: {_id}"
+        )
+        to["stages"][_id] = {"outputs": output}
+        return to
 
-    def is_skip(self, params: DictData | None = None) -> bool:
+    def is_skipped(self, params: DictData | None = None) -> bool:
         """Return true if condition of this stage do not correct.
 
         :param params: A parameters that want to pass to condition template.
+        :rtype: bool
         """
         params: DictData = params or {}
         if self.condition is None:
@@ -104,8 +231,8 @@ class BaseStage(BaseModel, ABC):
                 raise TypeError("Return type of condition does not be boolean")
             return not rs
         except Exception as err:
-            logging.error(str(err))
-            raise StageException(str(err)) from err
+            logging.error(f"({self.run_id}) [STAGE]: {err}")
+            raise StageException(f"{err.__class__.__name__}: {err}") from err
 
 
 class EmptyStage(BaseStage):
@@ -131,8 +258,10 @@ class EmptyStage(BaseStage):
         :param params: A context data that want to add output result. But this
             stage does not pass any output.
         """
-        stm: str = param2template(self.echo, params=params) or "..."
-        logging.info(f"[STAGE]: Empty-Execute: {self.name!r}: " f"( {stm} )")
+        logging.info(
+            f"({self.run_id}) [STAGE]: Empty-Execute: {self.name!r}: "
+            f"( {param2template(self.echo, params=params) or '...'} )"
+        )
         return Result(status=0, context={})
 
 
@@ -174,20 +303,28 @@ class BashStage(BaseStage):
         f_shebang: str = "bash" if sys.platform.startswith("win") else "sh"
         with open(f"./{f_name}", mode="w", newline="\n") as f:
             # NOTE: write header of `.sh` file
-            f.write(f"#!/bin/{f_shebang}\n")
+            f.write(f"#!/bin/{f_shebang}\n\n")
 
             # NOTE: add setting environment variable before bash skip statement.
             f.writelines([f"{k}='{env[k]}';\n" for k in env])
 
             # NOTE: make sure that shell script file does not have `\r` char.
-            f.write(bash.replace("\r\n", "\n"))
+            f.write("\n" + bash.replace("\r\n", "\n"))
 
+        # NOTE: Make this .sh file able to executable.
         make_exec(f"./{f_name}")
 
+        logging.debug(
+            f"({self.run_id}) [STAGE]: Start create `.sh` file and running a "
+            f"bash statement."
+        )
+
         yield [f_shebang, f_name]
 
+        # Note: Remove .sh file that use to run bash.
         Path(f"./{f_name}").unlink()
 
+    @handler_result()
     def execute(self, params: DictData) -> Result:
         """Execute the Bash statement with the Python build-in ``subprocess``
         package.
@@ -195,11 +332,11 @@ class BashStage(BaseStage):
         :param params: A parameter data that want to use in this execution.
         :rtype: Result
         """
-        bash: str = param2template(self.bash, params)
+        bash: str = param2template(dedent(self.bash), params)
         with self.__prepare_bash(
             bash=bash, env=param2template(self.env, params)
         ) as sh:
-            logging.info(f"[STAGE]: Shell-Execute: {sh}")
+            logging.info(f"({self.run_id}) [STAGE]: Shell-Execute: {sh}")
             rs: CompletedProcess = subprocess.run(
                 sh,
                 shell=False,
@@ -211,9 +348,11 @@ class BashStage(BaseStage):
                 rs.stderr.encode("utf-8").decode("utf-16")
                 if "\\x00" in rs.stderr
                 else rs.stderr
+            ).removesuffix("\n")
+            raise StageException(
+                f"Subprocess: {err}\nRunning Statement:\n---\n"
+                f"```bash\n{bash}\n```"
             )
-            logging.error(f"{err}\n\n```bash\n{bash}```")
-            raise StageException(f"{err}\n\n```bash\n{bash}```")
         return Result(
             status=0,
             context={
@@ -227,6 +366,15 @@ class BashStage(BaseStage):
 class PyStage(BaseStage):
     """Python executor stage that running the Python statement that receive
     globals nad additional variables.
+
+    Data Validate:
+        >>> stage = {
+        ...     "name": "Python stage execution",
+        ...     "run": 'print("Hello {x}")',
+        ...     "vars": {
+        ...         "x": "BAR",
+        ...     },
+        ... }
     """
 
     run: str = Field(
@@ -239,26 +387,26 @@ class PyStage(BaseStage):
         ),
     )
 
-    def set_outputs(self, output: DictData, params: DictData) -> DictData:
+    def set_outputs(self, output: DictData, to: DictData) -> DictData:
         """Set an outputs from the Python execution process to an input params.
 
         :param output: A output data that want to extract to an output key.
-        :param params: A context data that want to add output result.
+        :param to: A context data that want to add output result.
         :rtype: DictData
         """
         # NOTE: The output will fileter unnecessary keys from locals.
         _locals: DictData = output["locals"]
         super().set_outputs(
-            {k: _locals[k] for k in _locals if k != "__annotations__"},
-            params=params,
+            {k: _locals[k] for k in _locals if k != "__annotations__"}, to=to
         )
 
         # NOTE:
         #   Override value that changing from the globals that pass via exec.
         _globals: DictData = output["globals"]
-        params.update({k: _globals[k] for k in params if k in _globals})
-        return params
+        to.update({k: _globals[k] for k in to if k in _globals})
+        return to
 
+    @handler_result()
     def execute(self, params: DictData) -> Result:
         """Execute the Python statement that pass all globals and input params
         to globals argument on ``exec`` build-in function.
@@ -266,34 +414,66 @@ class PyStage(BaseStage):
         :param params: A parameter that want to pass before run any statement.
         :rtype: Result
         """
+        # NOTE: Replace the run statement that has templating value.
+        run: str = param2template(dedent(self.run), params)
+
         # NOTE: create custom globals value that will pass to exec function.
         _globals: DictData = (
             globals() | params | param2template(self.vars, params)
         )
         _locals: DictData = {}
-        try:
-            logging.info(f"[STAGE]: Py-Execute: {uuid.uuid4()}")
-            exec(param2template(self.run, params), _globals, _locals)
-        except Exception as err:
-            raise StageException(
-                f"{err.__class__.__name__}: {err}\nRunning Statement:\n---\n"
-                f"{self.run}"
-            ) from None
+
+        # NOTE: Start exec the run statement.
+        logging.info(f"({self.run_id}) [STAGE]: Py-Execute: {self.name}")
+        exec(run, _globals, _locals)
+
         return Result(
-            status=0,
-            context={"locals": _locals, "globals": _globals},
+            status=0, context={"locals": _locals, "globals": _globals}
         )
 
 
 @dataclass
 class HookSearch:
-    """Hook Search dataclass."""
+    """Hook Search dataclass that use for receive regular expression grouping
+    dict from searching hook string value.
+    """
 
     path: str
     func: str
     tag: str
 
 
+def extract_hook(hook: str) -> Callable[[], TagFunc]:
+    """Extract Hook function from string value to hook partial function that
+    does run it at runtime.
+
+    :param hook: A hook value that able to match with Task regex.
+    :rtype: Callable[[], TagFunc]
+    """
+    if not (found := Re.RE_TASK_FMT.search(hook)):
+        raise ValueError(
+            f"Hook {hook!r} does not match with hook format regex."
+        )
+
+    # NOTE: Pass the searching hook string to `path`, `func`, and `tag`.
+    hook: HookSearch = HookSearch(**found.groupdict())
+
+    # NOTE: Registry object should implement on this package only.
+    rgt: dict[str, Registry] = make_registry(f"{hook.path}")
+    if hook.func not in rgt:
+        raise NotImplementedError(
+            f"``REGISTER-MODULES.{hook.path}.registries`` does not "
+            f"implement registry: {hook.func!r}."
+        )
+
+    if hook.tag not in rgt[hook.func]:
+        raise NotImplementedError(
+            f"tag: {hook.tag!r} does not found on registry func: "
+            f"``REGISTER-MODULES.{hook.path}.registries.{hook.func}``"
+        )
+    return rgt[hook.func][hook.tag]
+
+
 class HookStage(BaseStage):
     """Hook executor that hook the Python function from registry with tag
     decorator function in ``utils`` module and run it with input arguments.
@@ -306,7 +486,7 @@ class HookStage(BaseStage):
     Data Validate:
         >>> stage = {
         ...     "name": "Task stage execution",
-        ...     "task": "tasks/function-name@tag-name",
+        ...     "uses": "tasks/function-name@tag-name",
         ...     "args": {
         ...         "FOO": "BAR",
         ...     },
@@ -314,37 +494,15 @@ class HookStage(BaseStage):
     """
 
     uses: str = Field(
-        description="A pointer that want to load function from registry",
+        description="A pointer that want to load function from registry.",
+    )
+    args: DictData = Field(
+        default_factory=dict,
+        description="An arguments that want to pass to the hook function.",
+        alias="with",
     )
-    args: DictData = Field(alias="with")
-
-    @staticmethod
-    def extract_hook(hook: str) -> Callable[[], TagFunc]:
-        """Extract Hook string value to hook function.
-
-        :param hook: A hook value that able to match with Task regex.
-        """
-        if not (found := Re.RE_TASK_FMT.search(hook)):
-            raise ValueError("Task does not match with task format regex.")
-
-        # NOTE: Pass the searching hook string to `path`, `func`, and `tag`.
-        hook: HookSearch = HookSearch(**found.groupdict())
-
-        # NOTE: Registry object should implement on this package only.
-        rgt: dict[str, Registry] = make_registry(f"{hook.path}")
-        if hook.func not in rgt:
-            raise NotImplementedError(
-                f"``REGISTER-MODULES.{hook.path}.registries`` does not "
-                f"implement registry: {hook.func!r}."
-            )
-
-        if hook.tag not in rgt[hook.func]:
-            raise NotImplementedError(
-                f"tag: {hook.tag!r} does not found on registry func: "
-                f"``REGISTER-MODULES.{hook.path}.registries.{hook.func}``"
-            )
-        return rgt[hook.func][hook.tag]
 
+    @handler_result()
     def execute(self, params: DictData) -> Result:
         """Execute the Hook function that already in the hook registry.
 
@@ -352,9 +510,8 @@ class HookStage(BaseStage):
         :type params: DictData
         :rtype: Result
         """
-        t_func: TagFunc = self.extract_hook(param2template(self.uses, params))()
-        if not callable(t_func):
-            raise ImportError("Hook caller function does not callable.")
+        t_func_hook: str = param2template(self.uses, params)
+        t_func: TagFunc = extract_hook(t_func_hook)()
 
         # VALIDATE: check input task caller parameters that exists before
         #   calling.
@@ -366,59 +523,68 @@ class HookStage(BaseStage):
             if ips.parameters[k].default == Parameter.empty
         ):
             raise ValueError(
-                f"Necessary params, ({', '.join(ips.parameters.keys())}), "
+                f"Necessary params, ({', '.join(ips.parameters.keys())}, ), "
                 f"does not set to args"
             )
-
         # NOTE: add '_' prefix if it want to use.
         for k in ips.parameters:
             if k.removeprefix("_") in args:
                 args[k] = args.pop(k.removeprefix("_"))
 
-        try:
-            logging.info(f"[STAGE]: Hook-Execute: {t_func.name}@{t_func.tag}")
-            rs: DictData = t_func(**param2template(args, params))
-        except Exception as err:
-            raise StageException(f"{err.__class__.__name__}: {err}") from err
+        logging.info(
+            f"({self.run_id}) [STAGE]: Hook-Execute: {t_func.name}@{t_func.tag}"
+        )
+        rs: DictData = t_func(**param2template(args, params))
 
-        # VALIDATE: Check the result type from hook function, it should be dict.
+        # VALIDATE:
+        #   Check the result type from hook function, it should be dict.
         if not isinstance(rs, dict):
-            raise StageException(
-                f"Return of hook function: {t_func.name}@{t_func.tag} does not "
-                f"serialize to result model, you should fix it to `dict` type."
+            raise TypeError(
+                f"Return type: '{t_func.name}@{t_func.tag}' does not serialize "
+                f"to result model, you change return type to `dict`."
             )
         return Result(status=0, context=rs)
 
 
 class TriggerStage(BaseStage):
-    """Trigger Pipeline execution stage that execute another pipeline object."""
+    """Trigger Pipeline execution stage that execute another pipeline object.
+
+    Data Validate:
+        >>> stage = {
+        ...     "name": "Trigger pipeline stage execution",
+        ...     "trigger": 'pipeline-name-for-loader',
+        ...     "params": {
+        ...         "run-date": "2024-08-01",
+        ...         "source": "src",
+        ...     },
+        ... }
+    """
 
     trigger: str = Field(description="A trigger pipeline name.")
-    params: DictData = Field(default_factory=dict)
+    params: DictData = Field(
+        default_factory=dict,
+        description="A parameter that want to pass to pipeline execution.",
+    )
 
+    @handler_result("Raise from TriggerStage")
     def execute(self, params: DictData) -> Result:
-        """Trigger execution.
+        """Trigger pipeline execution.
 
         :param params: A parameter data that want to use in this execution.
         :rtype: Result
         """
-        from .exceptions import PipelineException
         from .pipeline import Pipeline
 
-        try:
-            # NOTE: Loading pipeline object from trigger name.
-            pipe: Pipeline = Pipeline.from_loader(
-                name=self.trigger, externals={}
-            )
-            rs: Result = pipe.execute(
-                params=param2template(self.params, params)
-            )
-        except PipelineException as err:
-            _alias_stage: str = self.id or self.name
-            raise StageException(
-                f"Trigger Stage: {_alias_stage} get trigger pipeline exception."
-            ) from err
-        return rs
+        # NOTE: Loading pipeline object from trigger name.
+        _trigger: str = param2template(self.trigger, params=params)
+
+        # NOTE: Set running pipeline ID from running stage ID to external
+        #   params on Loader object.
+        pipe: Pipeline = Pipeline.from_loader(
+            name=_trigger, externals={"run_id": self.run_id}
+        )
+        logging.info(f"({self.run_id}) [STAGE]: Trigger-Execute: {_trigger!r}")
+        return pipe.execute(params=param2template(self.params, params))
 
 
 # NOTE: Order of parsing stage data