Allow context usage in user_defined_macros

[vanchaxy]

Description

Allow context usage in user_defined_macros. For this move set_current_context context manager from wrapping _execute_task to wrapping _execute_task_with_callbacks.

Use case/motivation

After this change, we will be able to use get_current_context while rendering template fields.

Example use case. We have dag with the first task that prepares configuration for this run and after this, all other task uses this config to render templates. We have no control over other task codes, as it's providers operators etc.

After this change we can define class like:

class XcomResult:
    def __init__(self, task_id):
        self.task_id = task_id
        self._result = None

    def __get__(self, item):
        if self._result is None:
            context = get_current_context()
            self._result = context['ti'].xcom_pull(task_ids=self.task_id)
        self._result.__get__(item)

user_defined_macros = {"config": XcomResult("prepare_config_task_id")}

Then we can use this in all templates as config.needed_key.

For now, if we want something similar we need to define function like:

def built_get_xcom_result(task_id):
    def get_config(ti):
        return ti.xcom_pull(task_ids=task_id)
    return get_config 

user_defined_macros = {"get_config": built_get_xcom_result("prepare_config_task_id")}

and use it in all templates as get_config(ti).needed_key

Related issues

No response

Are you willing to submit a PR?

• [ ] Yes I am willing to submit a PR!

Code of Conduct

• [X] I agree to follow this project's Code of Conduct

[boring-cyborg[bot]]

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

[potiuk]

You seem to know what you want to do - would you like to propose a PR for that? Discussing on actual PR would be much more productive and you could become one of the > 2K contributors to airflow this way.

[vanchaxy]

I will try to propose a PR. Just need some time to configure env and read guidelines.

[romsharon98]

This issue not reproduce in main branch. Therefore closing this issue.