Source code for dagster.core.definitions.decorators.hook

from functools import update_wrapper
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Set, Union, cast

from dagster import check
from dagster.core.errors import DagsterInvalidDefinitionError

from ...decorator_utils import split_function_parameters, validate_decorated_fn_positionals
from ..events import HookExecutionResult
from ..hook import HookDefinition

if TYPE_CHECKING:
    from dagster.core.execution.context.system import HookContext
    from dagster.core.events import DagsterEvent


class _Hook:
    def __init__(
        self, name: Optional[str] = None, required_resource_keys: Optional[Set[str]] = None
    ):
        self.name = check.opt_str_param(name, "name")
        self.required_resource_keys = check.opt_set_param(
            required_resource_keys, "required_resource_keys"
        )

    def __call__(self, fn: Callable[["HookContext", List["DagsterEvent"]], Any]) -> HookDefinition:

        check.callable_param(fn, "fn")

        if not self.name:
            self.name = fn.__name__

        expected_positionals = ["context", "event_list"]
        fn_positionals, _ = split_function_parameters(fn, expected_positionals)
        missing_positional = validate_decorated_fn_positionals(fn_positionals, expected_positionals)
        if missing_positional:
            raise DagsterInvalidDefinitionError(
                "'{hook_name}' decorated function does not have required positional "
                "parameter '{missing_param}'. Hook functions should only have keyword arguments "
                "that match input names and a first positional parameter named 'context' and "
                "a second positional parameter named 'event_list'.".format(
                    hook_name=fn.__name__, missing_param=missing_positional
                )
            )

        hook_def = HookDefinition(
            name=self.name,
            hook_fn=fn,
            required_resource_keys=self.required_resource_keys,
        )
        update_wrapper(cast(Callable[..., Any], hook_def), fn)
        return hook_def


def event_list_hook(
    name: Union[Optional[str], Callable[..., Any]] = None,
    required_resource_keys: Optional[Set[str]] = None,
) -> Union[HookDefinition, _Hook]:
    """Create a generic hook with the specified parameters from the decorated function.

    This decorator is currently used internally by Dagster machinery to support success_hook and
    failure_hook.

    The user-defined hook function requires two parameters:
    - A `context` object is passed as the first parameter. The context is an instance of
        :py:class:`context <HookContext>`, and provides access to system
        information, such as loggers (context.log), resources (context.resources), the solid
        (context.solid) and its execution step (context.step) which triggers this hook.
    - An `event_list` object is passed as the second paramter. It provides the full event list of the
        associated execution step.

    Args:
        name (Optional[str]): The name of this hook.
        required_resource_keys (Optional[Set[str]]): Keys for the resources required by the
            hook.

    Examples:

        .. code-block:: python

            @event_list_hook(required_resource_keys={'slack'})
            def slack_on_materializations(context, event_list):
                for event in event_list:
                    if event.event_type == DagsterEventType.ASSET_MATERIALIZATION:
                        message = '{solid} has materialized an asset {key}.'.format(
                            solid=context.solid.name,
                            key=event.asset_key
                        )
                        # send a slack message every time a materialization event occurs
                        context.resources.slack.send_message(message)


    """
    # This case is for when decorator is used bare, without arguments.
    # e.g. @event_list_hook versus @event_list_hook()
    if callable(name):
        check.invariant(required_resource_keys is None)
        return _Hook()(name)

    return _Hook(name=name, required_resource_keys=required_resource_keys)


[docs]def success_hook(
    name: Union[Optional[str], Callable[..., Any]] = None,
    required_resource_keys: Optional[Set[str]] = None,
) -> Union[
    Union[HookDefinition, _Hook],
    Callable[[Callable[["HookContext"], Any]], Union[HookDefinition, _Hook]],
]:
    """Create a hook on step success events with the specified parameters from the decorated function.

    Args:
        name (Optional[str]): The name of this hook.
        required_resource_keys (Optional[Set[str]]): Keys for the resources required by the
            hook.

    Examples:

        .. code-block:: python

            @success_hook(required_resource_keys={'slack'})
            def slack_message_on_success(context):
                message = 'solid {} succeeded'.format(context.solid.name)
                context.resources.slack.send_message(message)

            @success_hook
            def do_something_on_success(context):
                do_something()


    """

    def wrapper(fn: Callable[["HookContext"], Any]) -> Union[HookDefinition, _Hook]:

        check.callable_param(fn, "fn")

        expected_positionals = ["context"]
        fn_positionals, _ = split_function_parameters(fn, expected_positionals)
        missing_positional = validate_decorated_fn_positionals(fn_positionals, expected_positionals)
        if missing_positional:
            raise DagsterInvalidDefinitionError(
                "@success_hook '{hook_name}' decorated function does not have required positional "
                "parameter '{missing_param}'. Hook functions should only have keyword arguments "
                "that match input names and a first positional parameter named 'context'.".format(
                    hook_name=fn.__name__, missing_param=missing_positional
                )
            )

        if name is None or callable(name):
            _name = fn.__name__
        else:
            _name = name

        @event_list_hook(_name, required_resource_keys)
        def _success_hook(
            context: "HookContext", event_list: List["DagsterEvent"]
        ) -> HookExecutionResult:
            for event in event_list:
                if event.is_step_success:
                    fn(context)
                    return HookExecutionResult(hook_name=_name, is_skipped=False)

            # hook is skipped when fn didn't run
            return HookExecutionResult(hook_name=_name, is_skipped=True)

        return _success_hook

    # This case is for when decorator is used bare, without arguments, i.e. @success_hook
    if callable(name):
        check.invariant(required_resource_keys is None)
        return wrapper(name)

    return wrapper


[docs]def failure_hook(
    name: Optional[str] = None, required_resource_keys: Optional[Set[str]] = None
) -> Union[
    Union[HookDefinition, _Hook],
    Callable[[Callable[["HookContext"], Any]], Union[HookDefinition, _Hook]],
]:
    """Create a hook on step failure events with the specified parameters from the decorated function.

    Args:
        name (Optional[str]): The name of this hook.
        required_resource_keys (Optional[Set[str]]): Keys for the resources required by the
            hook.

    Examples:

        .. code-block:: python

            @failure_hook(required_resource_keys={'slack'})
            def slack_message_on_failure(context):
                message = 'solid {} failed'.format(context.solid.name)
                context.resources.slack.send_message(message)

            @failure_hook
            def do_something_on_failure(context):
                do_something()


    """

    def wrapper(fn: Callable[["HookContext"], Any]) -> Union[HookDefinition, _Hook]:
        check.callable_param(fn, "fn")

        expected_positionals = ["context"]
        fn_positionals, _ = split_function_parameters(fn, expected_positionals)
        missing_positional = validate_decorated_fn_positionals(fn_positionals, expected_positionals)
        if missing_positional:
            raise DagsterInvalidDefinitionError(
                "@failure_hook '{hook_name}' decorated function does not have required positional "
                "parameter '{missing_param}'. Hook functions should only have keyword arguments "
                "that match input names and a first positional parameter named 'context'.".format(
                    hook_name=fn.__name__, missing_param=missing_positional
                )
            )

        if name is None or callable(name):
            _name = fn.__name__
        else:
            _name = name

        @event_list_hook(_name, required_resource_keys)
        def _failure_hook(
            context: "HookContext", event_list: List["DagsterEvent"]
        ) -> HookExecutionResult:
            for event in event_list:
                if event.is_step_failure:
                    fn(context)
                    return HookExecutionResult(hook_name=_name, is_skipped=False)

            # hook is skipped when fn didn't run
            return HookExecutionResult(hook_name=_name, is_skipped=True)

        return _failure_hook

    # This case is for when decorator is used bare, without arguments, i.e. @failure_hook
    if callable(name):
        check.invariant(required_resource_keys is None)
        return wrapper(name)

    return wrapper