refactor: Update NotificationChannelV2 model to include Telegram notification settings

This commit is contained in:
Marvin Zhang
2024-07-25 22:58:39 +08:00
parent eef3eea777
commit 96b89d05b8
7 changed files with 81 additions and 268 deletions

View File

@@ -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),
},
})
}

View File

@@ -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},
},
},
})
}

View File

@@ -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)
}

View File

@@ -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:"-"`
}

View File

@@ -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"
)

View File

@@ -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"`

View File

@@ -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{}
}