如何在Airflow的SnowflakeOperator非模板字段中读取run_id

由于SnowflakeOperator的session_parameters字段默认不支持模板渲染，你可以通过以下几种方式实现需求：

方法1：在SQL脚本中设置QUERY_TAG #

直接在执行的SQL文件开头添加会话参数设置语句，利用SQL字段的模板化特性替换run_id：

ALTER SESSION SET QUERY_TAG = 'dagrun_{{run_id}}';

-- 原有业务SQL逻辑
SELECT * FROM your_table;

这种方式无需修改Operator代码，简单直接，执行SQL时Airflow会自动替换模板变量。

方法2：用PythonOperator动态构造Snowflake任务 #

通过PythonOperator获取DAG运行上下文里的run_id，再动态创建并执行SnowflakeOperator：

from airflow.operators.python import PythonOperator
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator

def execute_snowflake_task(**context):
    # 从上下文获取run_id
    run_id = context["run_id"]
    # 构造Snowflake任务
    snowflake_task = SnowflakeOperator(
        task_id="dynamic_snowflake_task",
        snowflake_conn_id="snowflake_conn",
        sql="resources/some.sql",
        warehouse="my_warehouse",
        database="my_database",
        role="my_role",
        session_parameters={"QUERY_TAG": f"dagrun_{run_id}"}
    )
    # 执行任务
    snowflake_task.execute(context)

# 在DAG中添加PythonOperator
setup_snowflake_task = PythonOperator(
    task_id="setup_snowflake_task",
    python_callable=execute_snowflake_task,
    provide_context=True
)

方法3：自定义支持模板化的SnowflakeOperator #

继承原SnowflakeOperator，将session_parameters加入模板字段列表，使其支持变量渲染：

from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator

class TemplatedSnowflakeOperator(SnowflakeOperator):
    # 扩展模板字段，添加session_parameters
    template_fields = SnowflakeOperator.template_fields + ("session_parameters",)

# 使用自定义Operator
snowflake_operator = TemplatedSnowflakeOperator(
    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}}"}
)

内容的提问来源于stack exchange，提问作者Sasan Ahmadi