feat: added modules

This commit is contained in:
Marvin Zhang
2024-06-14 15:42:50 +08:00
parent 4d0adcb6f0
commit c4d795f47f
626 changed files with 60104 additions and 0 deletions

View File

@@ -0,0 +1,25 @@
package delegate_test
import (
"context"
"github.com/crawlab-team/crawlab-db/mongo"
"go.mongodb.org/mongo-driver/bson"
"testing"
"time"
)
func SetupTest(t *testing.T) {
CleanupTest()
t.Cleanup(CleanupTest)
}
func CleanupTest() {
db := mongo.GetMongoDb("")
names, _ := db.ListCollectionNames(context.Background(), bson.M{})
for _, n := range names {
_, _ = db.Collection(n).DeleteMany(context.Background(), bson.M{})
}
// avoid caching
time.Sleep(200 * time.Millisecond)
}

View File

@@ -0,0 +1,362 @@
package delegate
import (
"encoding/json"
"github.com/crawlab-team/crawlab-db/errors"
"github.com/crawlab-team/crawlab-db/mongo"
errors2 "github.com/crawlab-team/crawlab/core/errors"
"github.com/crawlab-team/crawlab/core/event"
"github.com/crawlab-team/crawlab/core/interfaces"
"github.com/crawlab-team/crawlab/core/models/models"
"github.com/crawlab-team/crawlab/core/utils"
"github.com/crawlab-team/go-trace"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/primitive"
mongo2 "go.mongodb.org/mongo-driver/mongo"
"reflect"
"time"
)
func NewModelDelegate(doc interfaces.Model, args ...interface{}) interfaces.ModelDelegate {
switch doc.(type) {
case *models.Artifact:
return newModelDelegate(interfaces.ModelIdArtifact, doc, args...)
case *models.Tag:
return newModelDelegate(interfaces.ModelIdTag, doc, args...)
case *models.Node:
return newModelDelegate(interfaces.ModelIdNode, doc, args...)
case *models.Project:
return newModelDelegate(interfaces.ModelIdProject, doc, args...)
case *models.Spider:
return newModelDelegate(interfaces.ModelIdSpider, doc, args...)
case *models.Task:
return newModelDelegate(interfaces.ModelIdTask, doc, args...)
case *models.Job:
return newModelDelegate(interfaces.ModelIdJob, doc, args...)
case *models.Schedule:
return newModelDelegate(interfaces.ModelIdSchedule, doc, args...)
case *models.User:
return newModelDelegate(interfaces.ModelIdUser, doc, args...)
case *models.Setting:
return newModelDelegate(interfaces.ModelIdSetting, doc, args...)
case *models.Token:
return newModelDelegate(interfaces.ModelIdToken, doc, args...)
case *models.Variable:
return newModelDelegate(interfaces.ModelIdVariable, doc, args...)
case *models.TaskQueueItem:
return newModelDelegate(interfaces.ModelIdTaskQueue, doc, args...)
case *models.TaskStat:
return newModelDelegate(interfaces.ModelIdTaskStat, doc, args...)
case *models.SpiderStat:
return newModelDelegate(interfaces.ModelIdSpiderStat, doc, args...)
case *models.DataSource:
return newModelDelegate(interfaces.ModelIdDataSource, doc, args...)
case *models.DataCollection:
return newModelDelegate(interfaces.ModelIdDataCollection, doc, args...)
case *models.Result:
return newModelDelegate(interfaces.ModelIdResult, doc, args...)
case *models.Password:
return newModelDelegate(interfaces.ModelIdPassword, doc, args...)
case *models.ExtraValue:
return newModelDelegate(interfaces.ModelIdExtraValue, doc, args...)
case *models.Git:
return newModelDelegate(interfaces.ModelIdGit, doc, args...)
case *models.Role:
return newModelDelegate(interfaces.ModelIdRole, doc, args...)
case *models.UserRole:
return newModelDelegate(interfaces.ModelIdUserRole, doc, args...)
case *models.Permission:
return newModelDelegate(interfaces.ModelIdPermission, doc, args...)
case *models.RolePermission:
return newModelDelegate(interfaces.ModelIdRolePermission, doc, args...)
case *models.Environment:
return newModelDelegate(interfaces.ModelIdEnvironment, doc, args...)
case *models.DependencySetting:
return newModelDelegate(interfaces.ModelIdDependencySetting, doc, args...)
default:
_ = trace.TraceError(errors2.ErrorModelInvalidType)
return nil
}
}
func newModelDelegate(id interfaces.ModelId, doc interfaces.Model, args ...interface{}) interfaces.ModelDelegate {
// user
u := utils.GetUserFromArgs(args...)
// collection name
colName := models.GetModelColName(id)
// model delegate
d := &ModelDelegate{
id: id,
colName: colName,
doc: doc,
a: &models.Artifact{
Col: colName,
},
u: u,
}
return d
}
type ModelDelegate struct {
id interfaces.ModelId
colName string
doc interfaces.Model // doc to delegate
cd bson.M // current doc
od bson.M // original doc
a interfaces.ModelArtifact // artifact
u interfaces.User // user
}
// Add model
func (d *ModelDelegate) Add() (err error) {
return d.do(interfaces.ModelDelegateMethodAdd)
}
// Save model
func (d *ModelDelegate) Save() (err error) {
return d.do(interfaces.ModelDelegateMethodSave)
}
// Delete model
func (d *ModelDelegate) Delete() (err error) {
return d.do(interfaces.ModelDelegateMethodDelete)
}
// GetArtifact refresh artifact and return it
func (d *ModelDelegate) GetArtifact() (res interfaces.ModelArtifact, err error) {
if err := d.do(interfaces.ModelDelegateMethodGetArtifact); err != nil {
return nil, err
}
return d.a, nil
}
// Refresh model
func (d *ModelDelegate) Refresh() (err error) {
return d.refresh()
}
// GetModel return model
func (d *ModelDelegate) GetModel() (res interfaces.Model) {
return d.doc
}
func (d *ModelDelegate) ToBytes(m interface{}) (bytes []byte, err error) {
if m != nil {
return utils.JsonToBytes(m)
}
return json.Marshal(d.doc)
}
// do action given the model delegate method
func (d *ModelDelegate) do(method interfaces.ModelDelegateMethod) (err error) {
switch method {
case interfaces.ModelDelegateMethodAdd:
err = d.add()
case interfaces.ModelDelegateMethodSave:
err = d.save()
case interfaces.ModelDelegateMethodDelete:
err = d.delete()
case interfaces.ModelDelegateMethodGetArtifact, interfaces.ModelDelegateMethodRefresh:
err = d.refresh()
default:
return trace.TraceError(errors2.ErrorModelInvalidType)
}
if err != nil {
return err
}
// trigger event
eventName := GetEventName(d, method)
go event.SendEvent(eventName, d.doc)
return nil
}
// add model
func (d *ModelDelegate) add() (err error) {
if d.doc == nil {
return trace.TraceError(errors.ErrMissingValue)
}
if d.doc.GetId().IsZero() {
d.doc.SetId(primitive.NewObjectID())
}
col := mongo.GetMongoCol(d.colName)
if _, err = col.Insert(d.doc); err != nil {
return trace.TraceError(err)
}
if err := d.upsertArtifact(); err != nil {
return trace.TraceError(err)
}
return d.refresh()
}
// save model
func (d *ModelDelegate) save() (err error) {
// validate
if d.doc == nil || d.doc.GetId().IsZero() {
return trace.TraceError(errors.ErrMissingValue)
}
// collection
col := mongo.GetMongoCol(d.colName)
// current doc
docData, err := bson.Marshal(d.doc)
if err != nil {
trace.PrintError(err)
} else {
if err := bson.Unmarshal(docData, &d.cd); err != nil {
trace.PrintError(err)
}
}
// original doc
if err := col.FindId(d.doc.GetId()).One(&d.od); err != nil {
trace.PrintError(err)
}
// replace
if err := col.ReplaceId(d.doc.GetId(), d.doc); err != nil {
return trace.TraceError(err)
}
// upsert artifact
if err := d.upsertArtifact(); err != nil {
return trace.TraceError(err)
}
return d.refresh()
}
// delete model
func (d *ModelDelegate) delete() (err error) {
if d.doc.GetId().IsZero() {
return trace.TraceError(errors2.ErrorModelMissingId)
}
col := mongo.GetMongoCol(d.colName)
if err := col.FindId(d.doc.GetId()).One(d.doc); err != nil {
return trace.TraceError(err)
}
if err := col.DeleteId(d.doc.GetId()); err != nil {
return trace.TraceError(err)
}
return d.deleteArtifact()
}
// refresh model and artifact
func (d *ModelDelegate) refresh() (err error) {
if d.doc.GetId().IsZero() {
return trace.TraceError(errors2.ErrorModelMissingId)
}
col := mongo.GetMongoCol(d.colName)
fr := col.FindId(d.doc.GetId())
if err := fr.One(d.doc); err != nil {
return trace.TraceError(err)
}
return d.refreshArtifact()
}
// refresh artifact
func (d *ModelDelegate) refreshArtifact() (err error) {
if d.doc.GetId().IsZero() {
return trace.TraceError(errors2.ErrorModelMissingId)
}
col := mongo.GetMongoCol(interfaces.ModelColNameArtifact)
if err := col.FindId(d.doc.GetId()).One(d.a); err != nil {
return trace.TraceError(err)
}
return nil
}
// upsertArtifact
func (d *ModelDelegate) upsertArtifact() (err error) {
// skip
if d._skip() {
return nil
}
// validate id
if d.doc.GetId().IsZero() {
return trace.TraceError(errors.ErrMissingValue)
}
// mongo collection
col := mongo.GetMongoCol(interfaces.ModelColNameArtifact)
// assign id to artifact
d.a.SetId(d.doc.GetId())
// attempt to find artifact
if err := col.FindId(d.doc.GetId()).One(d.a); err != nil {
if err == mongo2.ErrNoDocuments {
// new artifact
d.a.GetSys().SetCreateTs(time.Now())
d.a.GetSys().SetUpdateTs(time.Now())
if d.u != nil && !reflect.ValueOf(d.u).IsZero() {
d.a.GetSys().SetCreateUid(d.u.GetId())
d.a.GetSys().SetUpdateUid(d.u.GetId())
}
_, err = col.Insert(d.a)
if err != nil {
return trace.TraceError(err)
}
return nil
} else {
// error
return trace.TraceError(err)
}
}
// existing artifact
d.a.GetSys().SetUpdateTs(time.Now())
if d.u != nil {
d.a.GetSys().SetUpdateUid(d.u.GetId())
}
// save new artifact
return col.ReplaceId(d.a.GetId(), d.a)
}
// deleteArtifact
func (d *ModelDelegate) deleteArtifact() (err error) {
// skip
if d._skip() {
return nil
}
if d.doc.GetId().IsZero() {
return trace.TraceError(errors.ErrMissingValue)
}
col := mongo.GetMongoCol(interfaces.ModelColNameArtifact)
d.a.SetId(d.doc.GetId())
d.a.SetObj(d.doc)
d.a.SetDel(true)
d.a.GetSys().SetDeleteTs(time.Now())
if d.u != nil {
d.a.GetSys().SetDeleteUid(d.u.GetId())
}
return col.ReplaceId(d.doc.GetId(), d.a)
}
func (d *ModelDelegate) hasChange() (ok bool) {
return !utils.BsonMEqual(d.cd, d.od)
}
func (d *ModelDelegate) _skip() (ok bool) {
switch d.id {
case
interfaces.ModelIdArtifact,
interfaces.ModelIdTaskQueue,
interfaces.ModelIdTaskStat,
interfaces.ModelIdSpiderStat,
interfaces.ModelIdResult,
interfaces.ModelIdPassword:
return true
default:
return false
}
}

View File

@@ -0,0 +1,37 @@
package delegate
import (
"github.com/crawlab-team/crawlab/core/constants"
"github.com/crawlab-team/crawlab/core/interfaces"
"time"
)
type ModelNodeDelegate struct {
n interfaces.Node
interfaces.ModelDelegate
}
func (d *ModelNodeDelegate) UpdateStatus(active bool, activeTs *time.Time, status string) (err error) {
d.n.SetActive(active)
if activeTs != nil {
d.n.SetActiveTs(*activeTs)
}
d.n.SetStatus(status)
return d.Save()
}
func (d *ModelNodeDelegate) UpdateStatusOnline() (err error) {
now := time.Now()
return d.UpdateStatus(true, &now, constants.NodeStatusOnline)
}
func (d *ModelNodeDelegate) UpdateStatusOffline() (err error) {
return d.UpdateStatus(false, nil, constants.NodeStatusOffline)
}
func NewModelNodeDelegate(n interfaces.Node) interfaces.ModelNodeDelegate {
return &ModelNodeDelegate{
n: n,
ModelDelegate: NewModelDelegate(n),
}
}

View File

@@ -0,0 +1,65 @@
package delegate_test
import (
"github.com/crawlab-team/crawlab-db/mongo"
"github.com/crawlab-team/crawlab/core/interfaces"
"github.com/crawlab-team/crawlab/core/models/delegate"
models2 "github.com/crawlab-team/crawlab/core/models/models"
"github.com/stretchr/testify/require"
"testing"
)
func TestNode_Add(t *testing.T) {
SetupTest(t)
n := &models2.Node{}
err := delegate.NewModelDelegate(n).Add()
require.Nil(t, err)
require.NotNil(t, n.Id)
// validate artifact
a, err := delegate.NewModelDelegate(n).GetArtifact()
require.Nil(t, err)
require.Equal(t, n.Id, a.GetId())
require.NotNil(t, a.GetSys().GetCreateTs())
require.NotNil(t, a.GetSys().GetUpdateTs())
}
func TestNode_Save(t *testing.T) {
SetupTest(t)
n := &models2.Node{}
err := delegate.NewModelDelegate(n).Add()
name := "test_node"
n.Name = name
err = delegate.NewModelDelegate(n).Save()
require.Nil(t, err)
err = mongo.GetMongoCol(interfaces.ModelColNameNode).FindId(n.Id).One(&n)
require.Nil(t, err)
require.Equal(t, name, n.Name)
}
func TestNode_Delete(t *testing.T) {
SetupTest(t)
n := &models2.Node{
Name: "test_node",
}
err := delegate.NewModelDelegate(n).Add()
require.Nil(t, err)
err = delegate.NewModelDelegate(n).Delete()
require.Nil(t, err)
var a models2.Artifact
col := mongo.GetMongoCol(interfaces.ModelColNameArtifact)
err = col.FindId(n.Id).One(&a)
require.Nil(t, err)
require.NotNil(t, a.Obj)
require.True(t, a.Del)
}

View File

@@ -0,0 +1,65 @@
package delegate_test
import (
"github.com/crawlab-team/crawlab-db/mongo"
"github.com/crawlab-team/crawlab/core/interfaces"
"github.com/crawlab-team/crawlab/core/models/delegate"
models2 "github.com/crawlab-team/crawlab/core/models/models"
"github.com/stretchr/testify/require"
"testing"
)
func TestRole_Add(t *testing.T) {
SetupTest(t)
p := &models2.Role{}
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 TestRole_Save(t *testing.T) {
SetupTest(t)
p := &models2.Role{}
err := delegate.NewModelDelegate(p).Add()
require.Nil(t, err)
name := "test_role"
p.Name = name
err = delegate.NewModelDelegate(p).Save()
require.Nil(t, err)
err = mongo.GetMongoCol(interfaces.ModelColNameRole).FindId(p.Id).One(&p)
require.Nil(t, err)
require.Equal(t, name, p.Name)
}
func TestRole_Delete(t *testing.T) {
SetupTest(t)
p := &models2.Role{
Name: "test_role",
}
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)
}

View File

@@ -0,0 +1,65 @@
package delegate_test
import (
"github.com/crawlab-team/crawlab-db/mongo"
"github.com/crawlab-team/crawlab/core/interfaces"
"github.com/crawlab-team/crawlab/core/models/delegate"
models2 "github.com/crawlab-team/crawlab/core/models/models"
"github.com/stretchr/testify/require"
"testing"
)
func TestProject_Add(t *testing.T) {
SetupTest(t)
p := &models2.Project{}
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 TestProject_Save(t *testing.T) {
SetupTest(t)
p := &models2.Project{}
err := delegate.NewModelDelegate(p).Add()
require.Nil(t, err)
name := "test_project"
p.Name = name
err = delegate.NewModelDelegate(p).Save()
require.Nil(t, err)
err = mongo.GetMongoCol(interfaces.ModelColNameProject).FindId(p.Id).One(&p)
require.Nil(t, err)
require.Equal(t, name, p.Name)
}
func TestProject_Delete(t *testing.T) {
SetupTest(t)
p := &models2.Project{
Name: "test_project",
}
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)
}

View File

@@ -0,0 +1,97 @@
package delegate_test
import (
"github.com/crawlab-team/crawlab-db/mongo"
"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/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

@@ -0,0 +1,20 @@
package delegate
import (
"fmt"
"github.com/crawlab-team/crawlab/core/interfaces"
)
func GetEventName(d *ModelDelegate, method interfaces.ModelDelegateMethod) (eventName string) {
return getEventName(d, method)
}
func getEventName(d *ModelDelegate, method interfaces.ModelDelegateMethod) (eventName string) {
if method == interfaces.ModelDelegateMethodSave {
hasChange := d.hasChange()
if hasChange {
method = interfaces.ModelDelegateMethodChange
}
}
return fmt.Sprintf("model:%s:%s", d.colName, method)
}