使用 Airflow 构建 ETL 流水线

欢迎来到使用 Apache Airflow 构建 ETL 流水线的世界。在本教程中，我们将专注于使用 Polygon API 拉取股票市场数据，对其进行转换，然后将其加载到 SQLite 数据库中，以便轻松访问和处理。让我们开始吧！

什么是 Apache Airflow 与 Airflow ETL？

Apache Airflow 被视为数据编排与流水线管理的行业标准。凭借其编排复杂工作流、管理任务依赖关系、重试失败任务以及提供详尽日志的能力，它在数据科学家、机器学习工程师和 AI 从业者中非常受欢迎。

Airflow ETL 指的是使用 Apache Airflow 管理 ETL 流程。简单回顾一下，ETL 是一种数据集成方式，包括从各类来源抽取数据，将其转换为适合分析的格式，并将其加载到最终目的地（如数据仓库）。 

配置 Airflow 开发环境 

在开始用 Airflow 构建 ETL 流水线之前，我们必须先配置开发环境。有关如何配置开发环境的详细信息，请参阅我们的 Apache Airflow 入门教程。

我们还需要安装 Astro CLI。请查看由 Astronomer 维护的 Astro CLI，并按照其详尽文档中的说明进行安装。

创建 Airflow 项目

在配置好环境并安装 Astro CLI 后，我们创建一个 Airflow 项目。具体做法是打开终端，在选定路径下创建新目录。

 ~/Documents/data-engineering/ETL-pipeline/ 技术支持

在该目录的根路径下，运行以下命令以创建所需资源：

astro dev init技术支持

目录内容大致如下。具体输出可能有所差异。

├── dags/
├── include/
├── plugins/
├── tests/
├── airflow_settings.yaml
├── Dockerfile
├── packages.txt
└── requirements.txt技术支持

要启动您的项目，请执行以下命令： 

astro dev start技术支持

Airflow 环境大约需要一分钟启动。完成后，在浏览器中访问 localhost: 8080，您将看到 Airflow 的 UI。

现在，您已经准备好开始用 Airflow 开发自己的 ETL 流水线了！

设计 ETL 流水线

在动手写第一行代码之前，花时间规划流水线的各个组件非常重要。尤其是，先明确数据来源以及最终加载的目的地是个好习惯。通过提前识别数据源和目的地，我们也能了解数据在流程中将如何被转换。

在本示例中，我们将设计一个数据流水线，从 Polygon API 拉取股票市场数据，然后对其进行转换并加载到 SQLite 数据库。在这个案例中，源系统是 Polygon API，目的地是 SQLite 数据库。我们用一张图来说明： 

来源与目的地图示

根据我们作为数据工程师的经验，要将数据加载到 SQLite 数据库，需要先将其从 JSON 转换为表格格式。我们计划在从 Polygon API 拉取数据后，使用原生 Python 和 pandas 库对其进行转换。让我们更新可视化以展示这一变化：

架构图

通过补充这些信息，我们创建了一张架构图，即对系统的整体可视化表示。我们看到该流水线中的三个逻辑步骤，分别对应流程中的 E、T 和 L。

我们还可以把这些任务转化为一个有向无环图（DAG），它是一种特定的配置，用来定义 Airflow 要执行的全部任务、其执行顺序以及彼此之间的依赖。想要进一步了解有向无环图，请学习我们的数据工程入门课程，其中详细回顾了 Airflow DAG。

在企业数据场景中，许多团队还会使用技术规格说明（tech specs）等文档来达成一致并记录设计选择。对于我们的案例，创建一张表来记录数据流水线的设计选择就足够了。 

数据流水线设计选择

请记住，我们的 DAG 关注的是任务层面的细节，可能不涵盖全部信息。我们可以用第二张表记录需要进一步明确的细节。我们的问题包括：

• 该 DAG 将以多高的频率运行？ 
• 如果流水线中的某个任务失败会怎样？ 
• 如果我们想收集更多股票的数据怎么办？

待明确的问题

让我们回顾一下完成的工作：我们创建了架构图，记录了如何将流水线拆分为 Airflow 任务，并识别了配置 DAG 所需的高层信息。

如果您想进一步学习数据流水线的设计、开发与测试，请查看 DataCamp 的 数据流水线入门。在这里，您将掌握用 Python 构建 ETL 流水线的基础，以及确保方案健壮、弹性与可复用的最佳实践。

使用 Airflow 构建 ETL 流水线

我们将按顺序推进各步骤来组织 ETL 流水线的构建。采用结构化的方法可确保每个阶段都能精准执行。 

使用 Airflow 抽取数据

在从 Polygon API 拉取数据之前，您需要先访问 Polygon 并选择 Create API Key 按钮创建一个 API 令牌。请注意，本教程无需订阅付费版本——免费层即可生成 API 密钥并提供所需的全部功能。请务必复制并保存您的 API 密钥。

创建好 API 密钥后，我们就可以使用 Airflow 从 Polygon API 开始抽取数据了。我们将参考之前创建的技术规格表，其中包含 DAG 配置的细节，以辅助编写代码。 

from airflow import DAG 
from datetime import datetime, timedelta

with DAG(
    dag_id="market_etl",
    start_date=datetime(2024, 1, 1, 9),
    schedule="@daily",
    catchup=True,
    max_active_runs=1,
    default_args={
        "retries": 3,
        "retry_delay": timedelta(minutes=5)
    }
) as dag:技术支持

让我们添加第一个任务。我们将使用 TaskFlow API 和 requests 模块从 Polygon API 拉取股票的开盘/收盘数据。 

import requests
...
@task()

def hit_polygon_api(**context):
    # Instantiate a list of tickers that will be pulled and looped over
    stock_ticker = "AMZN"
    # Set variables
    polygon_api_key= "<your-api-key>"
  ds = context.get("ds")

  # Create the URL
    url= f"<https://api.polygon.io/v1/open-close/{stock_ticker}/{ds}?adjusted=true&apiKey={polygon_api_key}>"
    response = requests.get(url)
    # Return the raw data
    return response.json()技术支持

您会注意到这段代码中的几点： 

• hit_polygon_api 函数使用 @task 装饰器进行装饰。该装饰器会将函数转换为可在 DAG 中执行的 Airflow 任务。
• hit_polygon_api 的函数签名中定义了 context 参数。稍后将使用它来提取 ds 键对应的值。
• context 是一个包含任务和 DAG 元数据的字典。
• 当从 context 字典中获取 ds 时，我们得到的是 data_interaval_end 的日期，格式为 YYYY-mm-dd。 
• 为确保在运行 DAG 时执行新任务，我们需要调用 hit_polygon_api。

将以上内容整合后，创建 ETL 流水线第一部分的代码如下。

from airflow import DAG
from airflow.decorators import task
from datetime import datetime, timedelta import requests

with DAG(
    dag_id="market_etl",
    start_date=datetime(2024, 1, 1, 9),
    schedule="@daily",
    catchup=False,
    max_active_runs=1,
    default_args={
        "retries": 3,
        "retry_delay": timedelta(minutes=5)
    }
) as dag:
    # Create a task using the TaskFlow API
    @task()
    def hit_polygon_api(**context):
        # Instantiate a list of tickers that will be pulled and looped over
        stock_ticker = "AMZN"
        # Set variables
        polygon_api_key = "<your-api-key>"
        ds = context.get("ds")
        # Create the URL
        url = f"<https://api.polygon.io/v1/open-close/{stock_ticker}/{ds}?adjusted=true&apiKey={polygon_api_key}>"
        response = requests.get(url)
        # Return the raw data
        return response.json()
    hit_polygon_api()技术支持

您会发现，当该代码以 2024 年 1 月 1 日为日期运行时，响应会包含一个错误。由于当天市场休市，Polygon 会返回一条说明异常的响应。我们将在下一步中处理这个问题。

使用 Airflow 转换数据

从 Polygon API 抽取到数据后，我们就可以对其进行转换。

为此，我们将使用 TaskFlow API 创建另一个任务。该任务命名为 flatten_market_data，具有参数 polygon_response（即 hit_polygon_api 函数返回的原始数据）和 **context。我们稍后将更仔细地查看 polygon_response。

我们要执行的实际转换非常直接：把 Polygon API 返回的 JSON 扁平化为一个列表。关键在于，如果响应中不存在某个键，我们会为其提供特定的默认值。

例如，如果响应中不存在 from 键，我们将使用 Airflow 的上下文提供一个默认值。这解决了先前提到的问题——当市场休市时，响应仅包含有限的键。接下来，我们将把该列表转换为 pandas DataFrame，并将其返回。结果任务如下：

@task
def flatten_market_data(polygon_response, **context):
    # Create a list of headers and a list to store the normalized data in
    columns = {
        "status": "closed",
        "from": context.get("ds"),
        "symbol": "AMZN",
        "open": None,
        "high": None,
        "low": None,
        "close": None,
        "volume": None
    }
    # Create a list to append the data to
    flattened_record = []
    for header_name, default_value in columns.items():
        # Append the data
        flattened_record.append(polygon_response.get(header_name, default_value))
    # Convert to a pandas DataFrame
    flattened_dataframe = pd.DataFrame([flattened_record], columns=columns.keys())
    return flattened_dataframe技术支持

我们需要在 hit_polygon_api 与 flatten_market_data 任务之间添加依赖。为此，我们将 DAG 中的代码更新为如下所示：

import pandas as pd
with DAG(
    dag_id="market_etl",
    start_date=datetime(2024, 1, 1, 9),
    schedule="@daily",
    catchup=True,
    max_active_runs=1,
    default_args={
        "retries": 3,
        "retry_delay": timedelta(minutes=5)
    }
) as dag:
    @task()
    def hit_polygon_api(**context):
        ...
    @task
    def flatten_market_data(polygon_response, **context):
        # Create a list of headers and a list to store the normalized data in
        columns = {
            "status": None,
            "from": context.get("ds"),
            "symbol": "AMZN",
            "open": None,
            "high": None,
            "low": None,
            "close": None,
            "volume": None
        }
        # Create a list to append the data to
        flattened_record = []
        for header_name, default_value in columns.items():
            # Append the data
            flattened_record.append(polygon_response.get(header_name, default_value))
        # Convert to a pandas DataFrame
        flattened_dataframe = pd.DataFrame([flattened_record], columns=columns.keys())
        return flattened_dataframe
# Set dependencies
    raw_market_data = hit_polygon_api()
    flatten_market_data(raw_market_data)技术支持

这里，hit_polygon_api 任务的返回值被存储在 raw_market_data 中。然后，raw_market_data 作为 polygon_response 参数传递给 flatten_market_data 任务。这样的代码不仅建立了 hit_polygon_api 与 flatten_market_data 之间的依赖关系，还允许在这两个任务之间共享数据。

尽管我们进行的转换相对简单，Airflow 也能执行更为复杂的各类数据处理。除了使用原生任务外，还可以轻松利用 Airflow 丰富的 provider 构建的 hooks 与 operators，来编排使用 AWS Lambda、DBT 等工具的转换。

使用 Airflow 加载数据

我们已来到 ETL 流水线的最后一步。我们计划使用 SQLite 数据库和一个用 TaskFlow API 定义的最终任务来完成此步骤。

与之前相同，在创建任务时我们将定义一个参数，命名为 flattened_dataframe。这使得 flatten_market_data 任务返回的数据可以传入新任务。

在编写将数据加载到 SQLite 数据库的代码之前，我们需要先在 Airflow UI 中创建一个连接。要打开连接页面，请按照以下步骤：

• 打开 Airflow UI
• 将鼠标悬停在 Admin 上
• 选择 Connections。 
• 点击 + 图标创建一个新连接。

随后您会进入如下界面：

Airflow 连接页面

填写连接页面时，请执行以下操作：

• 将 Connection Type 更改为 Sqlite。
• 将 Connection Id 设置为 "market_database_conn"。
• 在 Host 字段中添加 "/usr/local/airflow/market_data.db"。

该连接的配置应类似下图。完成后点击 Save。

未保存的指向 SQLite 的 Airflow 连接

现在我们已经创建了连接，可以在任务中使用 SqliteHook 来获取该信息。请看下面的代码。

from airflow.providers.sqlite.hooks.sqlite import SqliteHook技术支持

@task
def load_market_data(flattened_dataframe):
    # Pull the connection
    market_database_hook = SqliteHook("market_database_conn")
    market_database_conn = market_database_hook.get_sqlalchemy_engine()
    # Load the table to Postgres, replace if it exists
    flattened_dataframe.to_sql(
        name="market_data",
        con=market_database_conn,
        if_exists="append",
        index=False
    )
    # print(market_database_hook.get_records("SELECT * FROM market_data;"))技术支持

通过这段代码，我们创建了到上一步所指定 SQLite 数据库的连接。随后，使用 .get_sqlalchemy_engine() 方法从 hook 中检索连接引擎。调用 flattened_dataframe 的 .to_sql() 方法时，将该引擎作为 con 参数传入。

请注意，数据写入的表名为 market_data，若该表已存在，则会在其后追加。在测试时，我通常会通过检索并打印记录来检查数据是否已写入。您可以取消注释该任务最后一行来实现这一点。

把这些全部整合起来，我们的代码大致如下：

from airflow import DAG
from airflow.decorators import task
from airflow.providers.sqlite.hooks.sqlite import SqliteHook
from datetime import datetime, timedelta
import requests
import pandas as pd

with DAG(
    dag_id="market_etl",
    start_date=datetime(2024, 1, 1, 9),
    schedule="@daily",
    catchup=True,
    max_active_runs=1,
    default_args={
        "retries": 3,
        "retry_delay": timedelta(minutes=5)
    }
) as dag:
    # Create a task using the TaskFlow API
    @task()
    def hit_polygon_api(**context):
        # Instantiate a list of tickers that will be pulled and looped over
        stock_ticker = "AMZN"
        # Set variables
        polygon_api_key = "<your-api-key>"
        ds = context.get("ds")
        # Create the URL
        url = f"<https://api.polygon.io/v1/open-close/{stock_ticker}/{ds}?adjusted=true&apiKey={polygon_api_key}>"
        response = requests.get(url)
        # Return the raw data
        return response.json()
    @task
    def flatten_market_data(polygon_response, **context):
        # Create a list of headers and a list to store the normalized data in
        columns = {
            "status": None,
            "from": context.get("ds"),
            "symbol": "AMZN",
            "open": None,
            "high": None,
            "low": None,
            "close": None,
            "volume": None
        }
        # Create a list to append the data to
        flattened_record = []
        for header_name, default_value in columns.items():
            # Append the data
            flattened_record.append(polygon_response.get(header_name, default_value))
        # Convert to a pandas DataFrame
        flattened_dataframe = pd.DataFrame([flattened_record], columns=columns.keys())
        return flattened_dataframe
    @task
    def load_market_data(flattened_dataframe):
        # Pull the connection
        market_database_hook = SqliteHook("market_database_conn")
        market_database_conn = market_database_hook.get_sqlalchemy_engine()
        # Load the table to SQLite, append if it exists
        flattened_dataframe.to_sql(
            name="market_data",
            con=market_database_conn,
            if_exists="append",
            index=False
        )
    # Set dependencies between tasks
    raw_market_data = hit_polygon_api()
    transformed_market_data = flatten_market_data(raw_market_data)
    load_market_data(transformed_market_data)技术支持

同样地，我们更新了依赖关系，使 flatten_market_data 任务返回的数据可以传递给 load_market_data 任务。该 DAG 的图形视图大致如下：

ETL 流水线的图形视图

测试

现在您已经构建了第一个 Airflow DAG，是时候确保它能够正常工作了。常见的方式之一是端到端运行 DAG。

在 Airflow UI 中执行此操作时，导航到您的 DAG，将开关从蓝色切换为 active。由于 catchup 设置为 True，将会排队一个 DAG 运行并开始执行。如果任务执行成功，UI 中对应的方块会变为绿色。如果 DAG 中的所有任务都成功，DAG 将被标记为 success，并触发下一次运行。

如果某个任务失败，其状态将为 up for retry 并以黄色标记。此时，最好查看该任务的日志。方法是在网格视图中点击黄色方块并选择 Logs。您可以在此找到异常信息并开始排查。如果某任务失败的次数超过指定的重试次数，该任务与 DAG 的状态将被设置为 failed。

除了端到端测试 DAG，Airflow 还便于编写单元测试。使用 astro dev start 初始创建环境时，会为您创建一个 tests/ 目录。您可以在其中为 DAG 及其组件添加单元测试。

下面是一个针对我们 DAG 配置的单元测试。该测试验证定义时设置的若干参数，如 start_date、schedule 和 catchup。编写好测试后，导航到项目根目录并运行：

from airflow.models.dagbag import DagBag
from datetime import datetime
import pytz

def test_market_etl_config():
    # Pull the DAG
    market_etl_dag = DagBag().get_dag("market_etl")
    # Assert start date, schedule, and catchup
    assert market_etl_dag.start_date == datetime(2024, 3, 25, 9, tzinfo=pytz.UTC)
    assert market_etl_dag.schedule_interval == "@daily"
    assert market_etl_dag.catchup技术支持

astro dev pytest技术支持

该命令会执行 tests/ 目录中的所有单元测试。如果您只想运行单个测试，可以将文件路径作为参数传给上述命令。除了使用 Astro CLI 执行测试外，也可以使用任意 Python 测试运行器来编写与执行单元测试。

对于个人项目，编写单元测试有助于确保代码表现符合预期。在企业环境中，单元测试几乎是必需的。大多数数据团队会利用某种 CI/CD 工具来部署其 Airflow 项目。该流程通常涉及运行单元测试并验证其输出，以确保您编写的 DAG 已准备好投产。有关单元测试的更多信息，请参阅我们的 Pytest 单元测试实战教程，以及Python 测试入门。

进阶 Airflow 技巧与方法

我们已经构建了一个简单且有效的数据流水线，并使用了转换与持久化技术。在其他场景中，Airflow 也能通过 provider 构建与自定义的运算符来编排复杂工作流，处理 TB 级数据。

例如，S3ToSnowflakeOperator 与 DatabricksRunNowOperator 都能轻松与更大的数据栈集成。在兴趣项目中使用这类运算符会比较棘手。比如，要使用 S3ToSnowflakeOperator，您需要同时具备 AWS 与 Snowflake 账户，并为要传输数据的资源完成相应配置。

除了 ETL 工作流，Airflow 还支持 ELT 工作流。对于使用云数据仓库的团队而言，ELT 正日益成为行业标准。设计数据流水线时请牢记这一点。

在我们流水线的 加载 环节，我们创建了到 SQLite 数据库的连接，随后检索并用于数据持久化。连接（有时称为 Secrets）是 Airflow 的一个功能，旨在简化与源和目的地系统的交互。使用这些连接存储敏感信息（例如您的 Polygon API 密钥）可以增强代码的安全性。这种做法还可使您将凭据与代码库分离管理。在可行的情况下，建议广泛使用 Connections，以确保工作流既安全又有序。

您可能注意到，我们在 hit_polygon_api 和 flatten_market_data 任务中将股票代码 “AMZN” 硬编码了。这使我们能够针对单只股票进行抽取、转换与加载。但如果您想用同一套代码处理多只股票怎么办？幸运的是，可以轻松地动态生成 DAG。只需对代码做少量重构，就可以遍历股票代码列表，并将股票代码参数化。这有助于让您的 DAG 更加模块化与可移植。关于动态生成 DAG 的信息，请查看 Astronomer 的在 Airflow 中动态生成 DAG文档。

结语

恭喜您！您已经使用 Python、pandas 和 SQLite 构建了一个 Airflow DAG，实现了从 Polygon API 抽取、转换与加载股票市场数据。在此过程中，您还提升了绘制架构图与技术规格、创建 Airflow 连接和测试 DAG 的技能。随着您继续深入掌握 Airflow，不妨尝试更多高级技巧，让流水线更加健壮、弹性且可复用。  

想要进一步学习，请查看以下资源：

• Python 中的 Airflow 入门课程
• Apache Airflow 入门教程
• Airflow 与 Prefect 对比
• 主流 Airflow 替代方案
• 热门 Airflow 面试题