基于 Airflow 的面向对象实现

本文链接：https://blog.csdn.net/zfyzlldoom/article/details/146105262

背景

一切皆对象，本来是一句被 Java 洗脑的话，但随着这么多年 SRE 及 DevOps 经验，在一些工具的封装中，这句话的含金量越来越高。
最近对 ETL 很感兴趣，因为以前有过类似的想法，而业内这个专有名词对我的想法给予了最高的肯定，随便 Google 了下相关的框架，Airflow 就扑了过来。

Airflow

Airflow， Apache 基金会 flow 赫赫有名的开源框架，社区文档也很丰富，以 2.0 版本为分界线，之前上下文（context）的形式很主流，而后面以装饰器形式简化了使用也更易理解。目前的最新版本 2.10 更是兼容了 Python3.8+.

在粗略地看了一遍社区文档后发现 Airflow 对于面向对象集成相对较少，可能原因是这版本的 Airflow 从使用角度来说已经很简单了，@dag, @task, Context 完全够用，而且基本所有代码都在一个单独的文件实行。 本文仅针对那些对面向对象有执念的朋友一起探讨，如果说错或者冗余还望海涵！

用法

官方用例

先简单介绍 Airflow 推荐的用法

from airflow.decorators import dag, task

@dag(
    "test",
    schedule=None,
    catchup=False,
    tags=['test'],
)
def test_dag():
    """
    ### Test Dag Documentation
    This is a test dag to do a simple demo.
    :return:
    :rtype:
    """

    @task
    def extract():
        """
        ### Extract Task

        :return:
        :rtype:
        """

        return {"a": 2.0, "b": 3.0}

    @task
    def transform(order_data_dict: dict):
        """
        #### Transform task
        A simple Transform task which takes in the collection of order data and
        computes the total order value.
        """
        total_order_value = 0

        for value in order_data_dict.values():
            total_order_value += value

        return {"total": total_order_value}

    @task
    def load(order_data_dict: dict):
        """
        #### Load task
        A simple Load task which takes in the result of the Transform task and
        instead of saving it to end user review, just prints it out.
        """
        total_order_value = order_data_dict["total"]
        print(f"Total order value is: {total_order_value:.2f}")
        raise Exception("test on_failure_callback")

    order_data = extract()
    order_summary = transform(order_data)
    load(order_summary)


test_dag()

说明：

dag 是一个有序执行流，可以理解为定义任务执行顺序
task 是实际运行的任务，业务逻辑主要实现的地方
test_dag()很重要，是用来实例化dag，这样才能被 scheduler 发现，而且需要在最顶层声明
这样是基本够用的，但出于对 code 的简约化，如果我能够设计一个基类，让子类只需要在意 extract, transform,load 的具体实现，屏蔽掉 Dag 的定义和 task，这对于大型项目来说是会有帮助的。

面向对象

基类(BaseETLDag)

from abc import abstractmethod
from datetime import datetime

from airflow import DAG
from airflow.decorators import task
from airflow.models.baseoperator import chain

from interface.error import DagValueError
from interface.loader import BaseLoader


class BaseETLDag(BaseLoader):
    """ETL DAG Base Class"""

    def __init__(self, dag_id: str, schedule=None, tags: list = None,
                 start_date: datetime = datetime(2025, 1, 1),
                 catchup: bool = False,
                 **kwargs):
        self.dag_id = dag_id
        self.schedule = schedule
        self.tags = tags if tags else ["ETL"]
        self.start_date = start_date
        self.catchup = catchup
        self.dag = DAG(
            self.dag_id,
            schedule=self.schedule,
            catchup=self.catchup,
            tags=self.tags,
            start_date=self.start_date,
            **kwargs
        )

        self.dag_tasks_priority = {
            10: self.extract,
            20: self.transform,
            30: self.load
        }

    def create_dag(self):
        """Create DAG"""
        with self.dag:
            chain(*self.set_tasks_sequence())

    def set_tasks_sequence(self):
        return [v() for k, v in sorted(self.dag_tasks_priority.items(), key=lambda x: x[0])]

    @staticmethod
    @abstractmethod
    @task
    def extract(*args, **kwargs):
        raise NotImplementedError("subclass must implement extract method.")

    @staticmethod
    @abstractmethod
    @task
    def transform(**context):
        raise NotImplementedError("subclass must implement transform method.")

    @staticmethod
    @abstractmethod
    @task
    def load(**context):
        raise NotImplementedError("subclass must implement load method.")

说明：

BaseETLDag 定义了基类
三个抽象方法 extract, transform, load 主要给子类实现。这里注意: 当我们要用@task时，被装饰的方法只能是普通方法或者静态方法，所以这里要申明是@staticmethod
create_dag方法设置extract, transform, load顺序。所有@task装饰的方法默认都会传递 context，这是一个字典，里面有一个非常重要的 key -- ti(它的用法为data = context["ti"].xcom_pull(task_ids=upstream_task_id)), 他可以获取上游所有 task 的 Xcom value，简单说就是上游 task return 的 data. 这一点非常有用，这使得我们只需要定义 task 的顺序，而不需要在意值传递。

子类(CrawlerDag)

from airflow.decorators import task

from interface.dag import BaseETLDag
from interface.structure_data import ExtractData, TransformData


class CrawlerDag(BaseETLDag):

    def __init__(self, *args, **kwargs):
        super().__init__(dag_id="crawler", tags=["crawler", "etl"])
    
    @staticmethod
    @task
    def extract(*args, **kwargs) -> ExtractData:
        return ExtractData(**{"a": 2.0, "b": 3.0})

    @staticmethod
    @task
    def transform(**context) -> TransformData:
        transform_data = context["ti"].xcom_pull(task_ids="extract")
        return TransformData(**transform_data)

    @staticmethod
    @task
    def load(**context):

        print(f"Load Func.")


CrawlerDag().create_dag()