加入权限管理

backend/model/project.go

@@ -16,11 +16,12 @@ type Project struct {
 	Tags        []string      `json:"tags" bson:"tags"`
 
 	// 前端展示
-	Spiders []Spider `json:"spiders" bson:"spiders"`
+	Spiders  []Spider `json:"spiders" bson:"spiders"`
+	Username string   `json:"username" bson:"username"`
 
 	UserId   bson.ObjectId `json:"user_id" bson:"user_id"`
-	CreateTs time.Time `json:"create_ts" bson:"create_ts"`
-	UpdateTs time.Time `json:"update_ts" bson:"update_ts"`
+	CreateTs time.Time     `json:"create_ts" bson:"create_ts"`
+	UpdateTs time.Time     `json:"update_ts" bson:"update_ts"`
 }
 
 func (p *Project) Save() error {
@@ -90,15 +91,21 @@ func GetProject(id bson.ObjectId) (Project, error) {
 	return p, nil
 }
 
-func GetProjectList(filter interface{}, skip int, sortKey string) ([]Project, error) {
+func GetProjectList(filter interface{}, sortKey string) ([]Project, error) {
 	s, c := database.GetCol("projects")
 	defer s.Close()
 
 	var projects []Project
-	if err := c.Find(filter).Skip(skip).Limit(constants.Infinite).Sort(sortKey).All(&projects); err != nil {
+	if err := c.Find(filter).Sort(sortKey).All(&projects); err != nil {
 		debug.PrintStack()
 		return projects, err
 	}
+
+	for i, p := range projects {
+		// 获取用户名称
+		user, _ := GetUser(p.UserId)
+		projects[i].Username = user.Username
+	}
 	return projects, nil
 }
 

backend/model/schedule.go

@@ -29,6 +29,7 @@ type Schedule struct {
 
 	// 前端展示
 	SpiderName string `json:"spider_name" bson:"spider_name"`
+	Username   string `json:"user_name" bson:"user_name"`
 	Nodes      []Node `json:"nodes" bson:"nodes"`
 	Message    string `json:"message" bson:"message"`
 
@@ -83,6 +84,10 @@ func GetScheduleList(filter interface{}) ([]Schedule, error) {
 			schedule.SpiderName = spider.Name
 		}
 
+		// 获取用户名称
+		user, _ := GetUser(schedule.UserId)
+		schedule.Username = user.Username
+
 		schs = append(schs, schedule)
 	}
 	return schs, nil
@@ -92,11 +97,16 @@ func GetSchedule(id bson.ObjectId) (Schedule, error) {
 	s, c := database.GetCol("schedules")
 	defer s.Close()
 
-	var result Schedule
-	if err := c.FindId(id).One(&result); err != nil {
-		return result, err
+	var schedule Schedule
+	if err := c.FindId(id).One(&schedule); err != nil {
+		return schedule, err
 	}
-	return result, nil
+
+	// 获取用户名称
+	user, _ := GetUser(schedule.UserId)
+	schedule.Username = user.Username
+
+	return schedule, nil
 }
 
 func UpdateSchedule(id bson.ObjectId, item Schedule) error {
@@ -147,11 +157,11 @@ func RemoveSchedule(id bson.ObjectId) error {
 	return nil
 }
 
-func GetScheduleCount() (int, error) {
+func GetScheduleCount(filter interface{}) (int, error) {
 	s, c := database.GetCol("schedules")
 	defer s.Close()
 
-	count, err := c.Count()
+	count, err := c.Find(filter).Count()
 	if err != nil {
 		return 0, err
 	}

backend/model/spider.go

@@ -33,6 +33,7 @@ type Spider struct {
 	Remark      string        `json:"remark" bson:"remark"`             // 备注
 	Src         string        `json:"src" bson:"src"`                   // 源码位置
 	ProjectId   bson.ObjectId `json:"project_id" bson:"project_id"`     // 项目ID
+	IsPublic    bool          `json:"is_public" bson:"is_public"`       // 是否公开
 
 	// 自定义爬虫
 	Cmd string `json:"cmd" bson:"cmd"` // 执行命令
@@ -63,6 +64,7 @@ type Spider struct {
 	LastStatus  string                  `json:"last_status"`  // 最后执行状态
 	Config      entity.ConfigSpiderData `json:"config"`       // 可配置爬虫配置
 	LatestTasks []Task                  `json:"latest_tasks"` // 最近任务列表
+	Username    string                  `json:"username""`
 
 	// 时间
 	UserId   bson.ObjectId `json:"user_id" bson:"user_id"`
@@ -83,6 +85,7 @@ func (spider *Spider) Save() error {
 	}
 
 	if err := c.UpdateId(spider.Id, spider); err != nil {
+		log.Errorf(err.Error())
 		debug.PrintStack()
 		return err
 	}
@@ -182,10 +185,22 @@ func GetSpiderList(filter interface{}, skip int, limit int, sortStr string) ([]S
 			continue
 		}
 
+		// 获取用户
+		var user User
+		if spider.UserId.Valid() {
+			user, err = GetUser(spider.UserId)
+			if err != nil {
+				log.Errorf(err.Error())
+				debug.PrintStack()
+				continue
+			}
+		}
+
 		// 赋值
 		spiders[i].LastRunTs = task.CreateTs
 		spiders[i].LastStatus = task.Status
 		spiders[i].LatestTasks = latestTasks
+		spiders[i].Username = user.Username
 	}
 
 	count, _ := c.Find(filter).Count()
@@ -221,13 +236,21 @@ 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 {
+	var spider Spider
+	if err := c.Find(bson.M{"name": name}).One(&spider); err != nil && err != mgo.ErrNotFound {
 		log.Errorf("get spider error: %s, spider_name: %s", err.Error(), name)
 		//debug.PrintStack()
-		return result
+		return spider
 	}
-	return result
+
+	// 获取用户
+	var user User
+	if spider.UserId.Valid() {
+		user, _ = GetUser(spider.UserId)
+	}
+	spider.Username = user.Username
+
+	return spider
 }
 
 // 获取爬虫(根据ID)
@@ -253,6 +276,14 @@ func GetSpider(id bson.ObjectId) (Spider, error) {
 		}
 		spider.Config = config
 	}
+
+	// 获取用户名称
+	var user User
+	if spider.UserId.Valid() {
+		user, _ = GetUser(spider.UserId)
+	}
+	spider.Username = user.Username
+
 	return spider, nil
 }
 
@@ -324,11 +355,11 @@ func RemoveAllSpider() error {
 }
 
 // 获取爬虫总数
-func GetSpiderCount() (int, error) {
+func GetSpiderCount(filter interface{}) (int, error) {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
 
-	count, err := c.Count()
+	count, err := c.Find(filter).Count()
 	if err != nil {
 		return 0, err
 	}

backend/model/task.go

@@ -31,6 +31,7 @@ type Task struct {
 	// 前端数据
 	SpiderName string `json:"spider_name"`
 	NodeName   string `json:"node_name"`
+	Username   string `json:"username"`
 
 	UserId   bson.ObjectId `json:"user_id" bson:"user_id"`
 	CreateTs time.Time     `json:"create_ts" bson:"create_ts"`
@@ -128,6 +129,10 @@ func GetTaskList(filter interface{}, skip int, limit int, sortKey string) ([]Tas
 		if node, err := task.GetNode(); err == nil {
 			tasks[i].NodeName = node.Name
 		}
+
+		// 获取用户名称
+		user, _ := GetUser(task.UserId)
+		task.Username = user.Username
 	}
 	return tasks, nil
 }
@@ -156,6 +161,11 @@ func GetTask(id string) (Task, error) {
 		debug.PrintStack()
 		return task, err
 	}
+
+	// 获取用户名称
+	user, _ := GetUser(task.UserId)
+	task.Username = user.Username
+
 	return task, nil
 }