Support defining pre/post hook for operators in a DAG at the DAG level

[sighingnow]

Description

We propose to add such features to airflow, which will allow the users to define a hook at the DAG level, and been executed before the pre_execution and after the post_execution.

TODO: we will propose detailed API change if such change is acceptable.

Use case / motivation

Airflow supports the operator to define their custom pre_execution/post_execution method, which will be exeucted before/after the execution of the operator, respectively.

But is hard to inject such hooks to existing operators, without modifying the code. i.e., airfow doesn't support a global pre-hook/post-hook to run before/after execution of every operators.

Such kind of hooks could be used for preparing the required inputs and finalizing the produced outputs for every operators in the DAG, etc. For example, the underlying XCom backends may support local values more efficiently than remote value accessing, the pre/post hook that the DAG level could do a migration/replication before the execution of operators, and gain end-to-end performance improvement.

Are you willing to submit a PR?

Yes, we would like to submit a PR if this feature request is acceptable.

Related Issues

None.

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template!

[uranusjr]

A DAG-level solution feels limiting to me. While it sounds convenient, it is not actually very often a user would want the same hook to be executed for every task, so you’d probably end up with a lot of if-else in the function to handle different tasks in the end. A better approach to me would be to allow supplying the hook without subclassing, something like

def pre_exec(task: BaseOperator, context: Any) -> None:
    ...

def different_pre_exec(task: BaseOperator, context: Any) -> None:
    ...

with DAG(...) as dag:
    t1 = BashOperator(
        ...,
        before_exec_hook=pre_exec,
    )
    t2 = BashOperator(
        ...,
        before_exec_hook=pre_exec,
    )
    t3 = BashOperator(
        ...,
        before_exec_hook=different_pre_exec,
    )

[sighingnow]

In our proposal the pre/post execution hooks only care about the linage information. Which commonly should be the same logic for every operator.

Supplying with an extra optional arguments for every operator does work, but it means we need to update at every place where we define an instance of some operator. Supplying the hooks at the DAG level makes it possible by just modifying the DAG(...), which would be less disturbing for end users.

[mik-laj]

@sighingnow What do you think about using default_args to define one common callback per DAG? Does it allow you to address your user case?

airflow/airflow/example_dags/example_bash_operator.py

Line 34 in ab91520

default_args=args,

[sighingnow]

Does it allow you to address your user case?

Thanks for the information. I think that would work, it the default_args is available before executing the task. I need to do experiment on it and will repost it here once I have some results.

[sighingnow]

After some experiment we think setting on_{execute,success,failure}_callback in default_args for a DAG works for us.

Thank you @mik-laj

[sighingnow]

Close as won't fix.