Apache Airflow多表场景下XCom传递列参数至Snowflake脚本的问题

问题核心 #

批量将S3数据同步到Snowflake多表时，没法通过XCom把每个表对应的列列表传给SQL脚本——试了传task ID没效果，硬编码task ID又触发重复ID报错。

修正后的Python代码 #

from airflow.decorators import dag, task
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
from airflow.utils.task_group import TaskGroup
from datetime import datetime

tableList = ["table_a","table_b"]

@dag(schedule_interval=None, start_date=datetime(2024,1,1), catchup=False)
def test_dag():

    @task
    def get_value(table: str):
        # 替换为实际获取表列名的逻辑，比如查询Snowflake的information_schema
        # 示例返回模拟列列表
        cols = ["col1", "col2", "col3"]
        return {
            "id_value": f"ID: {table}",
            "date": datetime.today(),
            "columns": ", ".join(cols)  # 直接拼成SQL能用的列字符串
        }

    for tab in tableList:
        tg_id = f"tg_{tab}"
        with TaskGroup(group_id=tg_id) as tg:
            # 给get_value传入当前表名
            val_task = get_value(tab)
            # 配置Snowflake任务，通过XCom拉取对应表的列数据
            insertTable = SnowflakeOperator(
                task_id=f"copy_snowflake_{tab}",
                snowflake_conn_id="conn",
                sql="insertTableTemplate.sql",
                params={
                    "table": tab,
                    "format": "your_file_format",  # 替换成你的实际文件格式
                    "pattern": "your_file_pattern",  # 替换成你的实际文件匹配规则
                    # 指定TaskGroup ID拉取对应组内get_value的返回值，避免跨组冲突
                    "columns": "{{ task_instance.xcom_pull(task_ids='get_value', key='return_value', task_group_id='" + tg_id + "')['columns'] }}"
                }
            )
            val_task >> insertTable

dag = test_dag()

关键修复点 #

• 修正命名错误：原代码里tg_id未定义、ids参数传错，现在按表名生成唯一的TaskGroup和任务ID，避免重复
• 简化列处理：在get_value里直接把列列表拼成col1, col2格式的字符串，SQL模板直接用就行
• 精准拉取XCom：通过task_group_id指定拉取当前TaskGroup内get_value的返回值，不会搞混不同表的列数据

调整后的SQL脚本（insertTableTemplate.sql） #

copy into {{ params.table }}
from (
    select {{ params.columns }}
    from @stage.path 
    (file_format => '{{ params.format }}', pattern => '{{ params.pattern }}')
);

注：原SQL语法有误，copy into的正确写法是copy into <表名> from <子查询>，已修正

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