如何处理Airflow中SnowflakeCheckOperator的返回值以实现正确校验

我正在实验用Airflow监控Snowflake任务，想实现一个单任务DAG：向Snowflake推送SQL查询，校验返回数值是否大于设定阈值。

我的SQL语句：

query_check = """select COUNT(*)
FROM (select CASE WHEN NAME LIKE '%SW_PRODUCTFEED%' THEN 'PRODUCTFEED'
        ELSE NULL END AS TASKTREE_NAME
       , NAME 
       , STATE
       , ERROR_MESSAGE
       , SCHEDULED_TIME
       , QUERY_START_TIME
       , NEXT_SCHEDULED_TIME
  from table(TEST_DB.INFORMATION_SCHEMA.task_history())
  where TASKTREE_NAME IS NOT NULL
  qualify DENSE_RANK() OVER (PARTITION BY TASKTREE_NAME ORDER BY to_date(SCHEDULED_TIME) desc) < 3
  order by scheduled_time desc);"""

DAG及任务定义：

with dag:
    query1_exec = SnowflakeCheckOperator(
        task_id="snowflake_check_task_history",
        sql=query_check,
        params={
            "check_name": "number_rows",
            "check_statement": "count >=1"
        },
        conn_id="Snowflake_test"
    )

query1_exec

目前遇到的问题：Snowflake返回的结果是类似Record: {'COUNT(*)': 10}的字典对象，SnowflakeCheckOperator没有针对字典里的数值做校验，导致校验始终为真，想知道怎么处理返回值实现正确校验——比如修改返回格式，或者直接获取字典中的对应值。

方法1：给COUNT(*)设置别名，调整校验语句 #

SnowflakeCheckOperator的校验逻辑是基于查询结果的列名匹配的，所以先给COUNT(*)加一个明确的别名，再调整校验语句的对应列名：

1. 修改SQL语句，给COUNT(*)添加别名：

query_check = """select COUNT(*) as count_rows
FROM (select CASE WHEN NAME LIKE '%SW_PRODUCTFEED%' THEN 'PRODUCTFEED'
        ELSE NULL END AS TASKTREE_NAME
       , NAME 
       , STATE
       , ERROR_MESSAGE
       , SCHEDULED_TIME
       , QUERY_START_TIME
       , NEXT_SCHEDULED_TIME
  from table(TEST_DB.INFORMATION_SCHEMA.task_history())
  where TASKTREE_NAME IS NOT NULL
  qualify DENSE_RANK() OVER (PARTITION BY TASKTREE_NAME ORDER BY to_date(SCHEDULED_TIME) desc) < 3
  order by scheduled_time desc);"""

2. 调整SnowflakeCheckOperator的params参数：

params={
    "check_name": "number_rows",
    "check_statement": "count_rows >= 1"
}

这样Operator就能正确识别列名，提取对应数值完成校验。

方法2：自定义结果处理（适合复杂校验逻辑） #

如果需要更灵活的校验规则，可以拆分任务：先用SnowflakeOperator执行查询，再用PythonOperator处理结果并校验：

with dag:
    # 执行查询并将结果推送到XCom
    run_query = SnowflakeOperator(
        task_id="run_snowflake_query",
        sql=query_check,
        conn_id="Snowflake_test",
        do_xcom_push=True
    )

    # 自定义校验函数
    def validate_count(**context):
        # 从XCom拉取查询结果
        result = context['ti'].xcom_pull(task_ids='run_snowflake_query')
        # 提取字典中的COUNT(*)数值
        count = result[0]['COUNT(*)']
        if count < 1:
            raise ValueError(f"校验失败：返回行数{count}小于阈值1")

    # 执行校验任务
    validate_task = PythonOperator(
        task_id="validate_query_result",
        python_callable=validate_count,
        provide_context=True
    )

    # 设置任务依赖
    run_query >> validate_task

这种方式可以完全自定义校验逻辑，适配各种复杂场景。

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