Airflow中Task2失败时，如何回滚Task1的Snowflake SQL操作？

可以实现，以下是两种可行的方案：

方案1：利用Snowflake事务机制（单任务封装逻辑） #

由于Airflow的独立任务会创建新的Snowflake连接会话，跨任务共享事务会话难度较高，推荐将task1和task2的逻辑封装在同一个任务中，通过Python代码控制事务的提交与回滚：

from airflow.decorators import task
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook

@task
def task_combined_with_transaction():
    hook = SnowflakeHook(snowflake_conn_id="your_snowflake_conn_id")
    conn = hook.get_conn()
    cursor = conn.cursor()
    try:
        # 执行task1的业务SQL
        cursor.execute("BEGIN TRANSACTION;")
        cursor.execute("INSERT INTO target_table VALUES (1, 'sample_data');")
        
        # 执行task2的业务逻辑（示例为验证查询，可替换为你的实际逻辑）
        cursor.execute("SELECT COUNT(*) FROM target_table WHERE id = 1;")
        count_result = cursor.fetchone()[0]
        if count_result != 1:
            raise Exception("Task2 validation failed")
        
        # 所有步骤成功，提交事务
        cursor.execute("COMMIT;")
    except Exception as e:
        # 任意步骤失败，回滚事务
        cursor.execute("ROLLBACK;")
        # 抛出异常标记任务失败
        raise e
    finally:
        cursor.close()
        conn.close()

这种方案严格遵循事务ACID特性，确保task1和task2的操作要么全部成功提交，要么全部回滚。

方案2：反向操作回滚（多任务拆分场景） #

如果必须保持task1和task2为独立任务，可针对task1的操作编写反向回滚逻辑，当task2失败时触发执行：

1. task1：执行原始业务SQL（自动提交）

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

task1 = SnowflakeOperator(
    task_id="task1_execute_sql",
    sql="INSERT INTO target_table VALUES (1, 'sample_data');",
    snowflake_conn_id="your_snowflake_conn_id",
    autocommit=True
)

2. task2：你的业务验证/处理任务

task2 = SnowflakeOperator(
    task_id="task2_business_logic",
    sql="SELECT * FROM target_table WHERE id = 1;",
    snowflake_conn_id="your_snowflake_conn_id"
)

3. rollback_task：执行task1的反向操作（如删除插入的数据）

rollback_task = SnowflakeOperator(
    task_id="rollback_task1",
    sql="DELETE FROM target_table WHERE id = 1;",
    snowflake_conn_id="your_snowflake_conn_id",
    trigger_rule="all_failed"  # 仅当task2失败时执行
)

4. 设置依赖关系：

task1 >> task2
task2 >> rollback_task

注意事项 #

• 方案2需要确保task1的操作存在可执行的反向逻辑，且反向操作不会影响其他业务数据；
• 若task1是复杂的多步操作（如批量更新、关联操作），反向逻辑的编写会更复杂，此时方案1的事务封装更可靠。

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