如何在Airflow自定义Sensor中使用Jinja模板解析上下文变量

我正在构建一个自定义Sensor，目前通过context手动提取data_interval_start、data_interval_end等变量来拼接SQL：

def poke(self, context):
    data_interval_start= context['data_interval_start']
    data_interval_end= context['data_interval_end']
    query = f"""
    SELECT 
        count(*) as count 
    FROM `mytable`
    WHERE mydate >= "{data_interval_start}" and mydate < "{data_interval_end}"
    """
...

这种方式比较繁琐，每次需要新的上下文变量都得手动定义。我想实现类似MySQLToGCSOperator的效果——直接在SQL中使用Jinja模板（如下），无需额外代码就能自动替换变量：

SELECT 
    count(*) as count 
FROM `mytable`
WHERE mydate >= "{{data_interval_start}}" and mydate < "{{data_interval_end}}"

请问如何在自定义Sensor中实现这种自动模板渲染，不用手动用f-string拼接？

有两种简洁的实现方式，都能避免手动提取context变量：

方式一：手动调用Airflow模板渲染函数 #

直接使用Airflow内置的render_template工具，传入SQL模板和context即可完成渲染：

from airflow.utils.template import render_template
from airflow.sensors.base import BaseSensorOperator

class MyCustomSensor(BaseSensorOperator):
    def poke(self, context):
        # 定义带Jinja占位符的SQL模板
        sql_template = """
        SELECT 
            count(*) as count 
        FROM `mytable`
        WHERE mydate >= "{{data_interval_start}}" and mydate < "{{data_interval_end}}"
        """
        # 自动渲染模板，替换所有上下文变量
        rendered_query = render_template(sql_template, context)
        
        # 后续使用渲染后的query执行查询逻辑
        ...

方式二：利用Airflow的模板字段自动渲染（推荐） #

通过在自定义Sensor类中声明template_fields属性，让Airflow自动处理指定字段的模板渲染，用法和MySQLToGCSOperator完全一致：

1. 定义自定义Sensor类 #

from airflow.sensors.base import BaseSensorOperator

class MyCustomSensor(BaseSensorOperator):
    # 指定需要自动模板化的字段（这里是sql）
    template_fields = ("sql",)

    def __init__(self, sql, **kwargs):
        super().__init__(**kwargs)
        self.sql = sql

    def poke(self, context):
        # 此时self.sql已经被Airflow自动渲染完成，直接使用
        rendered_query = self.sql
        
        # 执行查询逻辑
        ...

2. 在DAG中使用Sensor #

my_sensor_task = MyCustomSensor(
    task_id="check_mytable_count",
    sql="""
        SELECT 
            count(*) as count 
        FROM `mytable`
        WHERE mydate >= "{{data_interval_start}}" and mydate < "{{data_interval_end}}"
        """,
    dag=dag
)

这种方式下，Airflow会在任务执行前自动渲染sql字段中的Jinja模板，你无需在poke方法中做任何额外处理。

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