mirror of
https://github.com/crawlab-team/crawlab.git
synced 2026-01-22 17:31:03 +01:00
612 lines
15 KiB
Go
612 lines
15 KiB
Go
package routes
|
||
|
||
import (
|
||
"bytes"
|
||
"crawlab/constants"
|
||
"crawlab/model"
|
||
"crawlab/services"
|
||
"crawlab/utils"
|
||
"encoding/csv"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/globalsign/mgo/bson"
|
||
"net/http"
|
||
)
|
||
|
||
type TaskListRequestData struct {
|
||
PageNum int `form:"page_num"`
|
||
PageSize int `form:"page_size"`
|
||
NodeId string `form:"node_id"`
|
||
SpiderId string `form:"spider_id"`
|
||
ScheduleId string `form:"schedule_id"`
|
||
Status string `form:"status"`
|
||
}
|
||
|
||
type TaskResultsRequestData struct {
|
||
PageNum int `form:"page_num"`
|
||
PageSize int `form:"page_size"`
|
||
}
|
||
|
||
// @Summary Get task list
|
||
// @Description Get task list
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param data body routes.TaskListRequestData true "req data"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks [get]
|
||
func GetTaskList(c *gin.Context) {
|
||
// 绑定数据
|
||
data := TaskListRequestData{}
|
||
if err := c.ShouldBindQuery(&data); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
if data.PageNum == 0 {
|
||
data.PageNum = 1
|
||
}
|
||
if data.PageSize == 0 {
|
||
data.PageSize = 10
|
||
}
|
||
|
||
// 过滤条件
|
||
query := bson.M{}
|
||
if data.NodeId != "" {
|
||
query["node_id"] = bson.ObjectIdHex(data.NodeId)
|
||
}
|
||
if data.SpiderId != "" {
|
||
query["spider_id"] = bson.ObjectIdHex(data.SpiderId)
|
||
}
|
||
// 根据任务状态获取task列表
|
||
if data.Status != "" {
|
||
query["status"] = data.Status
|
||
}
|
||
if data.ScheduleId != "" {
|
||
query["schedule_id"] = bson.ObjectIdHex(data.ScheduleId)
|
||
}
|
||
|
||
// 获取校验
|
||
query = services.GetAuthQuery(query, c)
|
||
|
||
// 获取任务列表
|
||
tasks, err := model.GetTaskList(query, (data.PageNum-1)*data.PageSize, data.PageSize, "-create_ts")
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
// 获取总任务数
|
||
total, err := model.GetTaskListTotal(query)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
c.JSON(http.StatusOK, ListResponse{
|
||
Status: "ok",
|
||
Message: "success",
|
||
Total: total,
|
||
Data: tasks,
|
||
})
|
||
}
|
||
|
||
// @Summary Get task
|
||
// @Description Get task
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param id path string true "task id"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks/{id} [get]
|
||
func GetTask(c *gin.Context) {
|
||
id := c.Param("id")
|
||
result, err := model.GetTask(id)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
HandleSuccessData(c, result)
|
||
}
|
||
|
||
// @Summary Put task
|
||
// @Description Put task
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks [put]
|
||
func PutTask(c *gin.Context) {
|
||
type TaskRequestBody struct {
|
||
SpiderId bson.ObjectId `json:"spider_id"`
|
||
RunType string `json:"run_type"`
|
||
NodeIds []bson.ObjectId `json:"node_ids"`
|
||
Param string `json:"param"`
|
||
}
|
||
|
||
// 绑定数据
|
||
var reqBody TaskRequestBody
|
||
if err := c.ShouldBindJSON(&reqBody); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
// 任务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: reqBody.SpiderId,
|
||
NodeId: node.Id,
|
||
Param: reqBody.Param,
|
||
UserId: services.GetCurrentUserId(c),
|
||
RunType: constants.RunTypeAllNodes,
|
||
ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
|
||
}
|
||
|
||
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: reqBody.SpiderId,
|
||
Param: reqBody.Param,
|
||
UserId: services.GetCurrentUserId(c),
|
||
RunType: constants.RunTypeRandom,
|
||
ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
|
||
}
|
||
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: reqBody.SpiderId,
|
||
NodeId: nodeId,
|
||
Param: reqBody.Param,
|
||
UserId: services.GetCurrentUserId(c),
|
||
RunType: constants.RunTypeSelectedNodes,
|
||
ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
|
||
}
|
||
|
||
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)
|
||
}
|
||
|
||
func PutBatchTasks(c *gin.Context) {
|
||
var tasks []model.Task
|
||
if err := c.ShouldBindJSON(&tasks); err != nil {
|
||
HandleError(http.StatusOK, c, err)
|
||
return
|
||
}
|
||
var taskIds []string
|
||
for _, t := range tasks {
|
||
if t.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: t.SpiderId,
|
||
NodeId: node.Id,
|
||
Param: t.Param,
|
||
UserId: services.GetCurrentUserId(c),
|
||
RunType: constants.RunTypeAllNodes,
|
||
ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
|
||
}
|
||
|
||
id, err := services.AddTask(t)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
taskIds = append(taskIds, id)
|
||
}
|
||
} else if t.RunType == constants.RunTypeRandom {
|
||
// 随机
|
||
t := model.Task{
|
||
SpiderId: t.SpiderId,
|
||
Param: t.Param,
|
||
UserId: services.GetCurrentUserId(c),
|
||
RunType: constants.RunTypeRandom,
|
||
ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
|
||
}
|
||
id, err := services.AddTask(t)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
taskIds = append(taskIds, id)
|
||
} else if t.RunType == constants.RunTypeSelectedNodes {
|
||
// 指定节点
|
||
for _, nodeId := range t.NodeIds {
|
||
t := model.Task{
|
||
SpiderId: t.SpiderId,
|
||
NodeId: bson.ObjectIdHex(nodeId),
|
||
Param: t.Param,
|
||
UserId: services.GetCurrentUserId(c),
|
||
RunType: constants.RunTypeSelectedNodes,
|
||
ScheduleId: bson.ObjectIdHex(constants.ObjectIdNull),
|
||
}
|
||
|
||
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
|
||
}
|
||
}
|
||
|
||
c.JSON(http.StatusOK, Response{
|
||
Status: "ok",
|
||
Message: "success",
|
||
Data: taskIds,
|
||
})
|
||
}
|
||
|
||
// @Summary Delete task
|
||
// @Description Delete task
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param status query string true "task status"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks_by_status [delete]
|
||
func DeleteTaskByStatus(c *gin.Context) {
|
||
status := c.Query("status")
|
||
|
||
//删除相应的日志文件
|
||
if err := services.RemoveLogByTaskStatus(status); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
//删除该状态下的task
|
||
if err := model.RemoveTaskByStatus(status); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
HandleSuccess(c)
|
||
}
|
||
|
||
// 删除多个任务
|
||
|
||
// @Summary Delete tasks
|
||
// @Description Delete tasks
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks [delete]
|
||
func DeleteSelectedTask(c *gin.Context) {
|
||
ids := make(map[string][]string)
|
||
if err := c.ShouldBindJSON(&ids); err != nil {
|
||
HandleError(http.StatusBadRequest, c, err)
|
||
return
|
||
}
|
||
list := ids["ids"]
|
||
for _, id := range list {
|
||
if err := services.RemoveLogByTaskId(id); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
if err := model.RemoveTask(id); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
}
|
||
HandleSuccess(c)
|
||
}
|
||
|
||
// 删除单个任务
|
||
|
||
// @Summary Delete task
|
||
// @Description Delete task
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param id path string true "task id"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /task/{id} [delete]
|
||
func DeleteTask(c *gin.Context) {
|
||
id := c.Param("id")
|
||
// 删除日志文件
|
||
if err := services.RemoveLogByTaskId(id); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
// 删除task
|
||
if err := model.RemoveTask(id); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
HandleSuccess(c)
|
||
}
|
||
|
||
func CancelSelectedTask(c *gin.Context) {
|
||
ids := make(map[string][]string)
|
||
if err := c.ShouldBindJSON(&ids); err != nil {
|
||
HandleError(http.StatusBadRequest, c, err)
|
||
return
|
||
}
|
||
list := ids["ids"]
|
||
for _, id := range list {
|
||
if err := services.CancelTask(id); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
}
|
||
HandleSuccess(c)
|
||
}
|
||
|
||
func RestartSelectedTask(c *gin.Context) {
|
||
ids := make(map[string][]string)
|
||
if err := c.ShouldBindJSON(&ids); err != nil {
|
||
HandleError(http.StatusBadRequest, c, err)
|
||
return
|
||
}
|
||
list := ids["ids"]
|
||
for _, id := range list {
|
||
if err := services.RestartTask(id, services.GetCurrentUserId(c)); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
}
|
||
HandleSuccess(c)
|
||
}
|
||
|
||
// @Summary Get task log
|
||
// @Description Get task log
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param id path string true "task id"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks/{id}/log [delete]
|
||
func GetTaskLog(c *gin.Context) {
|
||
type RequestData struct {
|
||
PageNum int `form:"page_num"`
|
||
PageSize int `form:"page_size"`
|
||
Keyword string `form:"keyword"`
|
||
}
|
||
id := c.Param("id")
|
||
var reqData RequestData
|
||
if err := c.ShouldBindQuery(&reqData); err != nil {
|
||
HandleErrorF(http.StatusBadRequest, c, "invalid request")
|
||
return
|
||
}
|
||
logItems, logTotal, err := services.GetTaskLog(id, reqData.Keyword, reqData.PageNum, reqData.PageSize)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, ListResponse{
|
||
Status: "ok",
|
||
Message: "success",
|
||
Data: logItems,
|
||
Total: logTotal,
|
||
})
|
||
}
|
||
|
||
// @Summary Get task error log
|
||
// @Description Get task error log
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param id path string true "task id"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks/{id}/error-log [delete]
|
||
func GetTaskErrorLog(c *gin.Context) {
|
||
id := c.Param("id")
|
||
u := services.GetCurrentUser(c)
|
||
errLogItems, err := services.GetTaskErrorLog(id, u.Setting.MaxErrorLog)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, Response{
|
||
Status: "ok",
|
||
Message: "success",
|
||
Data: errLogItems,
|
||
})
|
||
}
|
||
|
||
// @Summary Get task list
|
||
// @Description Get task list
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param data body routes.TaskResultsRequestData true "req data"
|
||
// @Param id path string true "task id"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks/{id}/results [get]
|
||
func GetTaskResults(c *gin.Context) {
|
||
id := c.Param("id")
|
||
// 绑定数据
|
||
data := TaskResultsRequestData{}
|
||
if err := c.ShouldBindQuery(&data); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
// 获取任务
|
||
task, err := model.GetTask(id)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
// 获取结果
|
||
results, total, err := task.GetResults(data.PageNum, data.PageSize)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
c.JSON(http.StatusOK, ListResponse{
|
||
Status: "ok",
|
||
Message: "success",
|
||
Data: results,
|
||
Total: total,
|
||
})
|
||
}
|
||
|
||
// @Summary Get task results
|
||
// @Description Get task results
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param id path string true "task id"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks/{id}/results/download [get]
|
||
func DownloadTaskResultsCsv(c *gin.Context) {
|
||
id := c.Param("id")
|
||
// 获取任务
|
||
task, err := model.GetTask(id)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
// 获取结果
|
||
results, _, err := task.GetResults(1, constants.Infinite)
|
||
if err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
// 字段列表
|
||
var columns []string
|
||
if len(results) == 0 {
|
||
columns = []string{}
|
||
} else {
|
||
item := results[0].(bson.M)
|
||
for key := range item {
|
||
columns = append(columns, key)
|
||
}
|
||
}
|
||
|
||
// 缓冲
|
||
bytesBuffer := &bytes.Buffer{}
|
||
|
||
// 写入UTF-8 BOM，避免使用Microsoft Excel打开乱码
|
||
bytesBuffer.WriteString("\xEF\xBB\xBF")
|
||
|
||
writer := csv.NewWriter(bytesBuffer)
|
||
|
||
// 写入表头
|
||
if err := writer.Write(columns); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
|
||
// 写入内容
|
||
for _, result := range results {
|
||
// 将result转换为[]string
|
||
item := result.(bson.M)
|
||
var values []string
|
||
for _, col := range columns {
|
||
value := utils.InterfaceToString(item[col])
|
||
values = append(values, value)
|
||
}
|
||
|
||
// 写入数据
|
||
if err := writer.Write(values); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
}
|
||
|
||
// 此时才会将缓冲区数据写入
|
||
writer.Flush()
|
||
|
||
// 设置下载的文件名
|
||
c.Writer.Header().Set("Content-Disposition", "attachment;filename=data.csv")
|
||
|
||
// 设置文件类型以及输出数据
|
||
c.Data(http.StatusOK, "text/csv", bytesBuffer.Bytes())
|
||
}
|
||
|
||
// @Summary Cancel task
|
||
// @Description Cancel task
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param id path string true "task id"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks/{id}/cancel [post]
|
||
func CancelTask(c *gin.Context) {
|
||
id := c.Param("id")
|
||
if err := services.CancelTask(id); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
HandleSuccess(c)
|
||
}
|
||
|
||
// @Summary Restart task
|
||
// @Description Restart task
|
||
// @Tags task
|
||
// @Produce json
|
||
// @Param Authorization header string true "Authorization token"
|
||
// @Param id path string true "task id"
|
||
// @Success 200 json string Response
|
||
// @Failure 400 json string Response
|
||
// @Router /tasks/{id}/restart [post]
|
||
func RestartTask(c *gin.Context) {
|
||
id := c.Param("id")
|
||
uid := services.GetCurrentUserId(c)
|
||
|
||
if err := services.RestartTask(id, uid); err != nil {
|
||
HandleError(http.StatusInternalServerError, c, err)
|
||
return
|
||
}
|
||
HandleSuccess(c)
|
||
}
|