精通Airflow任务: Sensors, Operators & Hooks !

原创已于 2024-12-27 10:59:26 修改 · 1.4k 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#python #airflow #数据工程 #数据科学

于 2024-12-04 16:01:13 首次发布

数据分析工程同时被 3 个专栏收录

284 篇文章

订阅专栏

python

259 篇文章

订阅专栏

Airflow

28 篇文章

订阅专栏

部署运行你感兴趣的模型镜像

欢迎回来!在我们之前的博客中，我们深入研究了Airflow的基本原理，探索了它的目的、功能和机制。现在，让我们进一步了解Airflow的任务，并掌握如何在数据到达服务器时自动触发DAG。

尽管DAG用于构造作业并定义作业将在其中执行的上下文，但DAG实际上并不执行计算。另一方面，Airflow任务是“完成需要的工作”。任务有两种：Sensors和Operators 。Operators 任务执行一些显式操作，Sensors暂停相关任务的执行，直到满足某个条件。理论上，Operators 能够执行任何Python函数。任何过程或数据结构的状态也可以使用Sensors来确定。

Airflow Operators

执行引擎通常扮演操作员的角色。根据Operators 类型，任务由Operators创建，执行特定任务。例如，创建PythonOperator可以执行某个Python方法。常用Operators类别有：

DummyOperator

这是什么都不做的Operators，主要用于在 DAG（有向无环图）中占位，构建任务流程的结构和依赖关系。

from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from datetime import datetime

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('dummy_operator_dag', default_args=default_args, schedule_interval='@once') as dag:
    start_task = DummyOperator(task_id='start_task')
    end_task = DummyOperator(task_id='end_task')
    start_task >> end_task

在这个 DAG 中，start_task和end_task只是作为占位符来展示任务的开始和结束，中间没有实际的操作内容。

PythonOperator：

作用：用于执行自定义的 Python 函数。可以将任意 Python 函数封装成一个 Airflow 任务，方便在 DAG 中进行调度。

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime

def print_hello():
    print("Hello, Airflow!")

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('python_operator_dag', default_args=default_args, schedule_interval='@once') as dag:
    python_task = PythonOperator(
        task_id='python_task',
        python_callable=print_hello
    )

在这个示例中，python_task会在执行时调用print_hello函数，从而在日志中打印出Hello, Airflow!。

BashOperator：

作用：用于执行 Bash 命令。如果你的任务涉及到在命令行中运行 Bash 脚本或者命令，就可以使用这个操作符。

from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from datetime import datetime

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('bash_operator_dag', default_args=default_args, schedule_interval='@once') as dag:
    bash_task = BashOperator(
        task_id='bash_task',
        bash_command='echo "Running a Bash command in Airflow"'
    )

这里的bash_task会在执行时运行指定的 Bash 命令echo "Running a Bash command in Airflow"，将相应的字符串输出到日志或者终端。

MySqlOperator（以 MySQL 为例）：

作用：用于在 MySQL 数据库中执行 SQL 语句，比如创建表、插入数据、查询数据等操作。

from airflow import DAG
from airflow.providers.mysql.operators.mysql import MySqlOperator
from datetime import datetime

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('mysql_operator_dag', default_args=default_args, schedule_interval='@once') as dag:
    mysql_task = MySqlOperator(
        task_id='mysql_task',
        mysql_conn_id='your_mysql_conn_id',  # 你的MySQL连接ID
        sql="CREATE TABLE IF NOT EXISTS test_table (id INT AUTO_INCREMENT PRIMARY KEY, name VARCHAR(255))"
    )

在这个示例中，mysql_task会通过指定的mysql_conn_id连接到 MySQL 数据库，并执行创建表的 SQL 语句。

FileTransferOperator（以 S3 文件传输为例）：

作用：用于在不同的存储系统之间传输文件，这里以在本地文件系统和 Amazon S3 之间传输文件为例。

from airflow import DAG
from airflow.providers.amazon.aws.transfer.s3_to_local import S3ToLocalFilesystemOperator
from datetime import datetime

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('file_transfer_operator_dag', default_args=default_args, schedule_interval='@once') as dag:
    file_transfer_task = S3ToLocalFilesystemOperator(
        task_id='file_transfer_task',
        aws_conn_id='your_aws_conn_id',  # 你的AWS连接ID
        s3_bucket='your_s3_bucket_name',
        s3_key='your_s3_file_key',
        local_file_path='local/path/to/save/file'
    )

这个示例中的file_transfer_task会从指定的 S3 存储桶（your_s3_bucket_name）和文件键（your_s3_file_key）获取文件，并将其传输到本地文件系统的指定路径（local/path/to/save/file）。S3对象存储可以是互联网，也可以本地搭建minio进行测试。

mailOperator：

作用：用于发送电子邮件。可以在任务执行成功、失败或者其他特定情况下发送邮件通知相关人员。

from airflow import DAG
from airflow.operators.email_operator import EmailOperator
from datetime import datetime

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('email_operator_dag', default_args=default_args, schedule_interval='@once') as dag:
    email_task = EmailOperator(
        task_id='email_task',
        to='recipient@example.com',
        subject='Airflow Task Notification',
        html_content='<p>Your Airflow task has completed.</p>'
    )

在这个示例中，email_task会在执行时发送一封邮件到recipient@example.com，邮件主题是Airflow Task Notification，内容是包含一段 HTML 格式的消息Your Airflow task has completed.。

Airflow Sensors

Apache Airflow有一些专门的Operators ，用来等待某些事情发生。它们被称为Sensors。在标记Sensors为成功并允许执行下游操作之前，Sensors首先检查是否满足预定条件。如果使用得当，可以让DAG成为比事件驱动工具表现的更优秀。

Sensors类型

SQL Sensor

用于查询数据库，等待查询结果满足特定条件（如返回结果大于 0）后再继续执行任务。
假设我们有一个 MySQL 数据库，我们要检查一个表中是否有新的数据记录（通过检查行数是否大于 0）。

from airflow import DAG
from airflow.sensors.sql import SqlSensor
from datetime import datetime
from airflow.operators.dummy_operator import DummyOperator

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('sql_sensor_dag', default_args=default_args, schedule_interval='@once') as dag:
    start_task = DummyOperator(task_id='start_task')
    wait_task = SqlSensor(
        task_id='wait_task',
        conn_id='your_mysql_conn_id',  # 你的MySQL连接ID
        sql="SELECT COUNT(*) FROM your_table",  # 检查表中的行数
        success_callback=lambda: print("数据已准备好！"),
        poke_interval=5,  # 每5秒检查一次
        timeout=60  # 超时时间为60秒
    )
    end_task = DummyOperator(task_id='end_task')

    start_task >> wait_task >> end_task

在这个示例中，SqlSensor会每隔 5 秒查询数据库（通过sql语句），如果查询结果（表的行数）大于 0 或者超时（60 秒）就会根据情况继续执行end_task，并且如果查询成功会调用success_callback函数打印消息。

TimeDeltaSensor

用于等待一个指定的时间间隔过去后再继续执行任务。

from airflow import DAG
from airflow.sensors.time_delta import TimeDeltaSensor
from datetime import datetime, timedelta
from airflow.operators.dummy_operator import DummyOperator

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('time_delta_sensor_dag', default_args=default_args, schedule_interval='@once') as dag:
    start_task = DummyOperator(task_id='start_task')
    wait_task = TimeDeltaSensor(
        task_id='wait_task',
        delta=timedelta(seconds=30)  # 等待30秒
    )
    end_task = DummyOperator(task_id='end_task')

    start_task >> wait_task >> end_task

在这个示例中，TimeDeltaSensor（wait_task）会等待 30 秒后才允许执行end_task。

TimeSensor

等待直到到达一个特定的时间点才继续执行任务。

from airflow import DAG
from airflow.sensors.time_sensor import TimeSensor
from datetime import datetime
from airflow.operators.dummy_operator import DummyOperator

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('time_sensor_dag', default_args=default_args, schedule_interval='@once') as dag:
    start_task = DummyOperator(task_id='start_task')
    wait_task = TimeSensor(
        task_id='wait_task',
        target_time='2024 - 12 - 04T00:01:00'  # 等待直到这个时间点
    )
    end_task = DummyOperator(task_id='end_task')

    start_task >> wait_task >> end_task

这里的TimeSensor会等待直到到达2024 - 12 - 04T00:01:00这个时间点才继续执行后续任务。

FileSensor

用于检查文件系统中是否存在指定的文件。

from airflow import DAG
from airflow.sensors.filesystem import FileSensor
from datetime import datetime
from airflow.operators.dummy_operator import DummyOperator

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('file_sensor_dag', default_args=default_args, schedule_interval='@once') as dag:
    start_task = DummyOperator(task_id='start_task')
    wait_task = FileSensor(
        task_id='wait_task',
        filepath='/path/to/your/file.txt'  # 检查此文件是否存在
    )
    end_task = DummyOperator(task_id='end_task')

    start_task >> wait_task >> end_task

在这个例子中，FileSensor（wait_task）会检查/path/to/your/file.txt这个文件是否存在，只有文件存在后才会执行end_task。

ExternalTaskSensor

用于等待另一个 Airflow 任务（可以在不同的 DAG 中）完成后再执行当前任务。假设我们有两个 DAG，一个是producer_dag，一个是consumer_dag。producer_dag中的任务完成后，consumer_dag中的ExternalTaskSensor才会让后续任务执行。

### producer_dag.py

from airflow import DAG
from datetime import datetime
from airflow.operators.dummy_operator import DummyOperator

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('producer_dag', default_args=default_args, schedule_interval='@once') as dag:
    producer_task = DummyOperator(task_id='producer_task')

### consumer_dag
from airflow import DAG
from airflow.sensors.external_task_sensor import ExternalTaskSensor
from datetime import datetime
from airflow.operators.dummy_operator import DummyOperator

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 12, 4)
}

with DAG('consumer_dag', default_args=default_args, schedule_interval='@once') as dag:
    wait_task = ExternalTaskSensor(
        task_id='wait_task',
        external_dag_id='producer_dag',  # 等待这个DAG中的任务完成
        external_task_id='producer_task'
    )
    consumer_task = DummyOperator(task_id='consumer_task')

    wait_task >> consumer_task

HttpSensor

等待来自HTTP端点的指定响应。当任务依赖于外部API的可用性时，或者当您需要确保API在进行下一个任务之前完成处理时，HttpSensor特别有用。

from datetime import datetime
from airflow import DAG
from airflow.providers.http.sensors.http import HttpSensor

with DAG(dag_id='http_sensor_example', start_date=datetime(2023,8, 1)) as dag:
    wait_for_api = HttpSensor(
        task_id='wait_for_api',
        http_conn id='http_conn',
        endpoint='/api/resource',
        method='GET'
        response_check=lambda response:response.status code == 200
        mode='poke'
        timeout=200
        poke_interval=60,
    )
    
wait_for_api >>task_A >> task B

这段代码整体上是在 Airflow 中创建一个包含 HttpSensor 任务的 DAG，通过传感器来检测特定 HTTP 端点的状态，只有当满足相应条件后才会依次触发后续的其他任务执行。

主要参数说明：

response_check：这是一个用于检查响应是否满足要求的函数（以 lambda 表达式形式给出）。在这里它判断返回的 response 的 status_code 是否等于 200（表示请求成功的常见 HTTP 状态码），如果等于 200 则认为满足继续执行后续任务的条件，否则会继续等待并按照 poke_interval 的间隔再次发起请求。
mode：设置传感器的工作模式，这里是 poke 模式，意味着传感器会按照一定的时间间隔（由 poke_interval 确定）主动去 “探测”（发送请求并检查响应）HTTP 端点的状态，直到满足条件或者达到 timeout 设定的超时时间为止。
timeout：以秒为单位设定整个传感器任务的超时时间，这里设置为 200 秒，即如果经过 200 秒后仍然没有满足 response_check 所设定的条件，任务将会结束（通常可能会导致后续依赖它的任务执行出现异常等情况，具体取决于整个 DAG 的错误处理机制）。
poke_interval：定义了传感器在 poke 模式下每次 “探测”（发送请求）的时间间隔，这里设置为 60 秒，也就是每隔 60 秒向指定的 HTTP 端点发送一次请求并检查响应是否满足条件。

Sensors最佳实践

超时处理：使用on_failure_callback或trigger_rule选项，你可以设置工作流，以便在Sensors超时的情况下重试操作、通知收件人或执行替代任务。
设置超时：为了避免传感器无限期地等待，这可能会导致资源消耗和操作停滞，请为它们配置合理的超时值。
时间戳间隔：为了平衡检查频率和工作流程的整体有效性，为Sensors设置适当的时间戳间隔。
监控Sensors性能：使用内置的Airflow监控功能，如工作持续时间图和甘特图，以保持对Sensors性能的关注。你可以使用它发现瓶颈并改进DAG流程。

Airflow Hooks

Airflow Hooks是可重复使用的部件，使其更容易与其他系统和服务进行通信。它们包括连接管理、身份验证和数据传输等操作，抽象API集成的差异，并提供用于与多个服务通信的标准化接口。Hooks的主要特性包括：

抽象：Hooks通过抽象掉api的复杂性，使与外部系统的连接一致且合乎逻辑成为可能。
可扩展性：为了满足数据管道集成的独特需求，可以简单地自定义或扩展挂钩。
模块化：因为钩子是自包含的和可重用的，所以它使组织代码和消除工作流中的重复变得更容易。

内置Airflow Hooks

Apache Airflow中提供了许多预构建的Hooks，用于与知名系统和服务进行通信，包括消息队列、云平台和数据库。常用的预构建钩子包括：

PostgresHook
BigQueryHook
SlackWebhook
S3Hook
HttpHook

from airflow.providers.postgres.hooks.postgres import PostgresHook
pg_hook = PostgresHook(conn_id='my_postgres_connection')
result = pg_hook.get_records('SELECT * FROM table')

自定义Airflow Hooks

如果预构建的Hooks不能满足你的需求，你可以扩展airflow.hooks.base.BaseHook类。要生成一个唯一的Hook，请执行以下操作：

在BaseHook类中包含任何必要的依赖项。
创建自己的派生自BaseHook的Hook类。
使用必要的技术与外部系统进行连接和通信。

连接MySQL

首先确保已经安装了 mysql-connector-python 库，用于连接 MySQL 数据库。

from airflow.hooks.base_hook import BaseHook
import mysql.connector

class CustomMySQLHook(BaseHook):
    """
    自定义的用于与 MySQL 数据库进行交互的 Hook 类
    """
    def __init__(self, conn_id):
        """
        初始化方法，接收连接 ID 作为参数

        参数:
            conn_id (str): Airflow 中配置的连接的 ID
    """
        super().__init__()
        self.conn_id = conn_id
        self.connection = self.get_connection(self.conn_id)

    def get_conn(self):
        """
        获取与 MySQL 数据库的连接对象

        返回:
            mysql.connector.connection.MySQLConnection: MySQL 数据库连接对象
        """
        # 获取连接配置中的主机地址、端口、用户名、密码、数据库名等信息
        host = self.connection.host
        port = self.connection.port
        username = self.connection.login
        password = self.connection.password
        database = self.connection.schema

        try:
            # 建立与 MySQL 数据库的连接
            conn = mysql.connector.connect(
                host=host,
                port=port,
                user=username,
                password=password,
                database=database
            )
            return conn
        except mysql.connector.Error as e:
            self.log.error(f"无法连接到 MySQL 数据库: {e}")
            raise

    def execute_query(self, query):
        """
        在 MySQL 数据库上执行给定的查询语句

        参数:
            query (str): 要执行的 SQL 查询语句

        返回:
            list: 查询结果，如果是查询数据（如 SELECT 语句），返回结果集的列表，否则返回 None
        """
        conn = self.get_conn()
        cursor = conn.cursor()
        try:
            cursor.execute(query)
            if query.upper().startswith('SELECT'):
                result = cursor.fetchall()
                return result
            else:
                conn.commit()
                return None
        except mysql.connector.Error as e:
            self.log.error(f"执行 SQL 查询失败: {e}")
            raise
        finally:
            cursor.close()
            conn.close()

在 Airflow 中，BaseHook类（自定义 Hook 继承自它）提供了get_connection方法。get_connection方法用于从 Airflow 的元数据库中根据传入的conn_id（连接标识符）来获取对应的连接配置信息，并将其封装成一个Connection对象返回。这个Connection对象包含了诸如主机地址（host）、端口（port）、用户名（login）、密码（password）、数据库名称（如果是数据库连接，对应schema属性）等一系列与具体连接相关的属性，方便后续在自定义 Hook 中基于这些属性去建立和外部系统的实际连接（例如像在之前示例中，利用这些信息去创建与 API 的会话对象或者与数据库的连接对象等）。

使用自定义hooks

假设已经在 Airflow 的连接配置中设置好了对应的 MySQL 连接（比如连接 ID 为 my_mysql_connection），可以在 DAG 任务中这样使用这个自定义 Hook：

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime

def query_mysql_database():
    """
    一个简单的 Python 函数，使用自定义 Hook 对 MySQL 数据库进行查询操作
    """
    hook = CustomMySQLHook(conn_id='my_mysql_connection')
    query = "SELECT * FROM some_table"
    result = hook.execute_query(query)
    # 这里可以对查询结果进行进一步处理，比如打印结果等
    print(result)

# 定义 DAG
with DAG('mysql_query_dag', start_date=datetime(2024, 12, 4), schedule_interval='@daily') as dag:
    task = PythonOperator(
        task_id='mysql_query_task',
        python_callable=query_mysql_database,
        dag=dag
    )

总结

总之，Apache Airflow的Sensors, Operators & Hooks作为数据编排的核心组件，使数据团队能够构建可靠、可扩展和可适应的数据管道。通过有效地组合这些组件并遵循最佳实践，组织可以在保持代码简单性、可维护性和效率的同时，保障数据管道正确且高效运行。随着数据环境的不断演变，掌握这些基本构建块无疑将成为任何数据工程或数据科学项目的关键资产。

您可能感兴趣的与本文相关的镜像