crawlab/core/models/delegate/model.go

package delegate

import (
	"encoding/json"
	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/crawlab/db/errors"
	"github.com/crawlab-team/crawlab/db/mongo"
	"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
	}
}