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

Author: Lee-W

dags/app/channel_reminder/dag.py

@@ -1,45 +1,45 @@
-"""
-Send Google Search Report to Discord
-"""
-
-from datetime import datetime, timedelta
-
-from airflow.sdk import Variable, dag, task
-from app import discord
-
-DEFAULT_ARGS = {
-    "owner": "David Jr.",
-    "depends_on_past": False,
-    "start_date": datetime(2022, 9, 15),
-    "retries": 2,
-    "retry_delay": timedelta(minutes=5),
-    "on_failure_callback": lambda x: "Need to send notification to Discord",
-}
-
-
-@dag(
-    default_args=DEFAULT_ARGS,
-    schedule="@yearly",
-    max_active_runs=1,
-    catchup=False,
-)
-def DISCORD_CHORES_REMINDER() -> None:
-    @task
-    def REMINDER_OF_THIS_TEAM() -> None:
-        webhook_url = Variable.get("DISCORD_CHORES_REMINDER_WEBHOOK")
-        discord.send_webhook_message(
-            webhook_url=webhook_url,
-            username="Data Team Airflow reminder",
-            msg=(
-                "<@&790739794148982796> <@&755827317904769184> <@&791157626099859487>\n"
-                "記得大會結束後，要有一個人負責去取消 Klaviyo 的訂閱，不然我們每個月會一直繳 $NTD2000 喔！"
-            ),
-        )
-
-    REMINDER_OF_THIS_TEAM()
-
-
-dag_obj = DISCORD_CHORES_REMINDER()
-
-if __name__ == "__main__":
-    dag_obj.test()
+# """
+# Send Google Search Report to Discord
+# """
+#
+# from datetime import datetime, timedelta
+#
+# from airflow.sdk import Variable, dag, task
+# from app import discord
+#
+# DEFAULT_ARGS = {
+#     "owner": "David Jr.",
+#     "depends_on_past": False,
+#     "start_date": datetime(2022, 9, 15),
+#     "retries": 2,
+#     "retry_delay": timedelta(minutes=5),
+#     "on_failure_callback": lambda x: "Need to send notification to Discord",
+# }
+#
+#
+# @dag(
+#     default_args=DEFAULT_ARGS,
+#     schedule="@yearly",
+#     max_active_runs=1,
+#     catchup=False,
+# )
+# def DISCORD_CHORES_REMINDER() -> None:
+#     @task
+#     def REMINDER_OF_THIS_TEAM() -> None:
+#         webhook_url = Variable.get("DISCORD_CHORES_REMINDER_WEBHOOK")
+#         discord.send_webhook_message(
+#             webhook_url=webhook_url,
+#             username="Data Team Airflow reminder",
+#             msg=(
+#                 "<@&790739794148982796> <@&755827317904769184> <@&791157626099859487>\n"
+#                 "記得大會結束後，要有一個人負責去取消 Klaviyo 的訂閱，不然我們每個月會一直繳 $NTD2000 喔！"
+#             ),
+#         )
+#
+#     REMINDER_OF_THIS_TEAM()
+#
+#
+# dag_obj = DISCORD_CHORES_REMINDER()
+#
+# if __name__ == "__main__":
+#     dag_obj.test()

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,8 +4,7 @@
 
 from datetime import datetime, timedelta
 
-from airflow.sdk import Variable, dag, task
-from app import discord
+from airflow.sdk import Asset, Metadata, Variable, dag, task
 from app.proposal_reminder.udf import get_proposal_summary
 
 DEFAULT_ARGS = {
@@ -18,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()