package controllers import ( errors2 "errors" "github.com/crawlab-team/crawlab/core/interfaces" "github.com/crawlab-team/crawlab/core/models/models" "github.com/crawlab-team/crawlab/core/models/service" "github.com/crawlab-team/crawlab/core/schedule" "github.com/crawlab-team/crawlab/core/spider/admin" "github.com/gin-gonic/gin" "github.com/juju/errors" "go.mongodb.org/mongo-driver/bson/primitive" ) type PostScheduleParams struct { Data models.Schedule `json:"data"` } func PostSchedule(c *gin.Context, params *PostScheduleParams) (response *Response[models.Schedule], err error) { s := params.Data u := GetUserFromContext(c) modelSvc := service.NewModelService[models.Schedule]() s.SetCreated(u.Id) s.SetUpdated(u.Id) id, err := modelSvc.InsertOne(s) if err != nil { return GetErrorResponse[models.Schedule](err) } s.Id = id if s.Enabled { scheduleSvc := schedule.GetScheduleService() if err := scheduleSvc.Enable(s, u.Id); err != nil { return GetErrorResponse[models.Schedule](err) } } return GetDataResponse(s) } type PutScheduleByIdParams struct { Id string `path:"id"` Data models.Schedule `json:"data"` } func PutScheduleById(c *gin.Context, params *PutScheduleByIdParams) (response *Response[models.Schedule], err error) { id, err := primitive.ObjectIDFromHex(params.Id) if err != nil { return GetErrorResponse[models.Schedule](errors.BadRequestf("invalid schedule id: %v", err)) } s := params.Data if s.Id != id { return GetErrorResponse[models.Schedule](errors2.New("id in path does not match id in body")) } modelSvc := service.NewModelService[models.Schedule]() err = modelSvc.ReplaceById(id, s) if err != nil { return GetErrorResponse[models.Schedule](err) } scheduleSvc := schedule.GetScheduleService() u := GetUserFromContext(c) if s.Enabled { if err := scheduleSvc.Enable(s, u.Id); err != nil { return GetErrorResponse[models.Schedule](err) } } else { if err := scheduleSvc.Disable(s, u.Id); err != nil { return GetErrorResponse[models.Schedule](err) } } return GetDataResponse(s) } type PostScheduleEnableDisableParams struct { Id string `path:"id"` } func PostScheduleEnable(c *gin.Context, params *PostScheduleEnableDisableParams) (response *VoidResponse, err error) { userId := GetUserFromContext(c).Id return postScheduleEnableDisableFunc(true, userId, params) } func PostScheduleDisable(c *gin.Context, params *PostScheduleEnableDisableParams) (response *VoidResponse, err error) { userId := GetUserFromContext(c).Id return postScheduleEnableDisableFunc(false, userId, params) } func postScheduleEnableDisableFunc(isEnable bool, userId primitive.ObjectID, params *PostScheduleEnableDisableParams) (response *VoidResponse, err error) { id, err := primitive.ObjectIDFromHex(params.Id) if err != nil { return GetErrorVoidResponse(errors.BadRequestf("invalid schedule id: %v", err)) } svc := schedule.GetScheduleService() s, err := service.NewModelService[models.Schedule]().GetById(id) if err != nil { return GetErrorVoidResponse(err) } if isEnable { err = svc.Enable(*s, userId) } else { err = svc.Disable(*s, userId) } if err != nil { return GetErrorVoidResponse(err) } return GetVoidResponse() } type PostScheduleRunParams struct { Id string `path:"id"` Mode string `json:"mode"` NodeIds []string `json:"node_ids"` Cmd string `json:"cmd"` Param string `json:"param"` Priority int `json:"priority"` } func PostScheduleRun(c *gin.Context, params *PostScheduleRunParams) (response *Response[[]primitive.ObjectID], err error) { userId := GetUserFromContext(c).Id return postScheduleRunFunc(params, userId) } func postScheduleRunFunc(params *PostScheduleRunParams, userId primitive.ObjectID) (response *Response[[]primitive.ObjectID], err error) { id, err := primitive.ObjectIDFromHex(params.Id) if err != nil { return GetErrorResponse[[]primitive.ObjectID](errors.BadRequestf("invalid schedule id: %v", err)) } var nodeIds []primitive.ObjectID for _, nodeId := range params.NodeIds { nodeId, err := primitive.ObjectIDFromHex(nodeId) if err != nil { return GetErrorResponse[[]primitive.ObjectID](errors.BadRequestf("invalid node id: %v", err)) } nodeIds = append(nodeIds, nodeId) } opts := interfaces.SpiderRunOptions{ Mode: params.Mode, NodeIds: nodeIds, Cmd: params.Cmd, Param: params.Param, Priority: params.Priority, ScheduleId: id, UserId: userId, } // schedule sch, err := service.NewModelService[models.Schedule]().GetById(id) if err != nil { return GetErrorResponse[[]primitive.ObjectID](err) } // spider s, err := service.NewModelService[models.Spider]().GetById(sch.SpiderId) if err != nil { return GetErrorResponse[[]primitive.ObjectID](err) } // user opts.UserId = userId // schedule tasks taskIds, err := admin.GetSpiderAdminService().Schedule(s.Id, &opts) if err != nil { return GetErrorResponse[[]primitive.ObjectID](err) } return GetDataResponse(taskIds) }