crawlab/backend/routes/system_tasks.go

package routes

import (
	"crawlab/constants"
	"crawlab/model"
	"crawlab/services"
	"crawlab/utils"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/globalsign/mgo/bson"
	"net/http"
)

func GetSystemScripts(c *gin.Context) {
	HandleSuccessData(c, utils.GetSystemScripts())
}

func PutSystemTask(c *gin.Context) {
	type TaskRequestBody struct {
		RunType string          `json:"run_type"`
		NodeIds []bson.ObjectId `json:"node_ids"`
		Script  string          `json:"script"`
	}

	// 绑定数据
	var reqBody TaskRequestBody
	if err := c.ShouldBindJSON(&reqBody); err != nil {
		HandleError(http.StatusBadRequest, c, err)
		return
	}

	// 校验脚本参数不为空
	if reqBody.Script == "" {
		HandleErrorF(http.StatusBadRequest, c, "script cannot be empty")
		return
	}

	// 校验脚本参数是否存在
	var allScripts = utils.GetSystemScripts()
	if !utils.StringArrayContains(allScripts, reqBody.Script) {
		HandleErrorF(http.StatusBadRequest, c, "script does not exist")
		return
	}

	// TODO: 校验脚本是否正在运行

	// 获取执行命令
	cmd := fmt.Sprintf("sh %s", utils.GetSystemScriptPath(reqBody.Script))

	// 任务ID
	var taskIds []string

	if reqBody.RunType == constants.RunTypeAllNodes {
		// 所有节点
		nodes, err := model.GetNodeList(nil)
		if err != nil {
			HandleError(http.StatusInternalServerError, c, err)
			return
		}
		for _, node := range nodes {
			t := model.Task{
				SpiderId:   bson.ObjectIdHex(constants.ObjectIdNull),
				NodeId:     node.Id,
				UserId:     services.GetCurrentUserId(c),
				RunType:    constants.RunTypeAllNodes,
				ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
				Type:       constants.TaskTypeSystem,
				Cmd:        cmd,
			}
			id, err := services.AddTask(t)
			if err != nil {
				HandleError(http.StatusInternalServerError, c, err)
				return
			}
			taskIds = append(taskIds, id)
		}
	} else if reqBody.RunType == constants.RunTypeRandom {
		// 随机
		t := model.Task{
			SpiderId:   bson.ObjectIdHex(constants.ObjectIdNull),
			UserId:     services.GetCurrentUserId(c),
			RunType:    constants.RunTypeRandom,
			ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
			Type:       constants.TaskTypeSystem,
			Cmd:        cmd,
		}
		id, err := services.AddTask(t)
		if err != nil {
			HandleError(http.StatusInternalServerError, c, err)
			return
		}
		taskIds = append(taskIds, id)
	} else if reqBody.RunType == constants.RunTypeSelectedNodes {
		// 指定节点
		for _, nodeId := range reqBody.NodeIds {
			t := model.Task{
				SpiderId:   bson.ObjectIdHex(constants.ObjectIdNull),
				NodeId:     nodeId,
				UserId:     services.GetCurrentUserId(c),
				RunType:    constants.RunTypeSelectedNodes,
				ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
				Type:       constants.TaskTypeSystem,
				Cmd:        cmd,
			}
			id, err := services.AddTask(t)
			if err != nil {
				HandleError(http.StatusInternalServerError, c, err)
				return
			}
			taskIds = append(taskIds, id)
		}
	} else {
		HandleErrorF(http.StatusInternalServerError, c, "invalid run_type")
		return
	}

	HandleSuccessData(c, taskIds)
}