Re: [D] Using nested (decorated) function in decorated tasks (plus integration with TaskFlow) [airflow]

["pykenny (via GitHub)" <gi...@apache.org>]

GitHub user pykenny added a comment to the discussion: Using nested (decorated) function in decorated tasks (plus integration with TaskFlow)

After some thoughts I figured out the idea comes needs on simplifying task functions:

1. Replace routine code that often runs at the start or at the end of pipeline tasks.
2. Similar to 1., plus further inject resource instantiated during the routine(s) into task function, such as logger, Airflow hooks, client from 3-rd party library or even customized package.

Simplifying utility implementation is common in Python programming using native decorator function (closure/inner function), `functools` helpers, or solutions such as [wrapt](https://github.com/GrahamDumpleton/wrapt/tree/develop), and I believe that's the idea behind defining DAG/Task/Task Group/TaskFlow with decorated function emerges. 

Certainly, these kind of work can be simplified with well designed class/function/context manager that meets the need, but I was thinking of the possibility to further integrate these utilities into Airflow's operator/decorator system, such as replacing the following code definition:

```python
from airflow.decorators import dag, task
from my_airflow_provider.decorators import 

@dag()
def my_dag():
    @task():
    def my_task_01():
        from my_task_package.commons import setup, cleanup

        setup()
        # ...
        cleanup()

    @task():
    def my_task_02():
        from my_task_package.commons import setup, cleanup

        setup()
        # ...
        cleanup()

    @task():
    def my_task_03():
        from my_task_package.commons import setup, cleanup

        setup()
        # ...
        cleanup()

    # ... and more tasks ...

dag = my_dag()
```

with:

```python
from airflow.decorators import dag, task
from my_airflow_provider.decorators.commons import setup_cleanup

@dag()
def my_dag():
    @task():
    @setup_cleanup()
    def my_task_01():
        ...

    @task():
    @setup_cleanup()
    def my_task_02():
        ....

    @task():
    @setup_cleanup()
    def my_task_03():
        ...

dag = my_dag()
```

Simply writing wrapper function that packs another one, and then hand it to DAG parser for sure is not a good idea, since Airflow has few (if any) knowledge about what's inside of the function, and it's likely to hit on many edge cases when nested closure gets too complex. Basically, Airflow decorator functions receive a Python function, or even an existing decorated task function (for instance, `@setup` and `@teardown` works on both native Python functions and decorator instances already wrapped by `@task`), wrap it into a decorator instance, and finally instantiate task instances when it gets called. 

Ideally we can work around with the mechanism and do what we want above on the wrapped instances. However, that requires the task to run in the same environment as its Airflow runner's, and in several cases it isn't, like running tasks under virtual environments (`VirtualenvPythonOperator`) or external environments (`ExternalPythonOperator`).

GitHub link: https://github.com/apache/airflow/discussions/45963#discussioncomment-12080962

----
This is an automatically sent email for commits@airflow.apache.org.
To unsubscribe, please send an email to: commits-unsubscribe@airflow.apache.org