通过代码为Airflow任务实例加备注遇多行结果错误求解析

问题背景 #

尝试通过代码给Airflow的Task Instance设置备注，运行代码时触发错误：

sqlalchemy.exc.MultipleResultsFound: Multiple rows were found when exactly one was required

疑惑为何查询会返回多行结果。

错误原因 #

1. 查询对象完全错误：你要修改的是Task Instance的备注，却查询了DagRun对象。DagRun是DAG的运行实例，一个DagRun对应多个TaskInstance（DAG里的每个任务都会生成一个TaskInstance）。当你用TaskInstance的字段过滤DagRun时，SQLAlchemy会自动关联两张表，返回所有关联符合条件TaskInstance的DagRun——如果同一个DagRun下有多个TaskInstance匹配条件，就会返回多行重复的DagRun记录，触发报错。
2. 表关联导致重复结果：session.query(DagRun).filter(TaskInstance.dag_id == ...)这种写法会让SQLAlchemy关联DagRun和TaskInstance表，只要TaskInstance满足过滤条件，对应的DagRun就会被返回。哪怕是同一个DagRun，只要关联了多个符合条件的TaskInstance，就会出现多行结果。

修正方案 #

方案一：直接从Context获取TaskInstance（最推荐） #

Airflow的任务上下文里已经自带当前的task_instance对象，不需要手动查询数据库，直接修改即可：

def ti_note():
    @task
    def set_my_note(**context):
        ti = context['task_instance']
        message = f"Note for dag {ti.dag_id}, task {ti.task_id}, and execution date {ti.execution_date}"
        ti.note = message
        # Airflow会自动处理对象的持久化，无需手动操作Session
    set_my_note_task = set_my_note()

方案二：正确查询TaskInstance对象 #

如果必须手动操作数据库，要确保查询的是TaskInstance而非DagRun，同时用更安全的查询方法：

from airflow.models import TaskInstance, settings

def ti_note():
    @task
    def set_my_note(**context):
        dag_id = context['dag'].dag_id
        task_id = context['task'].task_id
        execution_date = context['execution_date']

        message = f"Note for dag {dag_id}, task {task_id}, and execution date {execution_date}"

        session = settings.Session()
        # 明确查询TaskInstance对象，而非DagRun
        task_instance = session.query(TaskInstance).filter(
            TaskInstance.dag_id == dag_id,
            TaskInstance.task_id == task_id,
            TaskInstance.execution_date == execution_date
        ).one_or_none()  # 用one_or_none避免无匹配时触发异常

        if task_instance:
            task_instance.note = message
            session.commit()
        session.close()

    set_my_note_task = set_my_note()

补充说明 #

• 用one_or_none()替代one()能避免因数据异常（比如重复的TaskInstance记录）导致的报错，更健壮。
• Airflow官方推荐优先使用上下文提供的对象，减少手动数据库操作，避免Session管理不当引发的问题。

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