How to extend the template context in Airflow?

Viewed 1659

How can I add new keys to the context that is provided for templates (i.e. Macro)? I have looked in the source code for solution, but could not find out myself.

I have to adapt to external constraints in naming the template parameters (start_date and end_date). user_defined_macros and user_defined_filters and plugin macros are not applicable, because these will provide functions, not the evaluated values to be substituted (i.e {{start_date()}} instead of {{start_date}} must be written in the template).
I have tried to use Jinja's contextfunction decorator, but context can't be modified.

The get_template_context() method of TaskInstance class returns the dictionary in models/taskinstance.py, which becomes the context for templating. It seems that, there's no way to extend (update()) this dictionary by other means other than patching the source of Airflow, which I would like to avoid.

1 Answers

you can use out of the box templating for sql files. For example you can use:

select * from table1 where created between "{{params.start_date}}" and "{{params.end_date}}"

and give the parameters in your task definition:

task = PostgresOperator(task_id='task_name',
                               sql="sql/delete_and_copy.sql",
                               postgres_conn_id="postgres_connection",
                               database="dev",
                               params={
                                          'start_date': start_date,
                                          'end_date': end_date,
                                      },
                                      )

start_date and end_date can be calculated by a Python operator. https://airflow.apache.org/docs/apache-airflow/stable/tutorial.html#templating-with-jinja

Related