feat(discord): move all discord message sending to one dag and use asset to trigger

Author: Lee-W

dags/app/discord.py

@@ -1,18 +1,69 @@
+import logging
+from datetime import datetime, timedelta
+
 import requests
-import tenacity
+from airflow.sdk import Asset, AssetWatcher, Context, dag, task
+
+from triggers.finance_report import FinanceReportTrigger
 
-session = requests.session()
+logger = logging.getLogger(__name__)
 
-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),
+finance_report_asset = Asset(
+    name="finance_report",
+    watchers=[
+        AssetWatcher(
+            name="finance_report_watcher",
+            trigger=FinanceReportTrigger(
+                poke_interval=5,  # 60*60*24
+            ),
+        )
+    ],
 )
 
 
-@tenacity.retry(**RETRY_ARGS)  # type: ignore[call-overload]
-def send_webhook_message(webhook_url: str, username: str, msg: str) -> None:
-    session.post(
-        webhook_url,
-        json={"username": username, "content": msg},
+@dag(
+    schedule=(
+        finance_report_asset
+        | Asset.ref(name="CFP_summary")
+        | Asset.ref(name="kktix_order_report")
+    ),
+    start_date=datetime(2025, 6, 28),
+    catchup=False,
+    max_active_runs=1,
+    default_args={
+        "owner": "Wei Lee",
+        "depends_on_past": False,
+    },
+)
+def discord_message_notification():
+    """Send Discord Message"""
+
+    @task(
+        retries=10,
+        retry_delay=timedelta(seconds=10),
     )
+    def send_discord_message(**context: Context) -> None:
+        triggering_asset_events = context["triggering_asset_events"]
+        session = requests.session()
+        for asset_uri, asset_event in triggering_asset_events.items():
+            logger.info(f"Receive asset event from Asset uri={asset_uri}")
+            if asset_event.extra.get("from_trigger", False):  # type: ignore[attr-defined]
+                details = asset_event.extra["payload"]  # type: ignore[attr-defined]
+            else:
+                details = asset_event.extra  # type: ignore[attr-defined]
+
+            session.post(
+                details.get("webhook_url"),
+                json={
+                    "username": details.get("username"),
+                    "content": details.get("content"),
+                },
+            )
+
+    send_discord_message()
+
+
+dag_obj = discord_message_notification()
+
+if __name__ == "__main__":
+    dag_obj.test()

dags/app/finance_bot/dag.py

@@ -4,10 +4,8 @@
 
 from datetime import datetime, timedelta
 
-from airflow.sdk import Variable, dag, task
-
-from dags.app import discord
-from dags.app.proposal_reminder.udf import get_proposal_summary
+from airflow.sdk import Asset, Metadata, Variable, dag, task
+from app.proposal_reminder.udf import get_proposal_summary
 
 DEFAULT_ARGS = {
     "owner": "Henry Lee",
@@ -19,29 +17,32 @@
 }
 
 
+cfp_summary_asset = Asset(name="CFP_summary")
+
+
 @dag(
     default_args=DEFAULT_ARGS,
     schedule="0 16 * * *",  # At 16:00 (00:00 +8)
     max_active_runs=1,
     catchup=False,
 )
 def DISCORD_PROPOSAL_REMINDER_v3():
-    @task
-    def SEND_PROPOSAL_SUMMARY():
-        webhook_url = Variable.get("DISCORD_PROGRAM_REMINDER_WEBHOOK")
-
+    @task(outlets=[cfp_summary_asset])
+    def fetch_proposal_summary():
         summary = get_proposal_summary()
         n_talk = summary["num_proposed_talk"]
         n_tutorial = summary["num_proposed_tutorial"]
-        msg = f"目前投稿議程數: {n_talk}; 課程數: {n_tutorial}"
 
-        discord.send_webhook_message(
-            webhook_url=webhook_url,
-            username="Program talk reminder",
-            msg=msg,
+        yield Metadata(
+            cfp_summary_asset,
+            extra={
+                "webhook_url": Variable.get("DISCORD_PROGRAM_REMINDER_WEBHOOK"),
+                "username": "Program talk reminder",
+                "content": f"目前投稿議程數: {n_talk}; 課程數: {n_tutorial}",
+            },
         )
 
-    SEND_PROPOSAL_SUMMARY()
+    fetch_proposal_summary()
 
 
 dag_obj = DISCORD_PROPOSAL_REMINDER_v3()

dags/app/finance_bot/udf.py

@@ -4,10 +4,8 @@
 
 from datetime import datetime, timedelta
 
-from airflow.sdk import Variable, dag, task
-
-from dags.app import discord
-from dags.app.team_registration_bot.udf import (
+from airflow.sdk import Asset, Metadata, Variable, dag, task
+from app.team_registration_bot.udf import (
     _compose_discord_msg,
     _get_statistics_from_bigquery,
 )
@@ -22,21 +20,26 @@
 }
 
 
+kktix_order_report_asset = Asset(name="kktix_order_report")
+
+
 @dag(
     default_args=DEFAULT_ARGS,
     schedule="@daily",
     max_active_runs=1,
     catchup=False,
 )
 def KKTIX_DISCORD_BOT_FOR_TEAM_REGISTRATION():
-    @task
+    @task(outlets=[kktix_order_report_asset])
     def LOAD_TO_DISCORD():
-        webhook_url = Variable.get("discord_webhook_registration_endpoint")
         statistics = _get_statistics_from_bigquery()
-        discord.send_webhook_message(
-            webhook_url=webhook_url,
-            usernmae="KKTIX order report",
-            msg=_compose_discord_msg(statistics),
+        yield Metadata(
+            kktix_order_report_asset,
+            extra={
+                "webhook_url": Variable.get("discord_webhook_registration_endpoint"),
+                "username": "KKTIX order report",
+                "content": _compose_discord_msg(statistics),
+            },
         )
 
     LOAD_TO_DISCORD()