Apache Airflow 中的设置和拆卸

2025 年 6 月 11 日 | 阅读 9 分钟

引言

Apache Airflow 是一个强大的平台，用于以编程方式编写、调度和监控工作流。它允许用户将工作流定义为任务的有向无环图 (DAG)，这些任务可以分布式执行。Airflow 的关键功能之一是其处理复杂依赖关系、重试和错误处理的能力，使其成为编排数据管道的理想选择。

Apache Airflow 中的设置和拆卸

什么是设置和清理？

在 Apache Airflow 的上下文中

设置任务 在主工作流开始之前执行。这些任务会准备环境、初始化资源，并确保一切就绪以顺利执行。
清理任务 在主工作流完成后执行。它们负责清理资源、释放锁，并确保环境保持一致状态。

设置和清理的重要性

资源管理： 确保正确初始化和清理数据库连接和文件句柄等资源，防止泄露。
一致性： 确保在执行前正确设置环境，并在执行后正确清理。
错误处理： 有助于在工作流执行前识别和解决潜在问题，并确保在发生故障时进行清理。
可重现性： 确保工作流可以多次以一致的方式执行。

设置 Apache Airflow

在深入研究设置和清理任务的具体细节之前，让我们先了解设置 Apache Airflow 的过程。

安装

可以使用 pip 安装 Apache Airflow。建议在虚拟环境中安装 Airflow，以避免与其他 Python 包发生冲突。

配置

安装 Airflow 后，需要初始化 Airflow 数据库。此数据库将存储有关工作流的元数据，包括 DAG、任务和执行日志。

初始化数据库后，需要为 Airflow Web 界面创建一个管理员用户。

 
airflow users create \
    --username admin \
    --firstname Admin \
    --lastname User \
    --role Admin \
    --email admin@example.com   

数据库设置

默认情况下，Airflow 使用 SQLite 作为其数据库后端。但是，SQLite 不适合生产使用。对于生产环境，应使用更强大的数据库，如 PostgreSQL 或 MySQL。

要配置 Airflow 使用不同的数据库，需要修改 airflow.cfg 文件。找到 sql_alchemy_conn 设置并使用数据库的连接字符串进行更新。

更新配置后，需要重新初始化数据库。

Web 服务器和调度器

要启动 Airflow Web 服务器，请运行以下命令

Web 服务器提供了一个用户界面来管理和监控工作流。

要启动 Airflow 调度器，请运行以下命令

调度器负责根据定义的计划执行工作流。

创建基本 DAG

在实现设置和清理任务之前，我们需要创建一个基本 DAG。DAG 是函数的集合，这些函数以反映其依赖关系和关系的方式进行组织。

以下是一个简单的 DAG 示例，它包含三个任务：

 
from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from datetime import datetime
default_args = {
    'owner': 'airflow',
    'start_date': datetime(2023, 1, 1),
}
dag = DAG(
    'example_dag',
    default_args=default_args,
    description='A simple example DAG',
    schedule_interval='@daily',
)
start_task = DummyOperator(task_id='start_task', dag=dag)
middle_task = DummyOperator(task_id='middle_task', dag=dag)
end_task = DummyOperator(task_id='end_task', dag=dag)
start_task >> middle_task >> end_task   

在此示例中，DAG 由三个任务组成：start_task、middle_task 和 end_task。这些任务按顺序执行，start_task 首先运行，然后是 middle_task，最后是 end_task。

实现设置和清理任务

现在我们有了一个基本 DAG，让我们向其中添加设置和清理任务。

设置任务

设置任务在主工作流开始之前执行。它们负责准备环境和初始化资源。

以下是一个创建临时目录的设置任务示例：

 
import os
import tempfile
from airflow.operators.python_operator import PythonOperator
def create_temp_dir():
    temp_dir = tempfile.mkdtemp()
    os.environ['TEMP_DIR'] = temp_dir
    print(f"Created temporary directory: {temp_dir}")
setup_task = PythonOperator(
    task_id='setup_task',
    python_callable=create_temp_dir,
    dag=dag,
)
setup_task >> start_task   

在此示例中，create_temp_dir 函数创建一个临时目录，并将其路径存储在环境变量中。setup_task 是一个调用此函数的 PythonOperator。setup_task 在 start_task 之前执行。

清理任务

清理任务在主工作流完成后执行。它们负责清理资源并确保环境保持一致状态。

以下是一个删除设置任务创建的临时目录的清理任务示例：

 
import shutil
def delete_temp_dir():
    temp_dir = os.environ.get('TEMP_DIR')
    if temp_dir and os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
        print(f"Deleted temporary directory: {temp_dir}")
teardown_task = PythonOperator(
    task_id='teardown_task',
    python_callable=delete_temp_dir,
    dag=dag,
)
end_task >> teardown_task 

在此示例中，delete_temp_dir 函数从环境变量中检索临时目录的路径并将其删除。teardown_task 是一个调用此函数的 PythonOperator。teardown_task 在 end_task 之后执行。

示例：数据管道中的设置和清理

让我们考虑一个更复杂的示例，其中有一个处理 CSV 文件的数据管道。管道包含以下任务：

设置任务： 创建一个临时目录并下载 CSV 文件。
处理任务： 读取 CSV 文件，处理数据，并将结果写入新文件。
清理任务： 删除临时目录和下载的 CSV 文件。

以下是完整的 DAG：

 
import os
import tempfile
import shutil
import requests
import pandas as pd
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime
default_args = {
    'owner': 'airflow',
    'start_date': datetime(2023, 1, 1),
}
dag = DAG(
    'data_pipeline_dag',
    default_args=default_args,
    description='A data pipeline DAG with setup and teardown tasks',
    schedule_interval='@daily',
)
def create_temp_dir():
    temp_dir = tempfile.mkdtemp()
    os.environ['TEMP_DIR'] = temp_dir
    print(f"Created temporary directory: {temp_dir}")
def download_csv_file():
    temp_dir = os.environ.get('TEMP_DIR')
    url = 'https://example.com/data.csv'
    response = requests.get(url)
    with open(csv_file_path, 'wb') as f:
        f.write(response.content)
    print(f"Downloaded CSV file to: {csv_file_path}")
def process_csv_file():
    temp_dir = os.environ.get('TEMP_DIR')
    df = pd.read_csv(csv_file_path)
    # Process the data (e.g., filter, transform, etc.)
    processed_file_path = os.path.join(temp_dir, 'processed_data.csv')
    df.to_csv(processed_file_path, index=False)
    print(f"Processed data saved to: {processed_file_path}")
def delete_temp_dir():
    temp_dir = os.environ.get('TEMP_DIR')
    if temp_dir and os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
        print(f"Deleted temporary directory: {temp_dir}")
setup_task = PythonOperator(
    task_id='setup_task',
    python_callable=create_temp_dir,
    dag=dag,
)
download_task = PythonOperator(
    task_id='download_task',
    python_callable=download_csv_file,
    dag=dag,
)
process_task = PythonOperator(
    task_id='process_task',
    python_callable=process_csv_file,
    dag=dag,
)
teardown_task = PythonOperator(
    task_id='teardown_task',
    python_callable=delete_temp_dir,
    dag=dag,
)
setup_task >> download_task >> process_task >> teardown_task  

在此示例中，setup_task 创建临时目录，download_task 下载 CSV 文件，process_task 处理数据，teardown_task 删除临时目录。

设置和清理的最佳实践

幂等性：确保您的设置和清理任务是幂等的，也就是说，它们可以多次执行而不会引起问题。在重试或失败的情况下，这一点尤其重要。
错误处理：在设置和清理任务中实现适当的错误处理。如果设置任务失败，工作流不应继续。同样，如果清理任务失败，它不应使环境处于不一致状态。
资源清理：始终确保在清理任务中正确清理资源。这包括关闭数据库连接、删除临时文件和释放任何锁。
日志记录：使用日志记录来跟踪设置和清理任务的执行情况。这将有助于调试和监控工作流。
测试：彻底测试您的设置和清理任务，以确保它们按预期工作。这包括对边缘情况和失败场景的测试。

高级用例

动态设置和清理

在某些情况下，您可能需要根据输入参数或环境状态动态创建设置和清理任务。例如，您可能希望为工作流的每次执行创建一个具有唯一名称的临时目录。

以下是如何动态创建临时目录的示例：

 
def create_temp_dir(**kwargs):
    temp_dir = tempfile.mkdtemp()
    kwargs['ti'].xcom_push(key='temp_dir', value=temp_dir)
    print(f"Created temporary directory: {temp_dir}")
def delete_temp_dir(**kwargs):
    temp_dir = kwargs['ti'].xcom_pull(key='temp_dir')
    if temp_dir and os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
        print(f"Deleted temporary directory: {temp_dir}")
setup_task = PythonOperator(
    task_id='setup_task',
    python_callable=create_temp_dir,
    provide_context=True,
    dag=dag,
)
teardown_task = PythonOperator(
    task_id='teardown_task',
    python_callable=delete_temp_dir,
    provide_context=True,
    dag=dag,
)   

在此示例中，create_temp_dir 函数创建临时目录并将其路径存储在 XCom 中，XCom 是 Airflow 中用于在任务之间共享数据的一种机制。delete_temp_dir 函数从 XCom 中检索路径并删除目录。

条件清理

在某些情况下，您可能希望根据工作流的结果有条件地执行清理任务。例如，如果工作流失败，您可能希望跳过清理任务。

以下是如何有条件地执行清理任务的示例：

 
def conditional_teardown(**kwargs):
    if kwargs['ti'].xcom_pull(key='workflow_status') == 'success':
        delete_temp_dir(**kwargs)
    else:
        print("Skipping teardown due to workflow failure")
teardown_task = PythonOperator(
    task_id='teardown_task',
    python_callable=conditional_teardown,
    provide_context=True,
    dag=dag,
)   

在此示例中，conditional_teardown 函数检查工作流的状态（存储在 XCom 中），并且仅当工作流成功时才执行 delete_temp_dir 函数。

错误处理和清理

如果在工作流执行过程中发生错误，确保仍会执行清理任务以清理资源非常重要。您可以通过使用 Airflow 的 trigger_rule 参数来实现这一点。

这里有一个例子

 
teardown_task = PythonOperator(
    task_id='teardown_task',
    python_callable=delete_temp_dir,
    provide_context=True,
    trigger_rule='all_done',
    dag=dag,
)   

在此示例中，trigger_rule 参数设置为 all_done，这意味着无论之前的任务成功还是失败，都会执行清理任务。

监控和日志记录

监控和日志记录对于确保您的设置和清理任务按预期工作至关重要。Airflow 提供了多种监控和日志记录工具，包括：

Airflow Web 界面：Airflow Web 界面提供了 DAG 的可视化表示，包括每个任务的状态。您可以使用 Web 界面来监控设置和清理任务的执行情况。
日志：Airflow 存储每个任务执行的日志。您可以在 Airflow Web 界面中或直接在日志文件中查看日志。日志对于调试和排除设置和清理任务的问题非常有用。
指标：Airflow 提供可用于监控工作流性能的指标。您可以使用 Prometheus 和 Grafana 等工具来可视化这些指标。

使用 Airflow 插件自动化设置和清理

Airflow 允许用户通过自定义插件扩展其功能。您可以将可重用的设置和清理任务定义为插件的一部分，以确保多个 DAG 之间的一致性。这种方法在大规模环境中特别有用，因为多个工作流共享通用的设置和清理过程。

创建自定义设置和清理插件

导航到您的 Airflow 插件目录 ($AIRFLOW_HOME/plugins/)。
创建一个 Python 文件，例如 setup_teardown_plugin.py。
定义可重用的设置和清理函数。

这是一个例子：

 
from airflow.plugins_manager import AirflowPlugin
from airflow.operators.python_operator import PythonOperator
import os
import tempfile
import shutil
def setup_environment():
    """Creates a temporary directory and sets an environment variable."""
    temp_dir = tempfile.mkdtemp()
    os.environ["TEMP_DIR"] = temp_dir
    print(f"Setup: Created temporary directory at {temp_dir}")
def teardown_environment():
    """Deletes the temporary directory and cleans up resources."""
    temp_dir = os.environ.get("TEMP_DIR")
    if temp_dir and os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
        print(f"Teardown: Deleted temporary directory at {temp_dir}")
class SetupTeardownPlugin(AirflowPlugin):
    name = "setup_teardown_plugin"
    operators = [PythonOperator]   

在 DAG 中使用插件

插件可用后，您可以轻松地将设置和清理任务集成到您的 DAG 中：

 
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime
from setup_teardown_plugin import setup_environment, teardown_environment
default_args = {
    "owner": "airflow",
    "start_date": datetime(2023, 1, 1),
}
dag = DAG(
    "dag_with_plugin",
    default_args=default_args,
    description="DAG using a setup and teardown plugin",
    schedule_interval="@daily",
)
setup_task = PythonOperator(
    task_id="setup_task",
    python_callable=setup_environment,
    dag=dag,
)
teardown_task = PythonOperator(
    task_id="teardown_task",
    python_callable=teardown_environment,
    dag=dag,
)
setup_task >> teardown_task   

这种方法通过确保所有 DAG 共享标准化的设置和清理过程来提高可维护性。

在设置和清理中处理秘密信息

许多工作流需要访问敏感信息，如数据库凭据、API 密钥或 SSH 私钥。直接将它们存储在 DAG 或环境变量中存在安全风险。相反，Airflow 提供了连接和变量来安全地存储和检索秘密信息。

使用 Airflow 连接

Airflow 的连接允许您安全地存储凭据并在 DAG 中检索它们。

通过 Airflow Web UI (Admin -> Connections) 添加连接。
使用 AirflowHook 动态获取凭据。

示例

 
from airflow.hooks.base_hook import BaseHook
def setup_database_connection():
    """Retrieve database credentials from Airflow connections and establish a connection."""
    connection = BaseHook.get_connection("my_database")
    db_host = connection.host
    db_user = connection.login
    db_password = connection.password
    print(f"Connected to database at {db_host} as {db_user}")   

使用 Airflow 变量

Airflow 变量允许存储非敏感的配置值。

通过 Airflow Web UI (Admin -> Variables) 添加变量。
在 DAG 中检索值。

示例

 
from airflow.models import Variable
def setup_api_keys():
    """Retrieve API keys from Airflow Variables."""
    api_key = Variable.get("api_key")
    print(f"Retrieved API Key: {api_key}")   

通过利用连接和变量，您的设置和清理任务将保持安全且易于管理。

将设置和清理与 KubernetesPodOperator 集成

对于云原生部署，KubernetesPodOperator 允许在容器内执行设置和清理任务。

示例

 
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
setup_task = KubernetesPodOperator(
    task_id="setup_task",
    namespace="airflow",
    image="python:3.8",
    cmds=["python", "-c", "import os; os.makedirs('/tmp/workdir', exist_ok=True)"],
    name="setup_task",
    is_delete_operator_pod=True,
    dag=dag,
)
teardown_task = KubernetesPodOperator(
    task_id="teardown_task",
    namespace="airflow",
    image="python:3.8",
    cmds=["rm", "-rf", "/tmp/workdir"],
    name="teardown_task",
    is_delete_operator_pod=True,
    dag=dag,
)   

使用 Kubernetes、PodOperator，设置和清理任务在隔离的容器中运行，确保执行的干净性。

设置和清理任务的重要性： 确保适当的资源管理，维护一致的环境，并提高工作流的可重现性。

结果： 有助于在 Apache Airflow 中创建高效、可靠且易于维护的工作流，确保数据管道的干净且受控的执行。

下一个主题使用公共接口与外部服务和应用程序集成

Apache Airflow 中的设置和拆卸

引言