refactor: rewrite discord message sending through httphook

Author: Lee-W

dags/app/discord.py

@@ -1,21 +1,24 @@
 import logging
-from datetime import datetime, timedelta
+from datetime import datetime
 
 import requests
-from airflow.sdk import Asset, AssetWatcher, Context, dag, task
+import tenacity
+from airflow.providers.http.hooks.http import HttpHook
+from airflow.sdk import Asset, AssetWatcher, Context, Variable, dag, task
 
 from triggers.finance_report import FinanceReportTrigger
 
-logger = logging.getLogger(__name__)
+# get the airflow.task logger
+task_logger = logging.getLogger("airflow.task")
+
 
 finance_report_asset = Asset(
     name="finance_report",
     watchers=[
         AssetWatcher(
             name="finance_report_watcher",
             trigger=FinanceReportTrigger(
-                # poke_interval=86400,  # 60*60*24
-                poke_interval=5,  # 60*60*24
+                poke_interval=86400,  # 60*60*24
             ),
         )
     ],
@@ -37,31 +40,44 @@
     },
 )
 def discord_message_notification():
-    """Send Discord Message"""
+    """Send Discord Message."""
 
-    @task(
-        retries=10,
-        retry_delay=timedelta(seconds=10),
-    )
+    @task
     def send_discord_message(**context: Context) -> None:
         triggering_asset_events = context["triggering_asset_events"]
-        session = requests.session()
-        logger.info(f"Receive asset events {triggering_asset_events}")
         for asset_uri, asset_events in triggering_asset_events.items():
-            logger.info(f"Receive asset event from Asset uri={asset_uri}")
+            task_logger.info(f"Receive asset event from Asset uri={asset_uri}")
+
+            http_hook = HttpHook(method="POST", http_conn_id="discord_webhook")
             for asset_event in asset_events:  # type: ignore[attr-defined]
                 if asset_event.extra.get("from_trigger", False):
                     details = asset_event.extra["payload"]
                 else:
                     details = asset_event.extra
 
-                session.post(
-                    details.get("webhook_url"),
-                    json={
+                if not details:
+                    task_logger.error(
+                        f"Detail {details} cannot be empty. It's required to send discord message."
+                    )
+                    continue
+
+                task_logger.info("Start sending discord message")
+                endpoint = Variable.get(details.get("webhook_endpoint_key"))
+                http_hook.run_with_advanced_retry(
+                    endpoint=endpoint,
+                    data={
                         "username": details.get("username"),
                         "content": details.get("content"),
                     },
+                    _retry_args=dict(
+                        wait=tenacity.wait_random(min=1, max=10),
+                        stop=tenacity.stop_after_attempt(10),
+                        retry=tenacity.retry_if_exception_type(
+                            requests.exceptions.ConnectionError
+                        ),
+                    ),
                 )
+                task_logger.info("Discord message sent")
 
     send_discord_message()
 

dags/app/proposal_reminder/dag.py

@@ -2,7 +2,7 @@
 Send Proposal Summary to Discord
 """
 
-from airflow.sdk import Metadata, Variable, asset
+from airflow.sdk import Metadata, asset
 from app.proposal_reminder.udf import get_proposal_summary
 
 # DEFAULT_ARGS = {
@@ -16,19 +16,17 @@
 
 
 @asset(
-    name="proposal_count",
-    dag_id="proposal_count",
     schedule="0 16 * * *",  # At 16:00 (00:00 +8)
 )
-def DISCORD_PROPOSAL_REMINDER_v3(self):
+def proposal_count(self):
     summary = get_proposal_summary()
     n_talk = summary["num_proposed_talk"]
     n_tutorial = summary["num_proposed_tutorial"]
 
     yield Metadata(
         self,
         extra={
-            "webhook_url": Variable.get("DISCORD_PROGRAM_REMINDER_WEBHOOK"),
+            "webhook_endpoint_key": "DISCORD_PROGRAM_REMINDER_WEBHOOK",
             "username": "Program talk reminder",
             "content": f"目前投稿議程數: {n_talk}; 課程數: {n_tutorial}",
         },

dags/app/team_registration_bot/dag.py

@@ -2,10 +2,10 @@
 Send daily ordering metrics to discord channel
 """
 
-from airflow.sdk import Metadata, Variable, asset
+from airflow.sdk import Metadata, asset
 from app.team_registration_bot.udf import (
-    _compose_discord_msg,
-    _get_statistics_from_bigquery,
+    compose_discord_msg,
+    get_statistics_from_bigquery,
 )
 
 # DEFAULT_ARGS = {
@@ -18,18 +18,15 @@
 # }
 
 
-@asset(
-    name="registration_statistics",
-    dag_id="registration_statistics",
-    schedule="@daily",
-)
-def KKTIX_DISCORD_BOT_FOR_TEAM_REGISTRATION(self):
-    statistics = _get_statistics_from_bigquery()
+@asset(schedule="@daily")
+def registration_statistics(self):
+    # KKTIX_DISCORD_BOT_FOR_TEAM_REGISTRATION
+    statistics = get_statistics_from_bigquery()
     yield Metadata(
         self,
         extra={
-            "webhook_url": Variable.get("discord_webhook_registration_endpoint"),
+            "webhook_endpoint_key": "discord_webhook_registration_endpoint",
             "username": "KKTIX order report",
-            "content": _compose_discord_msg(statistics),
+            "content": compose_discord_msg(statistics),
         },
     )

dags/app/team_registration_bot/udf.py

@@ -11,7 +11,7 @@
 CLIENT = bigquery.Client(project=os.getenv("BIGQUERY_PROJECT"))
 
 
-def _get_statistics_from_bigquery() -> RowIterator:
+def get_statistics_from_bigquery() -> RowIterator:
     query_job = CLIENT.query(
         f"""
         WITH UNIQUE_RECORDS AS (
@@ -55,10 +55,10 @@ def _get_statistics_from_bigquery() -> RowIterator:
 }
 
 
-def _compose_discord_msg(payload) -> str:
+def compose_discord_msg(payload) -> str:
     msg = (
         f"Hi 這是今天 {datetime.now().date()} 的票種統計資料，"
-        "售票期結束後，請 follow README 的 `gcloud` 指令進去把 Airflow DAG 關掉\n\n"
+        "售票期結束後，請 follow README 的 `gcloud` 指令進去把 Airflow dag 關掉\n\n"
     )
     total = 0
     total_income = 0

triggers/finance_report.py

@@ -125,7 +125,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                 msg = format_diff_df_as_message(df_diff)
                 yield TriggerEvent(
                     {
-                        "webhook_url": Variable.get("discord_data_stratagy_webhook"),
+                        "webhook_endpoint_key": "discord_data_stratagy_webhook",
                         "username": "財務機器人",
                         "content": msg,
                     }