From 821383a677bed7702e2ee557c0bdb4db1677a737 Mon Sep 17 00:00:00 2001 From: Marvin Zhang Date: Thu, 18 Jul 2024 00:05:48 +0800 Subject: [PATCH] refactor: Update SendNotification function to handle old and new settings triggers --- .../server/model_base_service_v2_server.go | 1 + .../models/v2/notification_channel_v2.go | 10 +- .../models/v2/notification_setting_v2.go | 55 +++-- core/node/service/master_service_v2.go | 31 +-- core/notification/constants.go | 4 +- core/notification/{mobile.go => im.go} | 11 +- core/notification/mail.go | 46 +++-- core/notification/service_v2.go | 189 +++--------------- core/spider/admin/service_v2.go | 4 - 9 files changed, 98 insertions(+), 253 deletions(-) rename core/notification/{mobile.go => im.go} (70%) diff --git a/core/grpc/server/model_base_service_v2_server.go b/core/grpc/server/model_base_service_v2_server.go index 7f356ec4..59ef7d11 100644 --- 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), diff --git a/core/models/models/v2/notification_channel_v2.go b/core/models/models/v2/notification_channel_v2.go index 035baeb0..2f328b79 100644 --- 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"` + SMTPServer string `json:"smtp_server" bson:"smtp_server"` + SMTPPort string `json:"smtp_port" bson:"smtp_port"` + 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"` diff --git a/core/models/models/v2/notification_setting_v2.go b/core/models/models/v2/notification_setting_v2.go index 7164a505..6f36c3ec 100644 --- a/core/models/models/v2/notification_setting_v2.go +++ b/core/models/models/v2/notification_setting_v2.go @@ -1,38 +1,33 @@ 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"` - TemplateMarkdown string `json:"template_markdown,omitempty" bson:"template_markdown,omitempty"` - 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"` -} + Name string `json:"name" bson:"name"` + Description string `json:"description" bson:"description"` + Enabled bool `json:"enabled" bson:"enabled"` -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"` -} + Title string `json:"title,omitempty" bson:"title,omitempty"` + Template string `json:"template" bson:"template"` + TemplateMode string `json:"template_mode" bson:"template_mode"` + TemplateMarkdown string `json:"template_markdown,omitempty" bson:"template_markdown,omitempty"` + 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"` -type NotificationSettingMobile struct { - Webhook string `json:"webhook" bson:"webhook"` + TaskTrigger string `json:"task_trigger" bson:"task_trigger"` + TriggerTarget string `json:"trigger_target" bson:"trigger_target"` + Trigger string `json:"trigger" bson:"trigger"` + + 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"` + + ChannelIds []primitive.ObjectID `json:"channel_ids,omitempty" bson:"channel_ids,omitempty"` + Channels []NotificationChannelV2 `json:"channels,omitempty" bson:"-"` } diff --git a/core/node/service/master_service_v2.go b/core/node/service/master_service_v2.go index 1487547b..87f9c85d 100644 --- 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" @@ -30,14 +28,12 @@ import ( type MasterServiceV2 struct { // dependencies - cfgSvc interfaces.NodeConfigService - server *server.GrpcServerV2 - schedulerSvc *scheduler.ServiceV2 - handlerSvc *handler.ServiceV2 - scheduleSvc *schedule.ServiceV2 - notificationSvc *notification.ServiceV2 - spiderAdminSvc *admin.ServiceV2 - systemSvc *system.ServiceV2 + cfgSvc interfaces.NodeConfigService + server *server.GrpcServerV2 + schedulerSvc *scheduler.ServiceV2 + handlerSvc *handler.ServiceV2 + scheduleSvc *schedule.ServiceV2 + systemSvc *system.ServiceV2 // settings cfgPath string @@ -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() diff --git a/core/notification/constants.go b/core/notification/constants.go index a69491c3..b3904af8 100644 --- a/core/notification/constants.go +++ b/core/notification/constants.go @@ -1,8 +1,8 @@ package notification const ( - TypeMail = "mail" - TypeMobile = "mobile" + TypeMail = "mail" + TypeIM = "im" ) const ( diff --git a/core/notification/mobile.go b/core/notification/im.go similarity index 70% rename from core/notification/mobile.go rename to core/notification/im.go index 90a23870..3650a4db 100644 --- a/core/notification/mobile.go +++ b/core/notification/im.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) } diff --git a/core/notification/mail.go b/core/notification/mail.go index 289381aa..118ca4ae 100644 --- 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) diff --git a/core/notification/service_v2.go b/core/notification/service_v2.go index 44118333..ded20c93 100644 --- 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) } diff --git a/core/spider/admin/service_v2.go b/core/spider/admin/service_v2.go index cf821201..543fe453 100644 --- 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)