From aedb2ccf0cdd3e869827748b39df673743e39bfe Mon Sep 17 00:00:00 2001
From: yezi <yehaowei20060411@qq.com>
Date: Mon, 26 Feb 2024 06:32:04 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E4=BF=AE=E6=94=B9=E5=B7=A5=E4=BD=9C?=
 =?UTF-8?q?=E6=B5=81=E8=BF=90=E8=A1=8C=E5=A4=B1=E8=B4=A5=E9=80=9A=E7=9F=A5?=
 =?UTF-8?q?=E6=8E=A8=E9=80=81=E6=96=B9=E5=BC=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 config.example.yaml        |  6 +++++-
 utils/config.py            |  7 +++++++
 utils/config_generators.py |  3 +++
 utils/event_handlers.py    | 37 +++++++++++++++++++++++++++++++++++++
 4 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 utils/event_handlers.py

diff --git a/config.example.yaml b/config.example.yaml
index a2a47be..9da2905 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -5,4 +5,8 @@ mongodb:
   database: jfetcher
 log:
   save_level: DEBUG
-  print_level: DEBUG
\ No newline at end of file
+  print_level: DEBUG
+feishu_notification:
+  enabled: false
+  webhook_url: <webhook-url>
+  failure_card_id: <card-id>
\ No newline at end of file
diff --git a/utils/config.py b/utils/config.py
index 57d5fde..15b0242 100644
--- a/utils/config.py
+++ b/utils/config.py
@@ -7,10 +7,17 @@
 )
 
 
+class _FeishuNotification(Struct, **CONFIG_STRUCT_CONFIG):
+    enabled: bool = False
+    webhook_url: str = "<webhook-url>"
+    failure_card_id: str = "<card-id>"
+
+
 class _Config(Struct, **CONFIG_STRUCT_CONFIG):
     version: str = "v3.0.0"
     mongodb: MongoDBConfig = MongoDBConfig()
     log: LoggingConfig = LoggingConfig()
+    feishu_notification: _FeishuNotification = _FeishuNotification()
 
 
 CONFIG = load_or_save_default_config(_Config)
diff --git a/utils/config_generators.py b/utils/config_generators.py
index 499c6d3..d83758a 100644
--- a/utils/config_generators.py
+++ b/utils/config_generators.py
@@ -3,6 +3,7 @@
 from prefect.client.schemas.schedules import CronSchedule
 
 from utils.config import CONFIG
+from utils.event_handlers import on_failure_or_crashed
 
 
 def generate_flow_config(
@@ -17,6 +18,8 @@ def generate_flow_config(
         "retries": retries,
         "retry_delay_seconds": retry_delay_seconds,
         "timeout_seconds": timeout,
+        "on_failure": [on_failure_or_crashed],
+        "on_crashed": [on_failure_or_crashed],
     }
 
 
diff --git a/utils/event_handlers.py b/utils/event_handlers.py
new file mode 100644
index 0000000..531a813
--- /dev/null
+++ b/utils/event_handlers.py
@@ -0,0 +1,37 @@
+from httpx import AsyncClient
+from prefect.client.schemas.objects import Flow, FlowRun, State
+
+from utils.config import CONFIG
+
+CLIENT = AsyncClient(http2=True)
+
+
+async def on_failure_or_crashed(flow: Flow, flow_run: FlowRun, state: State) -> None:
+    if not CONFIG.feishu_notification.enabled:
+        return
+
+    response = await CLIENT.post(
+        url=CONFIG.feishu_notification.webhook_url,
+        json={
+            "msg_type": "interactive",
+            "card": {
+                "type": "template",
+                "data": {
+                    "template_id": CONFIG.feishu_notification.failure_card_id,
+                    "template_variable": {
+                        "status": flow_run.state_name,
+                        "flow_name": flow.name,
+                        "run_name": flow_run.name,
+                        "deployment_name": "[To-Do]",
+                        "start_time": flow_run.start_time.strftime(r"%Y-%m-%d %H:%M:%S")
+                        if flow_run.start_time
+                        else "[未知]",
+                        "error_message": state.message,
+                        "details_url": f"http://prod-server:4200/flow-runs/flow-run/{flow_run.id}",
+                    },
+                },
+            },
+        },
+    )
+
+    response.raise_for_status()