Написание ETL пайплайна при помощи airflow, rabbitmq и postgres

Всем привет! В данной статье мы рассмотрим как можно локально развернуть airflow с помощью docker desktop’a и PyCharm’a. Кроме него развернём ещё и следующее: rabbitMQ, postgreSQL, redis и MongoDB.

Первый делом, нужно установить себе на компьютер Docker desktop, PyCharm и PgAdmin4. Теперь нам нужен yaml файл, в котором будет прописано все образы, которые мы хотим поставить. Готовый файл для нашей задачи можете скачать по ссылке. Когда всё необходимое было скачано, можно приступить к развёртыванию airflow.

Создаём проект в PyCharm, и в корневую папку загружаем наш yaml файл. В терминале выполняем следующий код: docker compose up -d. Когда нажмёте enter, то у вас начнётся скачивание необходимых образов в docker desktop и сразу из них запустятся контейнеры, кроме этого, в проекте у вас создастся необходимая структура папок.

Когда всё скачается, у вас docker desktop должен выглядеть примерно так:

На данный момент, у нас уже локально развернулось всё, что я обещал, теперь будем устанавливать необходимые соединения.

Объединение airflow и postgres

Начнём с Postgres; зайдём в PgAdmin и подключимся к нашему серверу

Логин и пароль: airflow, airflow соответственно

После подключения у вас должно появиться в сервере две базы данных:

Для работы с airflow используется airflow. Для будущего написания ETL давайте сразу создадим таблицу:

Теперь займёмся установкой соединения между airflow и postgres. Для того, чтобы перейти в airflow, зайдём в docker desktop и нажмём на порт у airflow-webserver:

Нужно будет ввести логин и пароль: airflow и airflow

После этого, вы окажетесь на основной странице airflow:

Теперь, чтобы наш airflow мог взаимодействовать с postgres, необходимо создать соединение между ними в airflow:

Пароль указываем от нашей базы данных, в нашем случае airflow

Сохраняем соединение и теперь у нас налажена работа airflow и postgres.

Напишем самый простой dag, чтобы проверить, что всё работает. Для начала нужно установить дополнительные библиотеки в PyCharm: apache-airflow и apache-airflow-providers-postgres.

Напишем код:

from airflow import DAG from airflow.operators.python import PythonOperator from airflow.providers.postgres.hooks.postgres import PostgresHook   hook = PostgresHook(postgres_conn_id='Postgres')  request = "insert into public.data (name, birth_date, age) values ('Anton', '2004.11.14', 19)"   def insert():     hook.run(request)   with DAG(     dag_id="base_dag",     default_args={         "owner": "Kirill",     } ) as dag:      t = PythonOperator(         task_id='insert_postgres',         python_callable=insert         )

Как только мы закончим его писать, в airflow появится наш dag:

Запустим его и посмотрим, что получилось.

Как видим, dag отработал успешно и добавил новую строку в нашу таблицу.

Объединение airflow и rabbitmq

Начнём с настройки rabbitMQ. Для того, чтобы перейти в rabbitMQ, нужно нажать на его порт в docker desktop’е:

Логин и пароль: guest, guest. Перейдём в раздел «Queues» и создадим там новую очередь с названием queue.airflow:

Вернёмся в PyCharm и установим библиотеку airflow-provider-rabbitmq.

Перейдём в airflow в раздел с соединениями и добавим соединение с брокером, в поле с хостом нужно указать имя нашего контейнера с rabbit, то есть, rabbitmq_main:

После всего этого, настройка брокера закончена и мы можем приступать к взаимодействию питона, airflow и брокера. Напишем базовую программу:

Для этого в папке dags создадим папку tasks_broker. В ней создадим sensor.py с следующим кодом:

from rabbitmq_provider.sensors.rabbitmq import RabbitMQSensor   sensor = RabbitMQSensor(     task_id="sensor",     queue_name="queue.airflow",     rabbitmq_conn_id="RabbitMQ", ) 

Данный код является таской, которая принимает данные из брокера и передаёт в другую таску.

Также в папке tasks_broker создадим get_data.py с следующим содержанием:

import json from airflow.decorators import task   @task(task_id="get_data") def get_data(**kwargs) -> None:     message = json.loads(kwargs['task_instance'].xcom_pull(task_ids='sensor'))      print("#########################################################################################")     print(message)     print("#########################################################################################")

После этого, в папке dags создадим python_broker.py:

from airflow import DAG from tasks_broker.sensor import sensor from tasks_broker.get_data import get_data   with DAG(         dag_id="python_broker",         schedule_interval=None, ) as dag:     sensor >> get_data()

И у нас в airflow появится новый dag python_broker. Запустим его и он будет ждать до тех пор, пока из брокера не поступит сообщение. Чтобы отправить сообщение из брокера, нужно перейти в раздел «queues» и выбрать необходимую очередь, в нашем случае queue.airflow. Там выбрать «publish message», ввести нужное сообщение и отправить.

Как мы отправим сообщение, наш dag выполнится и мы можем убедиться в том, что всё прошло успешно, зайдя в логи второй таски:

Решение задачи ETL. Объединение airflow, rabbitmq и postgres

Давайте теперь решим задачу ETL, будет три таски:

1. Получаем данные о человеке в формате json из брокера (имя, пол, дату рождения, возраст, город проживания) и передаём его во вторую таску для обработки

2. Вторая таска отвечает за обработку полученных полей и дальнейшую передачу результата в третью таску (будем из пяти полей передавать только три (имя, дату рождения и возраст))

3. Третья таска будет загружать данные в postgres

По сути, все соединения у нас установлены, необходимость лишь заключается в правильно написанном коде. В папке dags создадим папку tasks_broker_postgres. Там создадим три файла extract.py, transformation.py и load.py.

extract.py:

from rabbitmq_provider.sensors.rabbitmq import RabbitMQSensor   extract = RabbitMQSensor(     task_id="extract",     queue_name="queue.airflow",     rabbitmq_conn_id="RabbitMQ", ) 

transformation.py:

import json from airflow.decorators import task   @task(task_id="transformation") def transformation(**kwargs) -> dict:     message = json.loads(kwargs['task_instance'].xcom_pull(task_ids='extract'))      new_message = {"name": message["name"],                    "birth_date": message["birth_date"],                    "age": message["age"]}      return new_message 

load.py:

from airflow.decorators import task from airflow.providers.postgres.hooks.postgres import PostgresHook   @task(task_id="load") def load(**kwargs) -> None:     message = kwargs['task_instance'].xcom_pull(task_ids='transformation')      hook = PostgresHook(postgres_conn_id='Postgres')      request = ("insert into public.data (name, birth_date, age) values ('" + message["name"] + "','" +                                                                             message["birth_date"] + "'," +                                                                             str(message["age"]) + ")")      hook.run(request) 

Теперь в папке dags создадим python_broker_postgres.py:

from airflow import DAG from tasks_broker_postgres.extract import extract from tasks_broker_postgres.transformation import transformation from tasks_broker_postgres.load import load   with DAG(         dag_id="python_broker_postgres",         schedule_interval=None, ) as dag:     extract >> transformation() >> load() 

Запустим dag в airflow, зайдём в брокер и отправим сообщение:

По логам смотрим, что все таски прошли успешно:

Зайдём теперь в базу данных:

Как видим, наше сообщение успешно обработалось и добавилось в базу данных.

Заключение

В данной статье мы рассмотрели взаимодействие python, airflow, postgres и rabbitmq.

Если вы не забыли, то в yaml файл входит ещё и MongoDB, работа с ней осуществляется аналогично работе с postgres. Для удобной работы с MongoDB можно скачать эту программу.

До новых публикаций!