package model import ( "crawlab/constants" "crawlab/database" "crawlab/entity" "crawlab/utils" "errors" "github.com/apex/log" "github.com/globalsign/mgo" "github.com/globalsign/mgo/bson" "gopkg.in/yaml.v2" "io/ioutil" "path/filepath" "runtime/debug" "time" ) type Env struct { Name string `json:"name" bson:"name"` Value string `json:"value" bson:"value"` } type Spider struct { Id bson.ObjectId `json:"_id" bson:"_id"` // 爬虫ID Name string `json:"name" bson:"name"` // 爬虫名称(唯一) DisplayName string `json:"display_name" bson:"display_name"` // 爬虫显示名称 Type string `json:"type" bson:"type"` // 爬虫类别 FileId bson.ObjectId `json:"file_id" bson:"file_id"` // GridFS文件ID Col string `json:"col" bson:"col"` // 结果储存位置 Site string `json:"site" bson:"site"` // 爬虫网站 Envs []Env `json:"envs" bson:"envs"` // 环境变量 Remark string `json:"remark" bson:"remark"` // 备注 Src string `json:"src" bson:"src"` // 源码位置 ProjectId bson.ObjectId `json:"project_id" bson:"project_id"` // 项目ID // 自定义爬虫 Cmd string `json:"cmd" bson:"cmd"` // 执行命令 // Scrapy 爬虫(属于自定义爬虫) IsScrapy bool `json:"is_scrapy" bson:"is_scrapy"` // 是否为 Scrapy 爬虫 SpiderNames []string `json:"spider_names" bson:"spider_names"` // 爬虫名称列表 // 可配置爬虫 Template string `json:"template" bson:"template"` // Spiderfile模版 // Git 设置 IsGit bool `json:"is_git" bson:"is_git"` // 是否为 Git GitUrl string `json:"git_url" bson:"git_url"` // Git URL GitBranch string `json:"git_branch" bson:"git_branch"` // Git 分支 GitHasCredential bool `json:"git_has_credential" bson:"git_has_credential"` // Git 是否加密 GitUsername string `json:"git_username" bson:"git_username"` // Git 用户名 GitPassword string `json:"git_password" bson:"git_password"` // Git 密码 GitAutoSync bool `json:"git_auto_sync" bson:"git_auto_sync"` // Git 是否自动同步 GitSyncFrequency string `json:"git_sync_frequency" bson:"git_sync_frequency"` // Git 同步频率 GitSyncError string `json:"git_sync_error" bson:"git_sync_error"` // Git 同步错误 // 长任务 IsLongTask bool `json:"is_long_task" bson:"is_long_task"` // 是否为长任务 // 前端展示 LastRunTs time.Time `json:"last_run_ts"` // 最后一次执行时间 LastStatus string `json:"last_status"` // 最后执行状态 Config entity.ConfigSpiderData `json:"config"` // 可配置爬虫配置 LatestTasks []Task `json:"latest_tasks"` // 最近任务列表 // 时间 CreateTs time.Time `json:"create_ts" bson:"create_ts"` UpdateTs time.Time `json:"update_ts" bson:"update_ts"` } // 更新爬虫 func (spider *Spider) Save() error { s, c := database.GetCol("spiders") defer s.Close() spider.UpdateTs = time.Now() // 兼容没有项目ID的爬虫 if spider.ProjectId.Hex() == "" { spider.ProjectId = bson.ObjectIdHex(constants.ObjectIdNull) } if err := c.UpdateId(spider.Id, spider); err != nil { debug.PrintStack() return err } return nil } // 新增爬虫 func (spider *Spider) Add() error { s, c := database.GetCol("spiders") defer s.Close() spider.Id = bson.NewObjectId() spider.CreateTs = time.Now() spider.UpdateTs = time.Now() if !spider.ProjectId.Valid() { spider.ProjectId = bson.ObjectIdHex(constants.ObjectIdNull) } if err := c.Insert(&spider); err != nil { return err } return nil } // 获取爬虫的任务 func (spider *Spider) GetTasks() ([]Task, error) { tasks, err := GetTaskList(bson.M{"spider_id": spider.Id}, 0, 10, "-create_ts") if err != nil { return tasks, err } return tasks, nil } // 爬虫最新的任务 func (spider *Spider) GetLastTask() (Task, error) { tasks, err := GetTaskList(bson.M{"spider_id": spider.Id}, 0, 1, "-create_ts") if err != nil { return Task{}, err } if tasks == nil { return Task{}, nil } return tasks[0], nil } // 爬虫正在运行的任务 func (spider *Spider) GetLatestTasks(latestN int) (tasks []Task, err error) { tasks, err = GetTaskList(bson.M{"spider_id": spider.Id}, 0, latestN, "-create_ts") if err != nil { return tasks, err } if tasks == nil { return tasks, err } return tasks, nil } // 删除爬虫 func (spider *Spider) Delete() error { s, c := database.GetCol("spiders") defer s.Close() return c.RemoveId(spider.Id) } // 获取爬虫列表 func GetSpiderList(filter interface{}, skip int, limit int, sortStr string) ([]Spider, int, error) { s, c := database.GetCol("spiders") defer s.Close() // 获取爬虫列表 var spiders []Spider if err := c.Find(filter).Skip(skip).Limit(limit).Sort(sortStr).All(&spiders); err != nil { debug.PrintStack() return spiders, 0, err } if spiders == nil { spiders = []Spider{} } // 遍历爬虫列表 for i, spider := range spiders { // 获取最后一次任务 task, err := spider.GetLastTask() if err != nil { log.Errorf(err.Error()) debug.PrintStack() continue } // 获取正在运行的爬虫 latestTasks, err := spider.GetLatestTasks(50) // TODO: latestN 暂时写死,后面加入数据库 if err != nil { log.Errorf(err.Error()) debug.PrintStack() continue } // 赋值 spiders[i].LastRunTs = task.CreateTs spiders[i].LastStatus = task.Status spiders[i].LatestTasks = latestTasks } count, _ := c.Find(filter).Count() return spiders, count, nil } // 获取所有爬虫列表 func GetSpiderAllList(filter interface{}) (spiders []Spider, err error) { spiders, _, err = GetSpiderList(filter, 0, constants.Infinite, "_id") if err != nil { return spiders, err } return spiders, nil } // 获取爬虫(根据FileId) func GetSpiderByFileId(fileId bson.ObjectId) *Spider { s, c := database.GetCol("spiders") defer s.Close() var result *Spider if err := c.Find(bson.M{"file_id": fileId}).One(&result); err != nil { log.Errorf("get spider error: %s, file_id: %s", err.Error(), fileId.Hex()) debug.PrintStack() return nil } return result } // 获取爬虫(根据名称) func GetSpiderByName(name string) Spider { s, c := database.GetCol("spiders") defer s.Close() var result Spider if err := c.Find(bson.M{"name": name}).One(&result); err != nil && err != mgo.ErrNotFound { log.Errorf("get spider error: %s, spider_name: %s", err.Error(), name) //debug.PrintStack() return result } return result } // 获取爬虫(根据ID) func GetSpider(id bson.ObjectId) (Spider, error) { s, c := database.GetCol("spiders") defer s.Close() // 获取爬虫 var spider Spider if err := c.FindId(id).One(&spider); err != nil { if err != mgo.ErrNotFound { log.Errorf("get spider error: %s, id: %id", err.Error(), id.Hex()) debug.PrintStack() } return spider, err } // 如果为可配置爬虫,获取爬虫配置 if spider.Type == constants.Configurable && utils.Exists(filepath.Join(spider.Src, "Spiderfile")) { config, err := GetConfigSpiderData(spider) if err != nil { return spider, err } spider.Config = config } return spider, nil } // 更新爬虫 func UpdateSpider(id bson.ObjectId, item Spider) error { s, c := database.GetCol("spiders") defer s.Close() var result Spider if err := c.FindId(id).One(&result); err != nil { debug.PrintStack() return err } if err := item.Save(); err != nil { return err } return nil } // 删除爬虫 func RemoveSpider(id bson.ObjectId) error { s, c := database.GetCol("spiders") defer s.Close() var result Spider if err := c.FindId(id).One(&result); err != nil { log.Errorf("find spider error: %s, id:%s", err.Error(), id.Hex()) debug.PrintStack() return err } if err := c.RemoveId(id); err != nil { log.Errorf("remove spider error: %s, id:%s", err.Error(), id.Hex()) debug.PrintStack() return err } // gf上的文件 s, gf := database.GetGridFs("files") defer s.Close() if result.FileId.Hex() != constants.ObjectIdNull { if err := gf.RemoveId(result.FileId); err != nil { log.Error("remove file error, id:" + result.FileId.Hex()) debug.PrintStack() } } return nil } // 删除所有爬虫 func RemoveAllSpider() error { s, c := database.GetCol("spiders") defer s.Close() var spiders []Spider err := c.Find(nil).All(&spiders) if err != nil { log.Error("get all spiders error:" + err.Error()) return err } for _, spider := range spiders { if err := RemoveSpider(spider.Id); err != nil { log.Error("remove spider error:" + err.Error()) } } return nil } // 获取爬虫总数 func GetSpiderCount() (int, error) { s, c := database.GetCol("spiders") defer s.Close() count, err := c.Count() if err != nil { return 0, err } return count, nil } // 获取爬虫定时任务 func GetConfigSpiderData(spider Spider) (entity.ConfigSpiderData, error) { // 构造配置数据 configData := entity.ConfigSpiderData{} // 校验爬虫类别 if spider.Type != constants.Configurable { return configData, errors.New("not a configurable spider") } // Spiderfile 目录 sfPath := filepath.Join(spider.Src, "Spiderfile") // 读取YAML文件 yamlFile, err := ioutil.ReadFile(sfPath) if err != nil { return configData, err } // 反序列化 if err := yaml.Unmarshal(yamlFile, &configData); err != nil { return configData, err } return configData, nil }