Skip to content
Advertisement

Read run_id in airflow operator for non templated fields

I am trying to read run_id for DAG in SnowflakeOperator to set a session parameter, query_tag. But it seems like the session parameter is not templated.

snowflake_operator = SnowflakeOperator(
    task_id='snowflake_task',
    snowflake_conn_id='snowflake_conn',
    sql='resources/some.sql',
    warehouse='my_warehouse',
    database='my_database',
    role='my_role',
    session_parameters={
        'QUERY_TAG': 'dagrun_{{run_id}}'
    }
)

How can I reference run_id and use it as an input here?

Advertisement

Answer

You need to make the non-templated field templated.

class MySnowflakeOperator(SnowflakeOperator):

    template_fields = (
        "session_parameters",
    ) + SnowflakeOperator.template_fields

Then you can use it as:

snowflake_operator = MySnowflakeOperator(
    task_id='snowflake_task',
    snowflake_conn_id='snowflake_conn',
    sql='resources/some.sql',
    warehouse='my_warehouse',
    database='my_database',
    role='my_role',
    session_parameters={
        'QUERY_TAG': 'dagrun_{{run_id}}'
    }
)
8 People found this is helpful
Advertisement