[Airflow] Telegram 알림 operator 만들기 (requests 사용)

버전정보

Airflow 2.5.3

Python 3.8

 

왜인지는 모르겠지만 2.5.3 버전을 사용하고 있는 환경에서는 telegram provider 설치는 되는데 connection type 리스트에는 안떴다 ㅠ (버전호환성도 다 맞춤)

2.6.3 버전으로 올리고 provider 재설치 후 다시 확인해보니 operator가 잘 동작하는 건 확인했는데.. 어쨌든 2.5.3 버전에서 사용하기 위해 request 라이브러리를 이용해서 post 요청을 보내는 방식으로 Webhook을 만들었다.

 

1. Telegram Connection 생성

이건 그냥 hook에서 connection 정보 어떻게 가져올지 상황에 맞게 설정하면 될 듯 하다.

나같은 경우는 그냥 http connection type을 사용했다.

airflow connection 추가

`Host`: https://api.telegram.org/ (텔레그램 api url 입력)

`Login`: chat_id 입력

`Password`: token값 입력. 왠지 토큰값은 암호화해줘야할 것 같아서 암호화된 값을 저장할 수 있는 password 란에 넣어줬다.

 

 

2. TelegramhookHook

import requests
from airflow.providers.http.hook.http import HttpHook

class TelegramhookHook(HttpHook):
	def __init__(
            self,
            http_conn_id = None,
            message = '',
            *args,
            **kwargs,
    	):
            super().__init__(http_conn_id=http_conn_id, *args, **kwargs)
            self.conn = self.get_connection(http_conn_id)
            self.telegram_url = self.conn.host
            self.message = message
            self.chat_id = self.conn.login
            self.token = self.conn.password

	def _build_telegram_data(self) -> str:
            return {
                        "chat_id": self.chat_id,
                        "text": self.message
            }

    def _build_full_url(self) -> str:
        if self.telegram_url[-1] == '/':
            return self.telegram_url + 'bot' + self.token + '/sendMessage'
        else:
            return self.telegram_url + '/bot' + self.token + '/sendMessage'


	def execute(self) -> None:
            telegram_message = self._build_telegram_data()
            url = self._build_full_url()

            push = requests.post(
                url = url,
                data = telegram_message
            )

 

HttpHook을 상속받음.

여러가지 값 설정 후에 `execute`에서 `requests.post()`로 요청하는 부분이다.

1에서 설정했다시피 conn에서 host, login, password 값을 가져와서 _build_full_url 에서 요청 url을 만든다.

근데 텔레그램이면 어차피 host값은 항상 같을거라서.. 그냥 귀찮게 매번 작성하지 않고 init에 추가해줘도 될 것 같다.

 

3. TelegramWebhookOperator

from typing import Any, Dict, Optional
from airflow.providers.http.hooks.http import HttpHook
from airflow.providers.http.operators.http import SimpleHttpOperator
from airflow.utils.decorators import apply_defaults
from TelegramhookHook import TelegramhookHook

class TelegramWebhookOperator(SimpleHttpOperator):
	@apply_defaults
    def __init__(
    	self,
        *,
        http_conn_id: str,
        message: str = '',
        **kwargs,
    ) -> None:
    	super().__init__(**kwargs)
        self.http_conn_id = http_conn_id
        self.message = message
        self.hook: Optional[TelegramhookHook] = None

	def execute(self, context: Dict[str, Any]) -> None:
    	self.hook = TelegramhookHook(
        	self.http_conn_id,
            self.message
        )
        self.hook.execute()

SimpleHttpOperator 을 상속받음.

http_conn_id, message 값을 받아서 TelegramhookHook에 전해주고 실행한다.

 

그냥 Task로 텔레그램 보내려면 여기서 TelegramWebhookOperator 사용하면 된다.

 

4. TelegramFailedNotification

텔레그램알림의 주요 사용처는 task 실패 시의 알람이었다.

그래서 아래와 같이 on_failure 함수를 만들었다.

`message_content` 부분에 필요한 task_instance 정보를 넣고

어차피 실패알림은 동일한 봇/동일한 채팅방으로 보낼거라 `telegram_conn`은 1에서 설정했던 커넥션id 그대로 넣어놨다.

추가로 airflow 내부적으로는 utc로 계산되기 때문에 execution_date도 utc timezone으로 리턴되었다.

한국시간과 맞추기 위해 9시간을 더해줬다.

from datetime import timedelta
from TelegramWebhookOperator import TelegramWebhookOperator

def on_failure(context):
	task_instance = context.get('task_instance')
    dag_id = task_instance.dag_id
    task_id = task_instance.task_id
    log_id = task_instance.run_id
    
    kst_execute_date = (context.get('execution_date') + timedelta(hours=9)).strftime('%Y%m%d %H%M%S')
    kst_next_execute_date = (context.get('next_execution_date') + timedelta(hours=9)).strftime('%Y%m%d %H%M%S')
    
    message_content = f"
    	필요한거 작성....
        ex)
		dag_id = {dag_id}
        task_id = {task_id}
        log_id = {log_id}
        실행시간 = {kst_execute_date}
        다음실행시간 = {kst_next_execute_date}
    "
    
    telegram_conn = "telegram_alert_conn_id"
    
    alert = TelegramWebhookOperator(
    	task_id = "AirflowFailedTelegramAlert",
        http_conn_id = telegram_conn
        message = message_content
    )
    
    return alert.execute(context=context)

 

 

참고: https://dydwnsekd.tistory.com/69