Airflow通用数据暂存Operator及非XCom大数据处理方案问询

实现1：封装通用的带数据自动迁入迁出的基类Operator #

你可以自行封装一个继承自BaseOperator的通用基类，把数据迁入、迁出逻辑放到pre_execute和post_execute钩子中，所有需要暂存数据的业务Operator直接继承这个基类即可，不需要每个Operator单独开发拷贝逻辑：

• 迁入逻辑：在pre_execute阶段，根据传入的源地址参数（比如S3/OSS/HDFS路径），自动把对应数据拉取到当前工作节点的动态临时目录（可以用tempfile.mkdtemp()生成，避免固定路径冲突，任务执行完自动清理）
• 核心执行：直接用子类的execute方法实现你的业务逻辑，比如DockerOperator的镜像启动、脚本执行等，只需要把临时目录挂载/传入到业务执行环境即可
• 迁出逻辑：在post_execute阶段，自动把临时目录下指定的输出文件上传到目标外部存储路径，异常场景下跳过迁出直接报错

简化的代码示例如下：

from airflow.models.baseoperator import BaseOperator
from airflow.utils.decorators import apply_defaults
import tempfile
import os
# 自行实现的存储工具类，封装S3/OSS/HDFS的上传下载通用逻辑
from utils import storage_download, storage_upload

class AutoStageDataOperator(BaseOperator):
    @apply_defaults
    def __init__(
        self,
        source_data_path: str,
        output_data_path: str,
        *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.source_data_path = source_data_path
        self.output_data_path = output_data_path
        self.temp_dir = None

    def pre_execute(self, context):
        # 生成唯一临时目录，避免多任务冲突
        self.temp_dir = tempfile.mkdtemp(prefix=f"airflow_stage_{self.task_id}_")
        # 自动拉取源数据到临时目录
        storage_download(self.source_data_path, self.temp_dir)
        # 把临时目录路径放到上下文供子类使用
        context["temp_stage_dir"] = self.temp_dir

    def post_execute(self, context, result=None):
        # 自动上传输出文件到外部存储
        output_local_path = os.path.join(self.temp_dir, "output")
        if os.path.exists(output_local_path):
            storage_upload(output_local_path, self.output_data_path)
        # 清理临时目录释放空间
        import shutil
        shutil.rmtree(self.temp_dir)

# 业务侧直接继承即可实现带暂存的DockerOperator
class StagedDockerOperator(AutoStageDataOperator, DockerOperator):
    def execute(self, context):
        # 把临时目录挂载到Docker容器内部即可
        self.volumes.append(f"{context['temp_stage_dir']}:/data")
        super().execute(context)

实现2：TaskFlow API拆分任务的轻量模式 #

如果不想封装基类，也可以用Airflow 2.x的TaskFlow API拆分三个步骤，用XCom传递临时目录标识，完全不需要修改现有Operator的逻辑：

from airflow.decorators import task, dag
from datetime import datetime

@dag(schedule_interval=None, start_date=datetime(2023, 1, 1))
def data_process_dag():
    @task
    def stage_in(source_path: str) -> str:
        temp_dir = tempfile.mkdtemp()
        storage_download(source_path, temp_dir)
        return temp_dir

    @task.docker(image="your_business_image:v1")
    def process_data(temp_dir: str):
        # 直接读写挂载的临时目录数据，处理完写回输出路径即可
        return {"status": "success"}

    @task
    def stage_out(temp_dir: str, target_path: str):
        storage_upload(f"{temp_dir}/output", target_path)
        shutil.rmtree(temp_dir)

    # 配置任务依赖
    temp_dir = stage_in("s3://your-bucket/source/")
    process_data(temp_dir)
    stage_out(temp_dir, "s3://your-bucket/target/")

注意：如果使用CeleryExecutor/KubernetesExecutor这类多节点执行器，需要保证三个任务被调度到同一个节点，可以通过queue参数把三个任务绑定到同一个专属队列，或者用KubernetesExecutor的affinity配置实现同节点调度，生产环境实测稳定性足够。

生产环境实操建议 #

• 所有临时目录都用动态生成的唯一路径，不要使用固定路径，避免多任务读写冲突
• 存储读写逻辑统一封装到工具类，不需要每个业务任务单独实现，支持S3/OSS/HDFS等多种存储后端无缝切换
• 大文件传输可以开启断点续传、分片上传能力，避免网络波动导致任务失败

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