Сделать собственные макросы Airflow расширять другие макросы

Есть ли способ сделать пользовательский макрос в Airflow, который сам вычисляется из других макросов?

from airflow import DAG
from airflow.operators.bash_operator import BashOperator

dag = DAG(
    'simple',
    schedule_interval='0 21 * * *',
    user_defined_macros={
        'next_execution_date': '{{ dag.following_schedule(execution_date) }}',
    },
)

task = BashOperator(
    task_id='bash_op',
    bash_command='echo "{{ next_execution_date }}"',
    dag=dag,
)

Здесь используется обратный перенос нового макроса Airflow v1.8 next_execution_date для работы в Airflow v1.7. К сожалению, этот шаблон отображается без расширения макроса:

$ airflow render simple bash_op 2017-08-09 21:00:00
    # ----------------------------------------------------------
    # property: bash_command
    # ----------------------------------------------------------
    echo "{{ dag.following_schedule(execution_date) }}"

ОТВЕТЫ

Ответ 1

Вот несколько решений:

1. Переопределите BashOperator чтобы добавить некоторые значения в контекст

class NextExecutionDateAwareBashOperator(BashOperator):
    def render_template(self, attr, content, context):
        dag = context['dag']
        execution_date = context['execution_date']
        context['next_execution_date'] = dag.following_schedule(execution_date)

        return super().render_templates(attr, content, context)
        # or in python 2:
        # return super(NextExecutionDateAwareBashOperator, self).render_templates(attr, content, context)

Хорошая часть этого подхода: вы можете записать некоторый повторяющийся код в свой пользовательский оператор.

Плохая часть: вы должны написать пользовательский оператор для добавления значений в контекст, прежде чем будут отображены шаблонные поля.

2. Сделайте ваши вычисления в пользовательском макросе

Макросы не обязательно являются значениями. Они могут быть функциями.

В твоем даге

def compute_next_execution_date(dag, execution_date):
    return dag.following_schedule(execution_date)

dag = DAG(
    'simple',
    schedule_interval='0 21 * * *',
    user_defined_macros={
        'next_execution_date': compute_next_execution_date,
    },
)

task = BashOperator(
    task_id='bash_op',
    bash_command='echo "{{ next_execution_date(dag, execution_date) }}"',
    dag=dag,
)

Хорошая часть: вы можете определить многократно используемые функции для обработки значений, доступных во время выполнения (значения XCom, свойства экземпляра задания, свойства экземпляра задачи и т.д.), И сделать свой результат функции доступным для визуализации шаблона.

Плохая часть (но не такая раздражающая): вам нужно импортировать такую функцию, как пользовательский макрос, в каждый тег, где это необходимо.

3. Назовите ваше выражение прямо в вашем шаблоне

Это самое простое (как уже упоминалось в ответе Ардана) и, вероятно, хорошее в вашем случае.

BashOperator(
    task_id='bash_op',
    bash_command='echo "{{ dag.following_schedule(execution_date) }}"',
    dag=dag,
)

Идеально подходит для простых звонков, как этот. И это некоторые другие объекты, напрямую доступные как макросы (например, task, task_instance и т.д.); даже некоторые стандартные модули доступны (например, macros.time ,...).

Ответ 2

user_defined_macros по умолчанию не обрабатываются как шаблоны. Если вы хотите сохранить шаблон в user_defined_macro (или если вы используете шаблон в переменной params), вы всегда можете повторно запустить функцию templating вручную:

class DoubleTemplatedBashOperator(BashOperator):
    def pre_execute(self, context):
        context['ti'].render_templates()

И это будет работать для шаблонов, которые также не ссылаются на другие параметры или UDM. Таким образом, вы можете иметь "двухмерные" шаблоны.

Или поместите UDM непосредственно в команду BashOperator вместо этого (самое простое решение):

BashOperator(
    task_id='bash_op',
    bash_command='echo "{{ dag.following_schedule(execution_date) }}"',
    dag=dag,
)