feat: Update notification trigger patterns and add alert trigger

The code changes modify the notification constants and models to update the trigger patterns for tasks and nodes. Additionally, a new trigger for alerts is added to the constants. This change allows for more flexible matching of notification triggers.
crawlab-team · Jul 31, 2024 · 7fe770a · 7fe770a
1 parent 86a7beb
commit 7fe770a
Show file tree

Hide file tree

Showing 6 changed files with 109 additions and 14 deletions.
diff --git a/core/constants/notification.go b/core/constants/notification.go
@@ -1,8 +1,8 @@
 package constants
 
 const (
-	NotificationTriggerTargetTask = "task"
-	NotificationTriggerTargetNode = "node"
+	NotificationTriggerPatternTask = "^task"
+	NotificationTriggerPatternNode = "^node"
 )
 
 const (
@@ -12,6 +12,7 @@ const (
 	NotificationTriggerNodeStatusChange = "node_status_change"
 	NotificationTriggerNodeOnline       = "node_online"
 	NotificationTriggerNodeOffline      = "node_offline"
+	NotificationTriggerAlert            = "alert"
 )
 
 const (

diff --git a/core/grpc/server/task_server_v2.go b/core/grpc/server/task_server_v2.go
@@ -168,8 +168,10 @@ func (svr TaskServerV2) SendNotification(_ context.Context, request *grpc.TaskSe
 
 	// settings
 	settings, err := service.NewModelServiceV2[models2.NotificationSettingV2]().GetMany(bson.M{
-		"enabled":        true,
-		"trigger_target": constants.NotificationTriggerTargetTask,
+		"enabled": true,
+		"trigger": bson.M{
+			"$regex": constants.NotificationTriggerPatternTask,
+		},
 	}, nil)
 	if err != nil {
 		return nil, trace.TraceError(err)

diff --git a/core/models/models/v2/notification_alert_v2.go b/core/models/models/v2/notification_alert_v2.go
@@ -15,4 +15,5 @@ type NotificationAlertV2 struct {
 	LastingSeconds                   int                `json:"lasting_seconds" bson:"lasting_seconds"`
 	TargetValue                      float32            `json:"target_value" bson:"target_value"`
 	Level                            string             `json:"level" bson:"level"`
+	TemplateKey                      string             `json:"template_key,omitempty" bson:"template_key,omitempty"`
 }
diff --git a/core/models/models/v2/notification_setting_v2.go b/core/models/models/v2/notification_setting_v2.go
@@ -17,9 +17,8 @@ type NotificationSettingV2 struct {
 	TemplateRichTextJson string `json:"template_rich_text_json,omitempty" bson:"template_rich_text_json,omitempty"`
 	TemplateTheme        string `json:"template_theme,omitempty" bson:"template_theme,omitempty"`
 
-	TaskTrigger   string `json:"task_trigger" bson:"task_trigger"`
-	TriggerTarget string `json:"trigger_target" bson:"trigger_target"`
-	Trigger       string `json:"trigger" bson:"trigger"`
+	TaskTrigger string `json:"task_trigger" bson:"task_trigger"`
+	Trigger     string `json:"trigger" bson:"trigger"`
 
 	SenderEmail          string   `json:"sender_email,omitempty" bson:"sender_email,omitempty"`
 	UseCustomSenderEmail bool     `json:"use_custom_sender_email,omitempty" bson:"use_custom_sender_email,omitempty"`

diff --git a/core/notification/entity.go b/core/notification/entity.go
@@ -3,9 +3,11 @@ package notification
 import "github.com/crawlab-team/crawlab/core/models/models/v2"
 
 type VariableData struct {
-	Task     *models.TaskV2     `json:"task"`
-	TaskStat *models.TaskStatV2 `json:"task_stat"`
-	Spider   *models.SpiderV2   `json:"spider"`
-	Node     *models.NodeV2     `json:"node"`
-	Schedule *models.ScheduleV2 `json:"schedule"`
+	Task     *models.TaskV2              `json:"task"`
+	TaskStat *models.TaskStatV2          `json:"task_stat"`
+	Spider   *models.SpiderV2            `json:"spider"`
+	Node     *models.NodeV2              `json:"node"`
+	Schedule *models.ScheduleV2          `json:"schedule"`
+	Alert    *models.NotificationAlertV2 `json:"alert"`
+	Metric   *models.MetricV2            `json:"metric"`
 }
diff --git a/core/notification/service_v2.go b/core/notification/service_v2.go
@@ -264,6 +264,43 @@ func (svc *ServiceV2) geContentWithVariables(template string, variables []entity
 			case "updated_by":
 				content = strings.ReplaceAll(content, v.GetKey(), svc.getUsernameById(vd.Schedule.UpdatedBy))
 			}
+
+		case "alert":
+			switch v.Name {
+			case "id":
+				content = strings.ReplaceAll(content, v.GetKey(), vd.Alert.Id.Hex())
+			case "name":
+				content = strings.ReplaceAll(content, v.GetKey(), vd.Alert.Name)
+			case "description":
+				content = strings.ReplaceAll(content, v.GetKey(), vd.Alert.Description)
+			case "enabled":
+				content = strings.ReplaceAll(content, v.GetKey(), fmt.Sprintf("%t", vd.Alert.Enabled))
+			case "metric_name":
+				content = strings.ReplaceAll(content, v.GetKey(), vd.Alert.MetricName)
+			case "operator":
+				content = strings.ReplaceAll(content, v.GetKey(), vd.Alert.Operator)
+			case "lasting_seconds":
+				content = strings.ReplaceAll(content, v.GetKey(), fmt.Sprintf("%d", vd.Alert.LastingSeconds))
+			case "target_value":
+				content = strings.ReplaceAll(content, v.GetKey(), svc.getFormattedTargetValue(vd.Alert))
+			case "level":
+				content = strings.ReplaceAll(content, v.GetKey(), vd.Alert.Level)
+			}
+
+		case "metric":
+			if vd.Metric == nil {
+				content = strings.ReplaceAll(content, v.GetKey(), "N/A")
+				continue
+			}
+			switch v.Name {
+			case "type":
+				content = strings.ReplaceAll(content, v.GetKey(), vd.Metric.Type)
+			case "node_id":
+				content = strings.ReplaceAll(content, v.GetKey(), vd.Metric.NodeId.Hex())
+			default:
+				content = strings.ReplaceAll(content, v.GetKey(), svc.getFormattedMetricValue(v.Name, vd.Metric))
+			}
+
 		}
 	}
 	return content
@@ -282,6 +319,10 @@ func (svc *ServiceV2) getVariableData(args ...any) (vd VariableData) {
 			vd.Node = arg.(*models.NodeV2)
 		case *models.ScheduleV2:
 			vd.Schedule = arg.(*models.ScheduleV2)
+		case *models.NotificationAlertV2:
+			vd.Alert = arg.(*models.NotificationAlertV2)
+		case *models.MetricV2:
+			vd.Metric = arg.(*models.MetricV2)
 		}
 	}
 	return vd
@@ -336,6 +377,53 @@ func (svc *ServiceV2) getFormattedTime(t time.Time) (res string) {
 	return t.Local().Format(time.DateTime)
 }
 
+func (svc *ServiceV2) getFormattedTargetValue(a *models.NotificationAlertV2) (res string) {
+	if strings.HasSuffix(a.MetricName, "_percent") {
+		return fmt.Sprintf("%.2f%%", a.TargetValue)
+	} else if strings.HasSuffix(a.MetricName, "_memory") {
+		return fmt.Sprintf("%dMB", int(a.TargetValue/(1024*1024)))
+	} else if strings.HasSuffix(a.MetricName, "_disk") {
+		return fmt.Sprintf("%dGB", int(a.TargetValue/(1024*1024*1024)))
+	} else if strings.HasSuffix(a.MetricName, "_rate") {
+		return fmt.Sprintf("%.2fMB/s", a.TargetValue/(1024*1024))
+	} else {
+		return fmt.Sprintf("%f", a.TargetValue)
+	}
+}
+
+func (svc *ServiceV2) getFormattedMetricValue(metricName string, m *models.MetricV2) (res string) {
+	switch metricName {
+	case "cpu_usage_percent":
+		return fmt.Sprintf("%.2f%%", m.CpuUsagePercent)
+	case "total_memory":
+		return fmt.Sprintf("%dMB", m.TotalMemory/(1024*1024))
+	case "available_memory":
+		return fmt.Sprintf("%dMB", m.AvailableMemory/(1024*1024))
+	case "used_memory":
+		return fmt.Sprintf("%dMB", m.UsedMemory/(1024*1024))
+	case "used_memory_percent":
+		return fmt.Sprintf("%.2f%%", m.UsedMemoryPercent)
+	case "total_disk":
+		return fmt.Sprintf("%dGB", m.TotalDisk/(1024*1024*1024))
+	case "available_disk":
+		return fmt.Sprintf("%dGB", m.AvailableDisk/(1024*1024*1024))
+	case "used_disk":
+		return fmt.Sprintf("%dGB", m.UsedDisk/(1024*1024*1024))
+	case "used_disk_percent":
+		return fmt.Sprintf("%.2f%%", m.UsedDiskPercent)
+	case "disk_read_bytes_rate":
+		return fmt.Sprintf("%.2fMB/s", m.DiskReadBytesRate/(1024*1024))
+	case "disk_write_bytes_rate":
+		return fmt.Sprintf("%.2fMB/s", m.DiskWriteBytesRate/(1024*1024))
+	case "network_bytes_sent_rate":
+		return fmt.Sprintf("%.2fMB/s", m.NetworkBytesSentRate/(1024*1024))
+	case "network_bytes_recv_rate":
+		return fmt.Sprintf("%.2fMB/s", m.NetworkBytesRecvRate/(1024*1024))
+	default:
+		return "N/A"
+	}
+}
+
 func (svc *ServiceV2) convertMarkdownToHtml(content string) (html string) {
 	return string(markdown.ToHTML([]byte(content), nil, nil))
 }
@@ -347,8 +435,10 @@ func (svc *ServiceV2) SendNodeNotification(node *models.NodeV2) {
 
 	// settings
 	settings, err := service.NewModelServiceV2[models.NotificationSettingV2]().GetMany(bson.M{
-		"enabled":        true,
-		"trigger_target": constants.NotificationTriggerTargetNode,
+		"enabled": true,
+		"trigger": bson.M{
+			"$regex": constants.NotificationTriggerPatternNode,
+		},
 	}, nil)
 	if err != nil {
 		log.Errorf("get notification settings error: %v", err)