Airflow

Use a custom function in Airflow templates

How to add a custom function to Airflow and use it in a template

Bartosz Mikulski 08 Nov 2020 – 1 min read

Adding a custom function to Airflow is quite simple. First, we have do define a function in Python, for example, this one:

def do_something_with_execution_date(execution_date):
    # Imagine that there is some useful code ;)
    ...

When the function is ready, we use the user_defined_macros parameter of the DAG object to pass a dictionary of custom functions:

dag = DAG(
    ...,
    user_defined_macros={
        'custom_function': do_something_with_execution_date,
    }
)

Now, we can use the custom function in any place that supports Airflow templates. Of course, only in the DAGs that have access to the functions.

{{ custom_function(execution_date) }};

Note that, I can pass parameters to the function and rename it by using a different name as the dictionary key.

Older post

Speed up counting the distinct elements in a Spark DataFrame