如何通过 Apache Airflow 将数据导入 Elasticsearch

作者：来自 Elastic Andre Luiz

了解如何通过 Apache Airflow 将数据导入 Elasticsearch。

Apache Airflow

Apache Airflow 是一个旨在创建、安排（schedule）和监控工作流的平台。它用于编排 ETL（Extract-Transform-Load）流程、数据管道和其他复杂工作流，提供灵活性和可扩展性。它的可视化界面和实时监控功能使管道管理更易于访问和高效，让你可以跟踪执行的进度和结果。以下是它的四个主要支柱：

动态：管道以 Python 定义，允许动态灵活地生成工作流。
可扩展：Airflow 可以与各种环境集成，可以创建自定义运算符，并可以根据需要执行特定代码。
优雅：管道以干净明确的方式编写。
可扩展：其模块化架构使用消息队列来编排任意数量的工作器。

在实践中，Airflow 可用于以下场景：

数据导入：编排将数据每日提取到 Elasticsearch 等数据库中。
日志监控：管理日志文件的收集和处理，然后在 Elasticsearch 中进行分析以识别错误或异常。
多种数据源集成：将来自不同系统（API、数据库、文件）的信息合并到 Elasticsearch 中的单个层中，简化搜索和报告。

DAG：Directed Acyclic Graphs - 有向无环图

在 Airflow 中，工作流由 DAG（有向无环图）表示。DAG 是一种定义任务执行顺序的结构。DAG 的主要特征是：

由独立任务组成：每个任务代表一个工作单元，旨在独立执行。
排序：任务的执行顺序在 DAG 中明确定义。
可重用性：DAG 旨在重复执行，促进流程自动化。

Airflow 的主要组件

Airflow 生态系统由多个组件组成，它们共同协作以协调任务：

调度程序 - scheduler：负责调度 DAG 并发送任务以供工作人员执行。
执行器 - Exectutor：管理任务的执行，将其委托给工作人员。
Web 服务器 - Webserver：提供与 DAG 和任务交互的图形界面。
Dags 文件夹 - Dags folder：我们存储用 Python 编写的 DAG 的文件夹。
元数据 - Metadata：作为工具存储库的数据库，由调度程序和执行器用于存储执行状态。

Apache Airflow 和 Elasticsearch

我们将演示如何使用 Apache Airflow 和 Elasticsearch 来协调任务并在 Elasticsearch 中索引结果。此演示的目标是创建一个任务管道来更新 Elasticsearch 索引中的记录。此索引包含电影数据库，用户可以在其中进行评分和分配评级。想象一个每天有数百个评级的场景，有必要保持评级记录更新。为此，将开发一个 DAG，它将每天执行，负责检索新的合并评级并更新索引中的记录。

在 DAG 流程中，我们将有一个获取评级的任务，然后是一个验证结果的任务。如果数据不存在，DAG 将被定向到失败任务。否则，数据将在 Elasticsearch 中编入索引。目标是通过一种带有负责计算分数的机制的方法检索评级，以更新索引中电影的评级字段。

使用 Apache Airflow 和 Elasticsearch 以及 Docker

要创建容器化环境，我们将使用 Apache Airflow 和 Docker。按照 “在 Docker 中运行 Airflow” 指南中的说明实际设置 Airflow。

至于 Elasticsearch，我将使用 Elastic Cloud 上的集群，但如果你愿意，也可以使用 Docker 配置 Elasticsearch。已经创建了一个包含电影目录的索引，其中电影数据已编入索引。这些电影的 “rating” 字段将被更新。

创建 DAG

通过 Docker 安装后，将创建一个文件夹结构，其中包括 dags 文件夹，我们必须将 DAG 文件放在该文件夹中，以便 Airflow 识别它们。

在此之前，我们需要确保安装了必要的依赖项。以下是此项目的依赖项：

pip install apache-airflow apache-airflow-providers-elasticsearch

我们将创建文件 update_ratings_movies.py 并开始编写任务代码。

现在，让我们导入必要的库：

from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.providers.elasticsearch.hooks.elasticsearch import ElasticsearchPythonHook

我们将使用 ElasticsearchPythonHook，这是一个通过抽象连接和使用外部 API 来简化 Airflow 和 Elasticsearch 集群之间集成的组件。

接下来，我们定义 DAG，并指定其主要参数：

dag_id：DAG 的名称。
start_date：DAG 的启动时间。
schedule：定义周期（在我们的例子中是每日）。
doc_md：将导入并显示在 Airflow 界面中的文档。

定义任务

现在，让我们定义 DAG 的任务。第一个任务将负责检索电影评级数据。我们将使用 PythonOperator，并将 task_id 设置为“get_movie_ratings”。python_callable 参数将调用负责获取 ratings 的函数。

get_ratings_operator = PythonOperator(
   task_id='get_movie_ratings',
   python_callable=get_movie_ratings_task
)

接下来，我们需要验证结果是否有效。为此，我们将使用带有 BranchPythonOperator 的条件。task_id 将为 “validate_result”，python_callable 将调用验证函数。op_args 参数将用于将上一个任务 “get_movie_ratings” 的结果传递给验证函数。

validate_result = BranchPythonOperator(
   task_id='validate_result',
   python_callable=validate_result,
   op_args=["{
  
  { task_instance.xcom_pull(task_ids='get_movie_ratings') }}"]
)

如果验证成功，我们将从 “get_movie_ratings” 任务中获取数据并将其索引到 Elasticsearch 中。为此，我们将创建一个新任务 “index_movie_ratings”，它将使用 PythonOperator。op_args 参数将 “get_movie_ratings” 任务的结果传递给索引函数。

index_ratings_operator = PythonOperator(
   task_id='index_movie_ratings',
   python_callable=index_movie_ratings_task,
   op_args=["{
  
  { task_instance.xcom_pull(task_ids='get_movie_ratings') }}"]
)

如果验证表明失败，DAG 将继续执行失败通知任务。在此示例中，我们只是打印一条消息，但在实际场景中，我们可以配置警报来通知失败。

failed_get_rating_operator = PythonOperator(
   task_id='failed_get_rating_operator',
   python_callable=lambda: print('Ratings were False, skipping indexing.')
)

最后，我们定义任务依赖关系，确保它们以正确的顺序执行：

get_ratings_operator >> validate_result >> [index_ratings_operator, failed_get_rating_operator]

以下是我们 DAG 的完整代码：

"""
DAG update Rating Movies
"""
import ast
import random

from airflow import DAG
from datetime import datetime

from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.providers.elasticsearch.hooks.elasticsearch import ElasticsearchPythonHook


def index_movie_ratings_task(movies):
   es_hook = ElasticsearchPythonHook(hosts=None,
                                     es_conn_args={
                                         "cloud_id": "cloud_id"
                                         "api_key": "api-key"
                                     })
   es_client = es_hook.get_conn
   actions = []
   for movie in ast.literal_eval(movies):
       actions.append(
           {
               "update": {
                   "_id": movie["id"],
                   "_index": "movies"
               }
           }
       )
       actions.append(
           {
               "doc": {
                   "rating": movie["rating"]
               },
               "doc_as_upsert": True
           }
       )
   result = es_client.bulk(operations=actions)
   print(f"Ingestion completed.")
   print(result)
   return True


def get_movie_ratings_task():
   movies = [
       {"id": i, "rating": round(random.uniform(1, 10), 1)}
       for i in range(1, 100)
   ]
   return movies

def validate_result(result):
   if not result:
       return 'failed_get_rating_operator'
   else:
       return 'index_movie_ratings'


with DAG(
       dag_id="update_ratings_movies_2024",
       start_date=datetime(2024, 12, 29),
       schedule="@daily",
       doc_md=__doc__,
):
   get_ratings_operator = PythonOperator(
       task_id='get_movie_ratings',
       python_callable=get_movie_ratings_task
   )

   validate_result = BranchPythonOperator(
       task_id='validate_result',
       python_callable=validate_result,
       op_args=["{
  
  { task_instance.xcom_pull(task_ids='get_movie_ratings') }}"],
       provide_context=True
   )

   index_ratings_operator = PythonOperator(
       task_id='index_movie_ratings',
       python_callable=index_movie_ratings_task,
       op_args=["{
  
  { task_instance.xcom_pull(task_ids='get_movie_ratings') }}"]
   )

   failed_get_rating_operator = PythonOperator(
       task_id='failed_get_rating_operator',
       python_callable=lambda: print('Ratings were False, skipping indexing.')
   )

get_ratings_operator >> validate_result >> [index_ratings_operator, failed_get_rating_operator]