Кастомные email-оповещения в Apache Airflow

Пролог

Каждый разработчик сталкивался, или непременно столкнется, с ситуацией, когда в бизнес-требованиях видишь “этакое желание”.

Этакое желание (каламбурное определение) – достичь чего-то невозможного или близкого к невозможному с помощью программирования.

Дорогой читатель, наверное, задаст риторический вопрос:
— Как это?! Разве чего-то нельзя достичь с помощью программирования?!
Вопрос, конечно, к месту, и ответ в большинстве случаев очевиден:
— Нет ничего невозможного… главное грамотно спроектировать техническое решение.
Но сейчас немного не об этом, а о целесообразности расхода ресурсов: трудозатрат, количество привлеченных специалистов и т.п.

Стоит ли овчинка выделки?

Увидев “этакое желание” хочется сразу пресечь его на корню и больше не возвращаться к обсуждению безумной идеи.
Но что если все-таки дать шанс и извлечь из этого выгоду в будущем?

Конкретный случай

Спойлер: стоит отметить, что дальше не будет rocket science и метапрограммирования, а будет описан подход к решению задачи.

Мне, как разработчику команды “Разработка и автоматизация загрузок данных” Газпромбанка, посчастливилось встретить в одном из технических заданий “этакое желание” — реализовать в загрузочном DAG кастомизированные email-оповещения.
Они должны собирать необходимую информацию на всех стадиях пайплайна, то есть в каждой из задач DAG’a, а в зависимости от успешного выполнения или при отклонении алгоритма, формировать тело email-сообщения. Например: в задаче происходит ошибка валидации данных и в этом случае должна сообщаться конкретная причина падения DAG. Информация должна быть полезной и понятной для бизнеса.

Для загрузок в нашей кодовой базе уже был реализован инструмент оповещения,
но оповещения выполнялись только при удачном сценарии отработки DAG, например,
с количеством обработанных файлов и записанных строк в базу данных.
Формат был стандартным и не конфигурировался.

Первое впечатление от увиденного требования у меня было крайне скептическое,
но на одной из ежедневных встреч с командой, обсудив все за и против, было принято решение реализовать данную задачу. Один из основных аргументов – создание универсального инструмента нотификации, который можно будет переиспользовать в других наших DAG’ах. Это переиспользование даст возможность получить другим Заказчикам более четкое представление о загрузках.
Чем подробнее алерты, тем меньше вопросов к нам.

Ближе к коду

Собрав команду на встрече “побрейнштормить” задачку, мы наметили итеративный подход к реализации: в первой итерации – основной механизм, во второй – доработка и исправление багов.

Умный и находчивый читатель, наверное, сразу догадался, что можно использовать встроенные в библиотеку Airflow Callbacks. Но это не наш случай, так как нужно обрабатывать ошибки по месту, конкретные случаи. Наиболее предпочтительный вариант — использовать  отдельный Task Airflow (далее dag_email_notification) с определенным TriggerRule = NONE_SKIPPED, так как необходимо обработать все задачи DAG’а со статусами success и failed.

С обработчиком определились, а как передавать информацию между задачами? – XCom, скажете вы, и будете правы. Но что и как передавать? Хотелось бы использовать лаконичный и понятный для эксплуатации объект, который сам бы отправился в XCom.
Мы решили, что для этого хорошо подходит — dataclass (далее EmailNotification) с несколькими методами. Dataclass хороший вариант, так как в XCom следует отправлять сериализуемые в JSON данные.

Рассмотрим, как выглядит наш объект:

@dataclass class EmailNotification:     task_description: str | None = None     lines: list[dict[str, t.Any] | str | BaseException] = field(default_factory=list)       XCOM_KEY: str = "email_notification"       @classmethod     def set_task_description(cls, task_description: str) -> None:         """         Добавление описания таски в тело сообщения.            :param task_description: Описание таски - дополнительно           раскрывает смысл/задачу.           :return: None           """         instance = cls.get_from_xcom()           instance.task_description = task_description           cls.push_into_xcom(instance)       @classmethod     def add_line(         cls,         message: str | BaseException | dict[str, t.Any],     ) -> None:         """         Добавление дополнительной строки в инстанс EmailNotification          и пуш в XCom.           :param message: Содержание строки с сообщением         :return: None           """         instance = cls.get_from_xcom()           instance.lines.append(message)           cls.push_into_xcom(instance)       @classmethod     def push_into_xcom(cls, instance: EmailNotification) -> None:         # util,который из context возвращает текущий TaskInstance         ti: TaskInstance = get_current_ti()           ti.xcom_push(key=cls.XCOM_KEY, value=instance)           @classmethod     def get_from_xcom(cls) -> EmailNotification:         """         Возвращает инстанс EmailNotification из XCOM или создает новый.           :return: EmailNotification           """         # util, который из context возвращает текущий TaskInstance         ti: TaskInstance = get_current_ti()           instance = ti.xcom_pull(key=cls.XCOM_KEY, task_ids=ti.task_id)           return instance or cls()

Итак, на борту имеется — основной метод add_line, который будет добавлять необходимую строку в тело сообщения, как мы и хотели. Classmethod выбран не случайным образом, а для лаконичного использования, исключается необходимость дополнительно создавать инстанс, инициализируя класс.
Основная работа скрыта в методе get_from_xcom: при добавлении строки в тело сообщения инстанс класса будет создан в XCom, а если объект уже есть, то добавится строка в существующий инстанс.

Как это выглядит на практике:

@task def first_task() -> None:     EmailNotification.set_task_description("Описание для task 1")       EmailNotification.add_line("Первая строка в task 1")       EmailNotification.add_line("Вторая строка в task 1")

По результатам выполнения задачи first_task, в XCom получаем следующий объект:

EmailNotification(     task_description="Описание для task 1",     lines=["Первая строка в task 1", "Вторая строка в task 1"],     XCOM_KEY="email_notification" )

Остается только получить все эти объекты из XCom задач со статусами success/failed с помощью обработчика dag_email_notification, собрать все это в одно тело сообщения, разделяя на блоки для наглядности, и разослать всем заинтересованным лицам.

Ну мёд, медятина!

Вот как выглядит наше сообщение:

Конечно, для красоты добавлены всякие “рюшечки” и “бантики”, например, если статус failed, фон блока задачи в HTML-разметке меняется на красный.
Автоматически определяются адресаты и тема сообщения в зависимости от статуса выполнения DAG.
Параметр message в методе add_line неспроста может быть: строкой, exception или словарем — обработчик dag_email_notification учитывает это и, в зависимости от типа данных, по-разному формируется содержимое блока с описанием процесса выполнения задачи в теле сообщения.

И что, всё?

Как сказал Тимлид:

“Попробуй протестить с Dynamic Task Mapping (далее Mapped Tasks), там могут быть нюансы…”.

И действительно, нюансы есть…

Первое, что идет не по плану — в EmailNotification.get_from_xcom после успешного выполнения первой Mapped Task остальные падают, и ничего непонятно.
В ходе деббагинга выявился фигурант дела — LazyXComAccess.

Кратко говоря, это последовательность XCom’ов с Mapped Tasks.
В get_from_xcom, чтобы получить инстанс класса из XCom, одного параметра task_id становится мало – у всех Mapped Tasks один и тот же task_id.
К счастью, XCom’ы внутри LazyXComAccess хранятся с использованием индексов.
Чтобы решить вопрос, добавим map_indexes=ti.map_index в вызов xcom_pull.
Также теперь стоит учитывать индексы при получении данных из XCom в обработчике dag_email_notification.

С Mapped Tasks сообщение может выглядеть так:

Напоследок

Хотелось бы показать больше кода, особенно обработчик dag_email_notification,
но цель статьи – не перегружать вас техническими деталями, а главное – продемонстрировать, как можно совместно с командой предоставить бизнесу эффективный инструмент для мониторинга и анализа процессов.