如何通过另一个DAG将Airflow中运行的DAG标记为失败？

我尝试创建一个可接收dag_id和execution_date作为参数的DAG，目的是将指定的DAG Run标记为失败，但始终无法实现目标。以下是我的代码，求有类似实践经验的用户帮忙排查问题：

from airflow import DAG, utils
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.python_operator import PythonOperator
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.utils.state import State

from airflow.models import DagRun
from datetime import datetime, timezone
import logging
logger = logging.getLogger(__file__)

default_args = {
    'owner': 'test',
    'start_date': utils.dates.days_ago(1),
    'retries': 2,
}

def mark_execution_as_failure(dag_id_param, execution_date_param):

    execution_param = datetime.strptime(execution_date_param, '%Y-%m-%dT%H:%M:%S').replace(tzinfo=timezone.utc)

    dagrun = DagRun.find(dag_id=dag_id_param, execution_date=execution_param)
    logger.info(dagrun)
    if dagrun:
        logger.info("DagRUN")
        logger.info(dagrun[0])
        dagrun[0].set_state(State.FAILED)
    else:
        logger.info(f"DAG run for dag_id {dag_id_param} and execution_date {execution_param} not found.")

with DAG(
    'mark_execution_failure',
    default_args=default_args,
    description='Mark an execution as failure',
    schedule_interval=None,
    catchup=False,
) as dag:

    start_operator = DummyOperator(task_id='start')

    dag_id_param = '{{ dag_run.conf.dag_id }}'
    execution_date_param = '{{ dag_run.conf.execution_date }}'

    logger.info(dag_id_param)
    logger.info(execution_date_param)

    task_kill_dag = PythonOperator(
        task_id="kill_dag",
        python_callable=mark_execution_as_failure,
        op_args=[dag_id_param, execution_date_param],
        provide_context=True,
    )

    start_operator >> task_kill_dag

问题排查与解决方案 #

你的代码主要存在两个关键问题，导致无法成功标记DAG Run为失败：

1. 数据库会话未提交：调用set_state()后，没有将变更提交到Airflow的元数据库，修改只存在于内存中，不会持久化。
2. 日期格式兼容性差：Airflow存储的execution_date可能包含毫秒（比如2024-05-20T10:00:00.000000+00:00），而你的代码只支持不带毫秒的格式，会导致解析失败或查找不到目标DAG Run。

以下是修复后的代码：

from airflow import DAG, utils
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.python_operator import PythonOperator
from airflow.utils.state import State
from airflow.models import DagRun
from datetime import datetime, timezone
import logging
from airflow.utils.session import create_session  # 导入会话管理工具

logger = logging.getLogger(__file__)

default_args = {
    'owner': 'test',
    'start_date': utils.dates.days_ago(1),
    'retries': 2,
}

def mark_execution_as_failure(dag_id_param, execution_date_param):
    # 兼容带毫秒和不带毫秒的日期格式
    execution_param = None
    for date_format in ['%Y-%m-%dT%H:%M:%S.%f', '%Y-%m-%dT%H:%M:%S']:
        try:
            execution_param = datetime.strptime(execution_date_param, date_format).replace(tzinfo=timezone.utc)
            break
        except ValueError:
            continue
    
    if not execution_param:
        logger.error(f"无法解析execution_date参数：{execution_date_param}")
        return

    # 使用会话操作数据库，确保修改被持久化
    with create_session() as session:
        # 直接通过查询获取目标DAG Run
        dagrun = session.query(DagRun).filter(
            DagRun.dag_id == dag_id_param,
            DagRun.execution_date == execution_param
        ).first()
        
        logger.info(f"查询到的DAG Run: {dagrun}")
        if dagrun:
            logger.info(f"将DAG Run [{dagrun.dag_id}] [{dagrun.execution_date}] 标记为FAILED状态")
            # 传入session参数并提交变更
            dagrun.set_state(State.FAILED, session=session)
            session.commit()
        else:
            logger.info(f"未找到dag_id={dag_id_param}且execution_date={execution_param}的DAG Run")

with DAG(
    'mark_execution_failure',
    default_args=default_args,
    description='标记指定DAG Run为失败',
    schedule_interval=None,
    catchup=False,
) as dag:
    start_operator = DummyOperator(task_id='start')

    task_kill_dag = PythonOperator(
        task_id="kill_dag",
        python_callable=mark_execution_as_failure,
        op_args=['{{ dag_run.conf.dag_id }}', '{{ dag_run.conf.execution_date }}'],
        provide_context=True,
    )

    start_operator >> task_kill_dag

额外注意事项 #

• 触发该DAG时，必须在conf中传入正确的参数，示例JSON如下：

  {"dag_id": "你的目标DAG ID", "execution_date": "2024-05-20T10:00:00"}
  

• 确保执行该DAG的账号拥有修改目标DAG状态的权限（RBAC模式下需要赋予DAG Runs的编辑权限）。
• 查看Airflow任务日志，确认是否存在参数解析错误或数据库操作异常的提示，便于进一步排查。

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