refactor: Update SendNotification function to handle old and new settings triggers

2026-01-21 17:21:09 +01:00 · 2024-07-18 00:05:48 +08:00
parent a5104663e8
commit 821383a677
9 changed files with 98 additions and 253 deletions
--- a/core/grpc/server/model_base_service_v2_server.go
+++ b/core/grpc/server/model_base_service_v2_server.go
@@ -27,6 +27,7 @@ var (
 		*new(models2.GitV2),
 		*new(models2.MetricV2),
 		*new(models2.NodeV2),
+		*new(models2.NotificationChannelV2),
 		*new(models2.NotificationSettingV2),
 		*new(models2.PermissionV2),
 		*new(models2.ProjectV2),
--- a/core/models/models/v2/notification_channel_v2.go
+++ b/core/models/models/v2/notification_channel_v2.go
@@ -3,15 +3,15 @@ package models
 type NotificationChannelV2 struct {
 	any                                `collection:"notification_channels"`
 	BaseModelV2[NotificationChannelV2] `bson:",inline"`
+	Type                               string `json:"type" bson:"type"`
 	Name                               string `json:"name" bson:"name"`
 	Description                        string `json:"description" bson:"description"`
-	Type                               string `json:"type" bson:"type"`
 	Provider                           string `json:"provider" bson:"provider"`
 	MailSettings                       struct {
 		SMTPServer   string `json:"smtp_server" bson:"smtp_server"`
 		SMTPPort     string `json:"smtp_port" bson:"smtp_port"`
-		SMTPFromEmailAddress  string `json:"smtp_from_email_address" bson:"smtp_from_email_address"`
-		SMTPFromEmailPassword string `json:"smtp_from_email_password" bson:"smtp_from_email_password"`
+		SMTPUser     string `json:"smtp_from_email_address" bson:"smtp_from_email_address"`
+		SMTPPassword string `json:"smtp_from_email_password" bson:"smtp_from_email_password"`
 	} `json:"mail_settings,omitempty" bson:"mail_settings,omitempty"`
 	IMSettings struct {
 		Webhook string `json:"webhook" bson:"webhook"`
--- a/core/models/models/v2/notification_setting_v2.go
+++ b/core/models/models/v2/notification_setting_v2.go
@@ -1,13 +1,14 @@
 package models

+import "go.mongodb.org/mongo-driver/bson/primitive"
+
 type NotificationSettingV2 struct {
 	any                                `collection:"notification_settings"`
 	BaseModelV2[NotificationSettingV2] `bson:",inline"`
-	Type                               string                    `json:"type" bson:"type"`
 	Name                               string `json:"name" bson:"name"`
 	Description                        string `json:"description" bson:"description"`
 	Enabled                            bool   `json:"enabled" bson:"enabled"`
-	Global                             bool                      `json:"global" bson:"global"`
+
 	Title                string `json:"title,omitempty" bson:"title,omitempty"`
 	Template             string `json:"template" bson:"template"`
 	TemplateMode         string `json:"template_mode" bson:"template_mode"`
@@ -15,24 +16,18 @@ type NotificationSettingV2 struct {
 	TemplateRichText     string `json:"template_rich_text,omitempty" bson:"template_rich_text,omitempty"`
 	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"`
-	Mail                               NotificationSettingMail   `json:"mail,omitempty" bson:"mail,omitempty"`
-	Mobile                             NotificationSettingMobile `json:"mobile,omitempty" bson:"mobile,omitempty"`
-}

-type NotificationSettingMail struct {
-	Server         string `json:"server" bson:"server"`
-	Port           string `json:"port,omitempty" bson:"port,omitempty"`
-	User           string `json:"user,omitempty" bson:"user,omitempty"`
-	Password       string `json:"password,omitempty" bson:"password,omitempty"`
-	SenderEmail    string `json:"sender_email,omitempty" bson:"sender_email,omitempty"`
-	SenderIdentity string `json:"sender_identity,omitempty" bson:"sender_identity,omitempty"`
-	To             string `json:"to,omitempty" bson:"to,omitempty"`
-	Cc             string `json:"cc,omitempty" bson:"cc,omitempty"`
-}
+	HasMail     bool   `json:"has_mail" bson:"has_mail"`
+	SenderEmail string `json:"sender_email" bson:"sender_email"`
+	SenderName  string `json:"sender_name" bson:"sender_name"`
+	MailTo      string `json:"mail_to" bson:"mail_to"`
+	MailCc      string `json:"mail_cc" bson:"mail_cc"`
+	MailBcc     string `json:"mail_bcc" bson:"mail_bcc"`

-type NotificationSettingMobile struct {
-	Webhook string `json:"webhook" bson:"webhook"`
+	ChannelIds []primitive.ObjectID    `json:"channel_ids,omitempty" bson:"channel_ids,omitempty"`
+	Channels   []NotificationChannelV2 `json:"channels,omitempty" bson:"-"`
 }
--- a/core/node/service/master_service_v2.go
+++ b/core/node/service/master_service_v2.go
@@ -12,9 +12,7 @@ import (
 	models2 "github.com/crawlab-team/crawlab/core/models/models/v2"
 	"github.com/crawlab-team/crawlab/core/models/service"
 	"github.com/crawlab-team/crawlab/core/node/config"
-	"github.com/crawlab-team/crawlab/core/notification"
 	"github.com/crawlab-team/crawlab/core/schedule"
-	"github.com/crawlab-team/crawlab/core/spider/admin"
 	"github.com/crawlab-team/crawlab/core/system"
 	"github.com/crawlab-team/crawlab/core/task/handler"
 	"github.com/crawlab-team/crawlab/core/task/scheduler"
@@ -35,8 +33,6 @@ type MasterServiceV2 struct {
 	schedulerSvc *scheduler.ServiceV2
 	handlerSvc   *handler.ServiceV2
 	scheduleSvc  *schedule.ServiceV2
-	notificationSvc *notification.ServiceV2
-	spiderAdminSvc  *admin.ServiceV2
 	systemSvc    *system.ServiceV2

 	// settings
@@ -77,12 +73,6 @@ func (svc *MasterServiceV2) Start() {
 	// start schedule service
 	go svc.scheduleSvc.Start()

-	// start notification service
-	go svc.notificationSvc.Start()
-
-	// start spider admin service
-	go svc.spiderAdminSvc.Start()
-
 	// wait for quit signal
 	svc.Wait()

@@ -361,15 +351,6 @@ func newMasterServiceV2() (res *MasterServiceV2, err error) {
 		return nil, err
 	}

-	// notification service
-	svc.notificationSvc = notification.GetNotificationServiceV2()
-
-	// spider admin service
-	svc.spiderAdminSvc, err = admin.GetSpiderAdminServiceV2()
-	if err != nil {
-		return nil, err
-	}
-
 	// system service
 	svc.systemSvc = system.GetSystemServiceV2()

--- a/core/notification/constants.go
+++ b/core/notification/constants.go
@@ -2,7 +2,7 @@ package notification

 const (
 	TypeMail = "mail"
-	TypeMobile = "mobile"
+	TypeIM   = "im"
 )

 const (
--- a/core/notification/mobile.go
+++ b/core/notification/mobile.go
@@ -2,6 +2,7 @@ package notification

 import (
 	"errors"
+	"github.com/crawlab-team/crawlab/core/models/models/v2"
 	"github.com/crawlab-team/crawlab/trace"
 	"github.com/imroc/req"
 	"strings"
@@ -12,7 +13,9 @@ type ResBody struct {
 	ErrMsg  string `json:"errmsg"`
 }

-func SendMobileNotification(webhook string, title string, content string) error {
+func SendIMNotification(s *models.NotificationSettingV2, ch *models.NotificationChannelV2, content string) error {
+	// TODO: compatibility with different IM providers
+
 	// request header
 	header := req.Header{
 		"Content-Type": "application/json; charset=utf-8",
@@ -22,7 +25,7 @@ func SendMobileNotification(webhook string, title string, content string) error
 	data := req.Param{
 		"msgtype": "markdown",
 		"markdown": req.Param{
-			"title":   title,
+			"title":   s.Title,
 			"text":    content,
 			"content": content,
 		},
@@ -32,7 +35,7 @@ func SendMobileNotification(webhook string, title string, content string) error
 		},
 		"text": content,
 	}
-	if strings.Contains(strings.ToLower(webhook), "feishu") {
+	if strings.Contains(strings.ToLower(ch.IMSettings.Webhook), "feishu") {
 		data = req.Param{
 			"msg_type": "text",
 			"content": req.Param{
@@ -42,7 +45,7 @@ func SendMobileNotification(webhook string, title string, content string) error
 	}

 	// perform request
-	res, err := req.Post(webhook, header, req.BodyJSON(&data))
+	res, err := req.Post(ch.IMSettings.Webhook, header, req.BodyJSON(&data))
 	if err != nil {
 		return trace.TraceError(err)
 	}
--- a/core/notification/mail.go
+++ b/core/notification/mail.go
@@ -14,23 +14,31 @@ import (
 	"strings"
 )

-func SendMail(m *models.NotificationSettingMail, to, cc, title, content string) error {
+func SendMail(s *models.NotificationSettingV2, ch *models.NotificationChannelV2, to, cc, bcc, title, content string) error {
+	// mail settings
+	ms := ch.MailSettings
+
 	// config
-	port, _ := strconv.Atoi(m.Port)
-	password := m.Password
-	SMTPUser := m.User
-	smtpConfig := smtpAuthentication{
-		Server:         m.Server,
-		Port:           port,
-		SenderEmail:    m.SenderEmail,
-		SenderIdentity: m.SenderIdentity,
-		SMTPPassword:   password,
-		SMTPUser:       SMTPUser,
+	port, err := strconv.Atoi(ms.SMTPPort)
+	if err != nil {
+		log.Errorf("failed to convert SMTP port to int: %v", err)
+		trace.PrintError(err)
+		return err
 	}
+	smtpConfig := smtpAuthentication{
+		Server:         ms.SMTPServer,
+		Port:           port,
+		SenderIdentity: s.SenderName,
+		SenderEmail:    s.SenderEmail,
+		SMTPUser:       ms.SMTPUser,
+		SMTPPassword:   ms.SMTPPassword,
+	}
+
 	options := sendOptions{
+		Subject: title,
 		To:      to,
 		Cc:      cc,
-		Subject: title,
+		Bcc:     bcc,
 	}

 	// convert html to text
@@ -80,9 +88,10 @@ type smtpAuthentication struct {

 // sendOptions are options for sending an email
 type sendOptions struct {
-	To      string
 	Subject string
+	To      string
 	Cc      string
+	Bcc     string
 }

 // send email
@@ -99,14 +108,6 @@ func send(smtpConfig smtpAuthentication, options sendOptions, htmlBody string, t
 		return errors.New("SMTP user is empty")
 	}

-	if smtpConfig.SenderIdentity == "" {
-		return errors.New("SMTP sender identity is empty")
-	}
-
-	if smtpConfig.SenderEmail == "" {
-		return errors.New("SMTP sender email is empty")
-	}
-
 	if options.To == "" {
 		return errors.New("no receiver emails configured")
 	}
@@ -134,6 +135,9 @@ func send(smtpConfig smtpAuthentication, options sendOptions, htmlBody string, t
 	if options.Cc != "" {
 		m.SetHeader("Cc", getRecipientList(options.Cc)...)
 	}
+	if options.Bcc != "" {
+		m.SetHeader("Bcc", getRecipientList(options.Bcc)...)
+	}

 	m.SetBody("text/plain", txtBody)
 	m.AddAlternative("text/html", htmlBody)
--- a/core/notification/service_v2.go
+++ b/core/notification/service_v2.go
@@ -7,6 +7,7 @@ import (
 	"github.com/crawlab-team/crawlab/core/entity"
 	"github.com/crawlab-team/crawlab/core/models/models/v2"
 	"github.com/crawlab-team/crawlab/core/models/service"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 	"regexp"
 	"strings"
 	"sync"
@@ -15,181 +16,45 @@ import (
 type ServiceV2 struct {
 }

-func (svc *ServiceV2) Start() {
-	// initialize data
-	if err := svc.initData(); err != nil {
-		log.Errorf("[NotificationServiceV2] initializing data error: %v", err)
-		return
-	}
-}
-
-func (svc *ServiceV2) initData() (err error) {
-	total, err := service.NewModelServiceV2[models.NotificationSettingV2]().Count(nil)
-	if err != nil {
-		return err
-	}
-	if total > 0 {
-		return nil
-	}
-
-	// data to initialize
-	settings := []models.NotificationSettingV2{
-		{
-			Type:        TypeMail,
-			Enabled:     true,
-			Name:        "任务通知（邮件）",
-			Description: "这是默认的邮件通知。您可以使用您自己的设置进行编辑。",
-			TaskTrigger: constants.NotificationTriggerTaskFinish,
-			Title:       "[Crawlab] 爬虫任务更新: {{$.status}}",
-			Template: `尊敬的 {{$.user.username}},
-
-请查看下面的任务数据。
-
-|键|值|
-|:-:|:--|
-|任务状态|{{$.status}}|
-|任务优先级|{{$.priority}}|
-|任务模式|{{$.mode}}|
-|执行命令|{{$.cmd}}|
-|执行参数|{{$.param}}|
-|错误信息|{{$.error}}|
-|节点|{{$.node.name}}|
-|爬虫|{{$.spider.name}}|
-|项目|{{$.spider.project.name}}|
-|定时任务|{{$.schedule.name}}|
-|结果数|{{$.:task_stat.result_count}}|
-|等待时间（秒）|{#{{$.:task_stat.wait_duration}}/1000#}|
-|运行时间（秒）|{#{{$.:task_stat.runtime_duration}}/1000#}|
-|总时间（秒）|{#{{$.:task_stat.total_duration}}/1000#}|
-|平均结果数/秒|{#{{$.:task_stat.result_count}}/({{$.:task_stat.total_duration}}/1000)#}|
-`,
-			Mail: models.NotificationSettingMail{
-				Server: "smtp.163.com",
-				Port:   "465",
-				To:     "{{$.user[create].email}}",
-			},
-		},
-		{
-			Type:        TypeMail,
-			Enabled:     true,
-			Name:        "Task Change (Mail)",
-			Description: "This is the default mail notification. You can edit it with your own settings",
-			TaskTrigger: constants.NotificationTriggerTaskFinish,
-			Title:       "[Crawlab] Task Update: {{$.status}}",
-			Template: `Dear {{$.user.username}},
-
-Please find the task data as below.
-
-|Key|Value|
-|:-:|:--|
-|Task Status|{{$.status}}|
-|Task Priority|{{$.priority}}|
-|Task Mode|{{$.mode}}|
-|Task Command|{{$.cmd}}|
-|Task Params|{{$.param}}|
-|Error Message|{{$.error}}|
-|Node|{{$.node.name}}|
-|Spider|{{$.spider.name}}|
-|Project|{{$.spider.project.name}}|
-|Schedule|{{$.schedule.name}}|
-|Result Count|{{$.:task_stat.result_count}}|
-|Wait Duration (sec)|{#{{$.:task_stat.wait_duration}}/1000#}|
-|Runtime Duration (sec)|{#{{$.:task_stat.runtime_duration}}/1000#}|
-|Total Duration (sec)|{#{{$.:task_stat.total_duration}}/1000#}|
-|Avg Results / Sec|{#{{$.:task_stat.result_count}}/({{$.:task_stat.total_duration}}/1000)#}|
-`,
-			Mail: models.NotificationSettingMail{
-				Server: "smtp.163.com",
-				Port:   "465",
-				To:     "{{$.user[create].email}}",
-			},
-		},
-		{
-			Type:        TypeMobile,
-			Enabled:     true,
-			Name:        "任务通知（移动端）",
-			Description: "这是默认的手机通知。您可以使用您自己的设置进行编辑。",
-			TaskTrigger: constants.NotificationTriggerTaskFinish,
-			Title:       "[Crawlab] 任务更新: {{$.status}}",
-			Template: `尊敬的 {{$.user.username}},
-
-请查看下面的任务数据。
-
- **任务状态**: {{$.status}}
- **任务优先级**: {{$.priority}}
- **任务模式**: {{$.mode}}
- **执行命令**: {{$.cmd}}
- **执行参数**: {{$.param}}
- **错误信息**: {{$.error}}
- **节点**: {{$.node.name}}
- **爬虫**: {{$.spider.name}}
- **项目**: {{$.spider.project.name}}
- **定时任务**: {{$.schedule.name}}
- **结果数**: {{$.:task_stat.result_count}}
- **等待时间（秒）**: {#{{$.:task_stat.wait_duration}}/1000#}
- **运行时间（秒）**: {#{{$.:task_stat.runtime_duration}}/1000#}
- **总时间（秒）**: {#{{$.:task_stat.total_duration}}/1000#}
- **平均结果数/秒**: {#{{$.:task_stat.result_count}}/({{$.:task_stat.total_duration}}/1000)#}`,
-			Mobile: models.NotificationSettingMobile{},
-		},
-		{
-			Type:        TypeMobile,
-			Enabled:     true,
-			Name:        "Task Change (Mobile)",
-			Description: "This is the default mobile notification. You can edit it with your own settings",
-			TaskTrigger: constants.NotificationTriggerTaskFinish,
-			Title:       "[Crawlab] Task Update: {{$.status}}",
-			Template: `Dear {{$.user.username}},
-
-Please find the task data as below.
-
- **Task Status**: {{$.status}}
- **Task Priority**: {{$.priority}}
- **Task Mode**: {{$.mode}}
- **Task Command**: {{$.cmd}}
- **Task Params**: {{$.param}}
- **Error Message**: {{$.error}}
- **Node**: {{$.node.name}}
- **Spider**: {{$.spider.name}}
- **Project**: {{$.spider.project.name}}
- **Schedule**: {{$.schedule.name}}
- **Result Count**: {{$.:task_stat.result_count}}
- **Wait Duration (sec)**: {#{{$.:task_stat.wait_duration}}/1000#}
- **Runtime Duration (sec)**: {#{{$.:task_stat.runtime_duration}}/1000#}
- **Total Duration (sec)**: {#{{$.:task_stat.total_duration}}/1000#}
- **Avg Results / Sec**: {#{{$.:task_stat.result_count}}/({{$.:task_stat.total_duration}}/1000)#}`,
-			Mobile: models.NotificationSettingMobile{},
-		},
-	}
-	_, err = service.NewModelServiceV2[models.NotificationSettingV2]().InsertMany(settings)
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
 func (svc *ServiceV2) Send(s *models.NotificationSettingV2, args ...any) {
 	content := svc.getContent(s, args...)
-	switch s.Type {
-	case TypeMail:
-		svc.SendMail(s, content)
-	case TypeMobile:
-		svc.SendMobile(s, content)
+
+	wg := sync.WaitGroup{}
+	wg.Add(len(s.ChannelIds))
+	for _, chId := range s.ChannelIds {
+		go func(chId primitive.ObjectID) {
+			defer wg.Done()
+			ch, err := service.NewModelServiceV2[models.NotificationChannelV2]().GetById(chId)
+			if err != nil {
+				log.Errorf("[NotificationServiceV2] get channel error: %v", err)
+				return
 			}
+			switch ch.Type {
+			case TypeMail:
+				svc.SendMail(s, ch, content)
+			case TypeIM:
+				svc.SendIM(s, ch, content)
+			}
+		}(chId)
+	}
+	wg.Wait()
 }

-func (svc *ServiceV2) SendMail(s *models.NotificationSettingV2, content string) {
+func (svc *ServiceV2) SendMail(s *models.NotificationSettingV2, ch *models.NotificationChannelV2, content string) {
 	// TODO: parse to/cc/bcc
+	mailTo := s.MailTo
+	mailCc := s.MailCc
+	mailBcc := s.MailBcc

 	// send mail
-	err := SendMail(&s.Mail, s.Mail.To, s.Mail.Cc, s.Title, content)
+	err := SendMail(s, ch, mailTo, mailCc, mailBcc, s.Title, content)
 	if err != nil {
 		log.Errorf("[NotificationServiceV2] send mail error: %v", err)
 	}
 }

-func (svc *ServiceV2) SendMobile(s *models.NotificationSettingV2, content string) {
-	err := SendMobileNotification(s.Mobile.Webhook, s.Title, content)
+func (svc *ServiceV2) SendIM(s *models.NotificationSettingV2, ch *models.NotificationChannelV2, content string) {
+	err := SendIMNotification(s, ch, content)
 	if err != nil {
 		log.Errorf("[NotificationServiceV2] send mobile notification error: %v", err)
 	}
--- a/core/spider/admin/service_v2.go
+++ b/core/spider/admin/service_v2.go
@@ -29,10 +29,6 @@ type ServiceV2 struct {
 	cfgPath string
 }

-func (svc *ServiceV2) Start() (err error) {
-	return nil
-}
-
 func (svc *ServiceV2) Schedule(id primitive.ObjectID, opts *interfaces.SpiderRunOptions) (taskIds []primitive.ObjectID, err error) {
 	// spider
 	s, err := service.NewModelServiceV2[models2.SpiderV2]().GetById(id)