use Airflow connection from a jinja template

Viewed 3723

I'm trying to pass DB params to BashOperator using environment variables, but I can't find any documentation/examples how to use a connection from a Jinja template.

So I'm looking for something similar to variables

echo {{ var.value.<variable_name> }}
2 Answers

For Airflow >= 2.2.0:

Assuming you have conn id test_conn you can use macros directly via:

{{ conn.test_conn }} so you get any connection attribute like:

{{ conn.test_conn.host }}, {{ conn.test_conn.login }}, {{ conn.test_conn.password }} and so on.

For Airflow < 2.2.0:

There is no ready to use macro however you can create custom macros to address this.

Connection example:

enter image description here

Creating the macros:

def get_host(conn_id):
    connection = BaseHook.get_connection(conn_id)
    return connection.host

def get_schema(conn_id):
    connection = BaseHook.get_connection(conn_id)
    return connection.schema

def get_login(conn_id):
    connection = BaseHook.get_connection(conn_id)
    return connection.login

Using them in a DAG:

def print_function(**context):
    print(f"host={context['host']} schema={context['schema']} login={context['login']}")

user_macros = {
    'get_host': get_host,
    'get_schema': get_schema,
    'get_login': get_login,
}

with DAG(
    dag_id='connection',
    default_args=default_args,
    schedule_interval=None,
    user_defined_macros=user_macros,
) as dag:

# Example how to use as function
python_op = PythonOperator( 
    task_id='python_task',
    provide_context=True,
    python_callable=print_function,
    op_kwargs={
        'host': get_host("test_conn"),
        'schema': get_schema("test_conn"),
        'login': get_login("test_conn"),
    }
)

# Example how to use as Jinja string
bash_op = BashOperator( 
    task_id='bash_task',
    bash_command='echo {{ get_host("test_conn") }} {{ get_schema("test_conn") }} {{ get_login("test_conn") }} ',
)

Rendering for PythonOperator example: enter image description here

enter image description here

Rendering for BashOperator example:

enter image description here

General Explnation: What this code does is creating a custom function func() to be used as user_defined_macros thus providing the ability to use it just like this macro was defined by Airflow itself. You can access the templating as: {{ func() }} as seen in the example the function allow accept parameters.

Note you can create such functions for all fields in the connection object.

be cautious with how you use it, passing passwords as text may not be a good idea.

My PR added the {{ conn.my_conn_id.login }} syntax and it will be available in airflow 2.2.0 (not released yet as of 2021-09-22).

See the unreleased documentation for templates reference here

For 2.1.4 and earlier:

Improving on previous answers,

Define macro per DAG: {{conn.<conn_id>}}

you can get conn.<connection_name>.host syntax by using the following macro:

class ConnectionGrabber:
    def __getattr__(self, name):
        return  Connection.get_connection_from_secrets(name)
dag = DAG(user_defined_macros={'connection': ConnectionGrabber()}

that injects the name connection into the jinja template context, connection is a ConnectionGrabber instance. This ConnectionGrabber provides dynamic/managed attributes, so when you request attribute my_conn_id (like connection.my_conn_id) it will perform a lookup of using the airflow.models.Connection.get_connection_from_secrets and return that, from there you can use the a.m.Connection attributes like host, login, password, etc.

Complete example where the connection mssql is accessed in jinja template with bash_command='echo {{connection.mssql.host }}':

from airflow.models import DAG,Connection
from airflow.operators.bash import BashOperator
from airflow.utils.dates import days_ago


class ConnectionGrabber:
    def __getattr__(self, name):
        return  Connection.get_connection_from_secrets(name)


args = {'owner': 'airflow', 'retries': 3, 'start_date': days_ago(2)}

dag = DAG(
    dag_id='test_connection',
    default_args=args,
    schedule_interval='0 0 * * *',
    dagrun_timeout=timedelta(minutes=60),
    user_defined_macros={'connection': ConnectionGrabber()}
)

task = BashOperator(task_id='read_connection', bash_command='echo {{connection.mssql.host }}', dag=dag)

Define macro in a plugin: {{macros.conn.value.<conn_id>}}

If you wanted to use this macro in all your DAGs you can wrap it in a plugin like so:

# $AIRFLOW_HOME/plugins/connection_macro.py
from airflow.plugins_manager import AirflowPlugin
from airflow.models import Connection


class ConnectionGrabber:
    __name__ = "value"
    def __str__(self):
        return self.__name__
    def __getattr__(self, name):
        return  Connection.get_connection_from_secrets(name)

class MacrosPlugin(AirflowPlugin):
    name = "conn"
    macros = [ConnectionGrabber()]

check that the plugin can be loaded with airflow plugins

airflow plugins
name | source                    | macros
=====+===========================+=======
conn | $PLUGINS_FOLDER/macros.py | value

Then you can use {{ macros.conn.value.<conn_id>.host }} in your jinja templates like so

   task = BashOperator(task_id='read_connection', bash_command='echo macros.conn.value.mssql.host = {{macros.conn.value.mssql.host }}', dag=dag)

I also opened an issue to add the conn.value.<conn_id> syntax natively

Related