Airflow2.x中PostgresOperator操作Redshift时如何传递查询结果到下一个任务

你可以通过XCom传递查询结果，无需使用临时表，有两种常用实现方式： #

方案1：使用高版本PostgresOperator内置XCom推送能力 #

apache-airflow-providers-postgres >= 4.1.0 版本的PostgresOperator已原生支持将查询结果推送到XCom，只需新增do_xcom_push=True参数即可：

第一个任务配置 #

retrieve_latest_query_task = PostgresOperator(
    sql='rs_warm-up_query-id.sql',
    postgres_conn_id='redshift',
    task_id='retrieve_latest_query_ids_from_metadata',
    do_xcom_push=True
)

第二个任务配置 #

PostgresOperator的sql参数支持Jinja模板渲染，直接通过模板语法拉取XCom中存储的查询ID列表即可，注意如果第一个查询返回多个ID，需要把WHERE query = {}调整为IN条件：

get_query_text_task = PostgresOperator(
    sql="""
    SELECT LISTAGG(CASE WHEN LEN (RTRIM(TEXT)) = 0 THEN TEXT ELSE RTRIM(TEXT) END,'') within group(ORDER BY SEQUENCE) AS TEXT
    FROM stl_querytext
    WHERE query IN (
        {% for query_item in ti.xcom_pull(task_ids='retrieve_latest_query_ids_from_metadata') %}
            {{ query_item[0] }}{% if not loop.last %},{% endif %}
        {% endfor %}
    );
    """,
    postgres_conn_id='redshift',
    task_id='get_full_query_text'
)

方案2：通过PythonOperator + PostgresHook自定义拉取（兼容低版本） #

如果你的Postgres Provider版本较低没有do_xcom_push参数，可以自己通过Hook拉取结果后推送到XCom：

1. 定义查询拉取函数 #

from airflow.providers.postgres.hooks.postgres import PostgresHook

def fetch_query_ids(**context):
    rs_hook = PostgresHook(postgres_conn_id='redshift')
    # 读取SQL文件内容
    with open('/path/to/rs_warm-up_query-id.sql', 'r', encoding='utf-8') as f:
        query_sql = f.read()
    # 执行查询并提取query ID列表
    query_result = rs_hook.get_records(query_sql)
    query_ids = [row[0] for row in query_result]
    # 返回值会自动推送到XCom
    return query_ids

2. 第一个任务替换为PythonOperator #

from airflow.operators.python import PythonOperator

retrieve_latest_query_task = PythonOperator(
    task_id='retrieve_latest_query_ids_from_metadata',
    python_callable=fetch_query_ids,
    provide_context=True
)

3. 第二个任务简化拉取逻辑 #

因为XCom中已经存储了纯ID列表，可以直接用Jinja过滤器格式化：

get_query_text_task = PostgresOperator(
    sql="""
    SELECT LISTAGG(CASE WHEN LEN (RTRIM(TEXT)) = 0 THEN TEXT ELSE RTRIM(TEXT) END,'') within group(ORDER BY SEQUENCE) AS TEXT
    FROM stl_querytext
    WHERE query IN ({{ ti.xcom_pull(task_ids='retrieve_latest_query_ids_from_metadata') | join(',') }});
    """,
    postgres_conn_id='redshift',
    task_id='get_full_query_text'
)

注意事项 #

• Airflow默认XCom单条数据最大为48KB，你的查询结果量级很小，完全满足使用需求，无需额外调整配置。
• 如果需要对每个query ID单独执行第二个查询，可以搭配Airflow动态任务映射实现批量执行。

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