如何在Airflow任务间传递CSV文件？代码拆分与XCom使用问题

问题分析 #

原代码是单任务完成数据拉取、CSV生成和S3上传，现在要拆分为两个独立任务：

1. 拉取Firebase用户数据并生成CSV内容
2. 将CSV内容上传至S3

你尝试用XCom传递但失败，核心原因大概率是传递了本地文件路径而非文件内容——Airflow任务可能运行在不同Worker节点，本地文件无法跨节点共享。由于文件体积小，直接传递CSV字符串是可行方案。

完整实现代码 #

import csv
import io
import firebase_admin
from firebase_admin import credentials, auth
import boto3
from datetime import datetime
import json
from airflow.models import Variable
from airflow.decorators import dag, task
from airflow.utils.dates import days_ago

# DAG基础配置
default_args = {
    'owner': 'airflow',
    'retries': 1,
}

@dag(
    default_args=default_args,
    schedule_interval=None,
    start_date=days_ago(1),
    catchup=False,
    tags=['firebase', 's3']
)
def firebase_users_to_s3():

    @task
    def fetch_firebase_users_and_generate_csv():
        # 加载Airflow变量
        aws_access_key_id = Variable.get('RIVERY_AWS_ACCESS_KEY_ID')
        aws_secret_access_key = Variable.get('RIVERY_AWS_SECRET_ACCESS_KEY')
        bucket = Variable.get('RIVERY_BUCKET')
        firebase_key = 'firebase/key/firebase.json'

        # 初始化S3会话并获取Firebase密钥
        session = boto3.Session(aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key)
        s3 = session.resource('s3')
        obj = s3.Object(bucket, firebase_key)
        file_content = obj.get()['Body'].read().decode('utf-8')
        json_data = json.loads(file_content)

        # 初始化Firebase（避免重复初始化）
        if not firebase_admin._apps:
            cred = credentials.Certificate(json_data)
            firebase_admin.initialize_app(cred)

        # 拉取所有Firebase用户
        all_users = auth.list_users()
        users_data = []
        for user in all_users.iterate_all():
            user_data = {
                'uid': user.uid,
                'email': user.email,
                'display_name': user.display_name,
                'creation_timestamp': user.user_metadata.creation_timestamp,
                'provider_data': [{
                    'provider_id': provider.provider_id,
                    'email': provider.email,
                    'phone_number': provider.phone_number,
                    'photo_url': provider.photo_url,
                    'uid': provider.uid,
                    'last_sign_in_timestamp': user.user_metadata.last_sign_in_timestamp
                } for provider in user.provider_data],
            }
            users_data.append(user_data)

        # 内存中生成CSV内容（不写本地文件）
        csv_fields = ['uid', 'email', 'display_name', 'creation_timestamp', 'provider_id', 'provider_email', 
                      'provider_phone_number', 'provider_photo_url', 'provider_uid', 'last_sign_in_timestamp']
        output = io.StringIO()
        writer = csv.DictWriter(output, fieldnames=csv_fields)
        writer.writeheader()

        for user in users_data:
            for provider_data in user['provider_data']:
                user_row = {
                    'uid': user['uid'],
                    'email': user['email'],
                    'display_name': user['display_name'],
                    'creation_timestamp': user['creation_timestamp'],
                    'provider_id': provider_data['provider_id'],
                    'provider_email': provider_data['email'],
                    'provider_phone_number': provider_data['phone_number'],
                    'provider_photo_url': provider_data['photo_url'],
                    'provider_uid': provider_data['uid'],
                    'last_sign_in_timestamp': provider_data['last_sign_in_timestamp'],
                }
                writer.writerow(user_row)
        
        # 返回CSV字符串，自动推送到XCom
        return output.getvalue()

    @task
    def upload_csv_to_s3(csv_content: str):
        # 加载Airflow变量
        aws_access_key_id = Variable.get('RIVERY_AWS_ACCESS_KEY_ID')
        aws_secret_access_key = Variable.get('RIVERY_AWS_SECRET_ACCESS_KEY')
        bucket = Variable.get('RIVERY_BUCKET')
        s3_key = 'firebase/auth/firebase_auth_from_airflow.csv'

        # 将CSV字符串转为字节流，直接上传S3
        s3_client = boto3.client('s3', aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key)
        s3_client.put_object(
            Bucket=bucket,
            Key=s3_key,
            Body=csv_content.encode('utf-8'),
            ContentType='text/csv'
        )
        print(f"文件已成功上传至 {bucket}/{s3_key}")

    # 定义任务依赖
    csv_content = fetch_firebase_users_and_generate_csv()
    upload_csv_to_s3(csv_content)

# 实例化DAG
firebase_users_to_s3_dag = firebase_users_to_s3()

关键改进点 #

1. XCom传递内容而非路径：第一个任务直接返回CSV字符串，Airflow自动将其存入XCom；第二个任务通过函数参数接收内容（Airflow自动从XCom拉取），无需手动调用xcom_push/xcom_pull。
2. 内存处理CSV：用io.StringIO在内存中生成和读取CSV，避免本地文件依赖，解决跨Worker节点的文件共享问题。
3. 修复原代码错误：
   • 修正了CSV文件名的格式化错误（原代码(today)未替换为实际日期，这里直接去掉本地文件，用内存生成）
   • 修复了bucket_name未定义的问题，统一使用bucket变量
   • 添加Firebase重复初始化检查，避免任务重试时抛出初始化错误
4. 任务解耦：两个任务职责单一，便于独立调试和维护。

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