From 96b89d05b8b807150628f6f24cce786f32d2d635 Mon Sep 17 00:00:00 2001 From: Marvin Zhang Date: Thu, 25 Jul 2024 22:58:39 +0800 Subject: [PATCH] refactor: Update NotificationChannelV2 model to include Telegram notification settings --- core/models/common/index_service.go | 145 ------------------ core/models/common/index_service_v2.go | 25 +++ core/models/delegate/model_user_role_test.go | 97 ------------ .../models/v2/notification_request_v2.go | 12 +- core/notification/constants.go | 5 + core/notification/im.go | 36 ++--- core/notification/service_v2.go | 29 +++- 7 files changed, 81 insertions(+), 268 deletions(-) delete mode 100644 core/models/common/index_service.go delete mode 100644 core/models/delegate/model_user_role_test.go diff --git a/core/models/common/index_service.go b/core/models/common/index_service.go deleted file mode 100644 index ff75b92e..00000000 --- a/core/models/common/index_service.go +++ /dev/null @@ -1,145 +0,0 @@ -package common - -import ( - "github.com/crawlab-team/crawlab/core/constants" - "github.com/crawlab-team/crawlab/core/interfaces" - "github.com/crawlab-team/crawlab/db/mongo" - "go.mongodb.org/mongo-driver/bson" - mongo2 "go.mongodb.org/mongo-driver/mongo" - "go.mongodb.org/mongo-driver/mongo/options" -) - -func CreateIndexes() { - // artifacts - mongo.GetMongoCol(interfaces.ModelColNameArtifact).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"_col": 1}}, - {Keys: bson.M{"_del": 1}}, - {Keys: bson.M{"_tid": 1}}, - }) - - // tags - mongo.GetMongoCol(interfaces.ModelColNameTag).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"col": 1}}, - {Keys: bson.M{"name": 1}}, - }) - - // nodes - mongo.GetMongoCol(interfaces.ModelColNameNode).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"key": 1}}, // key - {Keys: bson.M{"name": 1}}, // name - {Keys: bson.M{"is_master": 1}}, // is_master - {Keys: bson.M{"status": 1}}, // status - {Keys: bson.M{"enabled": 1}}, // enabled - {Keys: bson.M{"active": 1}}, // active - }) - - // projects - mongo.GetMongoCol(interfaces.ModelColNameNode).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"name": 1}}, - }) - - // spiders - mongo.GetMongoCol(interfaces.ModelColNameSpider).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"name": 1}}, - {Keys: bson.M{"type": 1}}, - {Keys: bson.M{"col_id": 1}}, - {Keys: bson.M{"project_id": 1}}, - }) - - // tasks - mongo.GetMongoCol(interfaces.ModelColNameTask).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"spider_id": 1}}, - {Keys: bson.M{"status": 1}}, - {Keys: bson.M{"node_id": 1}}, - {Keys: bson.M{"schedule_id": 1}}, - {Keys: bson.M{"type": 1}}, - {Keys: bson.M{"mode": 1}}, - {Keys: bson.M{"priority": 1}}, - {Keys: bson.M{"parent_id": 1}}, - {Keys: bson.M{"has_sub": 1}}, - {Keys: bson.M{"create_ts": -1}}, - }) - - // task stats - mongo.GetMongoCol(interfaces.ModelColNameTaskStat).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"create_ts": 1}}, - }) - - // schedules - mongo.GetMongoCol(interfaces.ModelColNameSchedule).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"name": 1}}, - {Keys: bson.M{"spider_id": 1}}, - {Keys: bson.M{"enabled": 1}}, - }) - - // users - mongo.GetMongoCol(interfaces.ModelColNameUser).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"username": 1}}, - {Keys: bson.M{"role": 1}}, - {Keys: bson.M{"email": 1}}, - }) - - // settings - mongo.GetMongoCol(interfaces.ModelColNameSetting).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"key": 1}}, - }) - - // tokens - mongo.GetMongoCol(interfaces.ModelColNameToken).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"name": 1}}, - }) - - // variables - mongo.GetMongoCol(interfaces.ModelColNameVariable).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"key": 1}}, - }) - - // data sources - mongo.GetMongoCol(interfaces.ModelColNameDataSource).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"name": 1}}, - }) - - // data collections - mongo.GetMongoCol(interfaces.ModelColNameDataCollection).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"name": 1}}, - }) - - // extra values - mongo.GetMongoCol(interfaces.ModelColNameExtraValues).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.M{"oid": 1}}, - {Keys: bson.M{"m": 1}}, - {Keys: bson.M{"t": 1}}, - {Keys: bson.M{"m": 1, "t": 1}}, - {Keys: bson.M{"oid": 1, "m": 1, "t": 1}}, - }) - - // roles - mongo.GetMongoCol(interfaces.ModelColNameRole).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.D{{"key", 1}}, Options: options.Index().SetUnique(true)}, - }) - - // user role relations - mongo.GetMongoCol(interfaces.ModelColNameUserRole).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.D{{"user_id", 1}, {"role_id", 1}}, Options: options.Index().SetUnique(true)}, - {Keys: bson.D{{"role_id", 1}, {"user_id", 1}}, Options: options.Index().SetUnique(true)}, - }) - - // permissions - mongo.GetMongoCol(interfaces.ModelColNamePermission).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.D{{"key", 1}}, Options: options.Index().SetUnique(true)}, - }) - - // role permission relations - mongo.GetMongoCol(interfaces.ModelColNameRolePermission).MustCreateIndexes([]mongo2.IndexModel{ - {Keys: bson.D{{"role_id", 1}, {"permission_id", 1}}, Options: options.Index().SetUnique(true)}, - {Keys: bson.D{{"permission_id", 1}, {"role_id", 1}}, Options: options.Index().SetUnique(true)}, - }) - - // cache - mongo.GetMongoCol(constants.CacheColName).MustCreateIndexes([]mongo2.IndexModel{ - { - Keys: bson.M{constants.CacheColTime: 1}, - Options: options.Index().SetExpireAfterSeconds(3600 * 24), - }, - }) -} diff --git a/core/models/common/index_service_v2.go b/core/models/common/index_service_v2.go index 9db29c88..2805ad3e 100644 --- a/core/models/common/index_service_v2.go +++ b/core/models/common/index_service_v2.go @@ -177,4 +177,29 @@ func CreateIndexesV2() { }, }, }) + + // notification requests + mongo.GetMongoCol(service.GetCollectionNameByInstance(models2.NotificationRequestV2{})).MustCreateIndexes([]mongo2.IndexModel{ + { + Keys: bson.D{ + {"created_ts", -1}, + }, + Options: (&options.IndexOptions{}).SetExpireAfterSeconds(60 * 60 * 24 * 7), + }, + { + Keys: bson.D{ + {"channel_id", 1}, + }, + }, + { + Keys: bson.D{ + {"setting_id", 1}, + }, + }, + { + Keys: bson.D{ + {"status", 1}, + }, + }, + }) } diff --git a/core/models/delegate/model_user_role_test.go b/core/models/delegate/model_user_role_test.go deleted file mode 100644 index d55b90c6..00000000 --- a/core/models/delegate/model_user_role_test.go +++ /dev/null @@ -1,97 +0,0 @@ -package delegate_test - -import ( - "github.com/crawlab-team/crawlab/core/interfaces" - "github.com/crawlab-team/crawlab/core/models/common" - "github.com/crawlab-team/crawlab/core/models/delegate" - models2 "github.com/crawlab-team/crawlab/core/models/models" - "github.com/crawlab-team/crawlab/db/mongo" - "github.com/spf13/viper" - "github.com/stretchr/testify/require" - "go.mongodb.org/mongo-driver/bson/primitive" - "testing" -) - -func init() { - viper.Set("mongo.db", "crawlab_test") - common.CreateIndexes() -} - -func TestUserRole_Add(t *testing.T) { - SetupTest(t) - - p := &models2.UserRole{} - - err := delegate.NewModelDelegate(p).Add() - require.Nil(t, err) - require.NotNil(t, p.Id) - - a, err := delegate.NewModelDelegate(p).GetArtifact() - require.Nil(t, err) - require.Equal(t, p.Id, a.GetId()) - require.NotNil(t, a.GetSys().GetCreateTs()) - require.NotNil(t, a.GetSys().GetUpdateTs()) -} - -func TestUserRole_Save(t *testing.T) { - SetupTest(t) - - p := &models2.UserRole{ - UserId: primitive.NewObjectID(), - RoleId: primitive.NewObjectID(), - } - - err := delegate.NewModelDelegate(p).Add() - require.Nil(t, err) - - uid := primitive.NewObjectID() - rid := primitive.NewObjectID() - p.UserId = uid - p.RoleId = rid - err = delegate.NewModelDelegate(p).Save() - require.Nil(t, err) - - err = mongo.GetMongoCol(interfaces.ModelColNameUserRole).FindId(p.Id).One(&p) - require.Nil(t, err) - require.Equal(t, uid, p.UserId) - require.Equal(t, rid, p.RoleId) -} - -func TestUserRole_Delete(t *testing.T) { - SetupTest(t) - - p := &models2.UserRole{} - - err := delegate.NewModelDelegate(p).Add() - require.Nil(t, err) - - err = delegate.NewModelDelegate(p).Delete() - require.Nil(t, err) - - var a models2.Artifact - col := mongo.GetMongoCol(interfaces.ModelColNameArtifact) - err = col.FindId(p.Id).One(&a) - require.Nil(t, err) - require.NotNil(t, a.Obj) - require.True(t, a.Del) -} - -func TestUserRole_AddDuplicates(t *testing.T) { - SetupTest(t) - - uid := primitive.NewObjectID() - rid := primitive.NewObjectID() - p := &models2.UserRole{ - UserId: uid, - RoleId: rid, - } - p2 := &models2.UserRole{ - UserId: uid, - RoleId: rid, - } - - err := delegate.NewModelDelegate(p).Add() - require.Nil(t, err) - err = delegate.NewModelDelegate(p2).Add() - require.NotNil(t, err) -} diff --git a/core/models/models/v2/notification_request_v2.go b/core/models/models/v2/notification_request_v2.go index 1130cc92..cc2836d7 100644 --- a/core/models/models/v2/notification_request_v2.go +++ b/core/models/models/v2/notification_request_v2.go @@ -5,8 +5,12 @@ import "go.mongodb.org/mongo-driver/bson/primitive" type NotificationRequestV2 struct { any `collection:"notification_requests"` BaseModelV2[NotificationRequestV2] `bson:",inline"` - Status string `json:"status" bson:"status"` - Error string `json:"error,omitempty" bson:"error,omitempty"` - SettingId primitive.ObjectID `json:"setting_id" bson:"setting_id"` - ChannelId primitive.ObjectID `json:"channel_id" bson:"channel_id"` + Status string `json:"status" bson:"status"` + Error string `json:"error,omitempty" bson:"error,omitempty"` + Title string `json:"title" bson:"title"` + Content string `json:"content" bson:"content"` + SettingId primitive.ObjectID `json:"setting_id" bson:"setting_id"` + ChannelId primitive.ObjectID `json:"channel_id" bson:"channel_id"` + Setting *NotificationSettingV2 `json:"setting,omitempty" bson:"-"` + Channel *NotificationChannelV2 `json:"channel,omitempty" bson:"-"` } diff --git a/core/notification/constants.go b/core/notification/constants.go index 6c483d4b..d8de0738 100644 --- a/core/notification/constants.go +++ b/core/notification/constants.go @@ -24,3 +24,8 @@ const ( ChannelIMProviderDingtalk = "dingtalk" // https://open.dingtalk.com/document/orgapp/custom-robot-access ChannelIMProviderLark = "lark" // https://www.larksuite.com/hc/en-US/articles/099698615114-use-webhook-triggers ) + +const ( + StatusSuccess = "success" + StatusError = "error" +) diff --git a/core/notification/im.go b/core/notification/im.go index 3af47cd4..52a1c420 100644 --- a/core/notification/im.go +++ b/core/notification/im.go @@ -17,16 +17,15 @@ type ResBody struct { } func SendIMNotification(ch *models.NotificationChannelV2, title, content string) error { - // TODO: compatibility with different IM providers switch ch.Provider { case ChannelIMProviderLark: return sendIMLark(ch, title, content) - case ChannelIMProviderSlack: - return sendIMSlack(ch, title, content) case ChannelIMProviderDingtalk: return sendIMDingTalk(ch, title, content) case ChannelIMProviderWechatWork: return sendIMWechatWork(ch, title, content) + case ChannelIMProviderSlack: + return sendIMSlack(ch, title, content) case ChannelIMProviderTelegram: return sendIMTelegram(ch, title, content) case ChannelIMProviderDiscord: @@ -221,23 +220,6 @@ func sendIMLark(ch *models.NotificationChannelV2, title, content string) error { return nil } -func sendIMSlack(ch *models.NotificationChannelV2, title, content string) error { - data := req.Param{ - "blocks": []req.Param{ - {"type": "header", "text": req.Param{"type": "plain_text", "text": title}}, - {"type": "section", "text": req.Param{"type": "mrkdwn", "text": convertMarkdownToSlack(content)}}, - }, - } - resBody, err := performIMRequestWithJson[ResBody](ch.WebhookUrl, data) - if err != nil { - return err - } - if resBody.ErrCode != 0 { - return errors.New(resBody.ErrMsg) - } - return nil -} - func sendIMDingTalk(ch *models.NotificationChannelV2, title string, content string) error { data := req.Param{ "msgtype": "markdown", @@ -273,6 +255,20 @@ func sendIMWechatWork(ch *models.NotificationChannelV2, title string, content st return nil } +func sendIMSlack(ch *models.NotificationChannelV2, title, content string) error { + data := req.Param{ + "blocks": []req.Param{ + {"type": "header", "text": req.Param{"type": "plain_text", "text": title}}, + {"type": "section", "text": req.Param{"type": "mrkdwn", "text": convertMarkdownToSlack(content)}}, + }, + } + _, err := performIMRequest(ch.WebhookUrl, data) + if err != nil { + return err + } + return nil +} + func sendIMTelegram(ch *models.NotificationChannelV2, title string, content string) error { type ResBody struct { Ok bool `json:"ok"` diff --git a/core/notification/service_v2.go b/core/notification/service_v2.go index 407e61e6..82d8076a 100644 --- a/core/notification/service_v2.go +++ b/core/notification/service_v2.go @@ -38,7 +38,7 @@ func (svc *ServiceV2) Send(s *models.NotificationSettingV2, args ...any) { case TypeMail: svc.SendMail(s, ch, title, content) case TypeIM: - svc.SendIM(ch, title, content) + svc.SendIM(s, ch, title, content) } }(chId) } @@ -56,13 +56,15 @@ func (svc *ServiceV2) SendMail(s *models.NotificationSettingV2, ch *models.Notif if err != nil { log.Errorf("[NotificationServiceV2] send mail error: %v", err) } + go svc.saveRequest(s, ch, title, content, err) } -func (svc *ServiceV2) SendIM(ch *models.NotificationChannelV2, title, content string) { +func (svc *ServiceV2) SendIM(s *models.NotificationSettingV2, ch *models.NotificationChannelV2, title, content string) { err := SendIMNotification(ch, title, content) if err != nil { log.Errorf("[NotificationServiceV2] send mobile notification error: %v", err) } + go svc.saveRequest(s, ch, title, content, err) } func (svc *ServiceV2) getContent(s *models.NotificationSettingV2, ch *models.NotificationChannelV2, args ...any) (content string) { @@ -361,6 +363,29 @@ func (svc *ServiceV2) SendNodeNotification(node *models.NodeV2) { } } +func (svc *ServiceV2) saveRequest(s *models.NotificationSettingV2, ch *models.NotificationChannelV2, title, content string, err error) { + status := StatusSuccess + errMsg := "" + if err != nil { + status = StatusError + errMsg = err.Error() + } + r := models.NotificationRequestV2{ + Status: status, + Error: errMsg, + SettingId: s.Id, + ChannelId: ch.Id, + Title: title, + Content: content, + } + r.SetCreatedAt(time.Now()) + r.SetUpdatedAt(time.Now()) + _, err = service.NewModelServiceV2[models.NotificationRequestV2]().InsertOne(r) + if err != nil { + log.Errorf("[NotificationServiceV2] save request error: %v", err) + } +} + func newNotificationServiceV2() *ServiceV2 { return &ServiceV2{} }