From ffd5725f3b6bbf508b61f73d00633980fb021a31 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=99=88=E6=99=AF=E9=98=B3?= <1656488874@qq.com>
Date: Mon, 30 Sep 2019 14:43:48 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E7=88=AC=E8=99=AB=E7=B1=BB?=
 =?UTF-8?q?=E5=9E=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/entity/spider.go |  6 ++++++
 backend/main.go          |  1 +
 backend/model/spider.go  | 34 +++++++++++++++++++++++++++++++++-
 backend/routes/spider.go | 14 ++++++++++++++
 4 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 backend/entity/spider.go

diff --git a/backend/entity/spider.go b/backend/entity/spider.go
new file mode 100644
index 00000000..8129ff40
--- /dev/null
+++ b/backend/entity/spider.go
@@ -0,0 +1,6 @@
+package entity
+
+type SpiderType struct {
+	Type  string `json:"type"`
+	Count int    `json:"count"`
+}
diff --git a/backend/main.go b/backend/main.go
index 47196fe5..5d95dbaf 100644
--- a/backend/main.go
+++ b/backend/main.go
@@ -132,6 +132,7 @@ func main() {
 			authGroup.POST("/spiders/:id/file", routes.PostSpiderFile)   // 爬虫目录写入
 			authGroup.GET("/spiders/:id/dir", routes.GetSpiderDir)       // 爬虫目录
 			authGroup.GET("/spiders/:id/stats", routes.GetSpiderStats)   // 爬虫统计数据
+			authGroup.GET("/spider/types", routes.GetSpiderTypes)        // 爬虫类型
 			// 任务
 			authGroup.GET("/tasks", routes.GetTaskList)                                 // 任务列表
 			authGroup.GET("/tasks/:id", routes.GetTask)                                 // 任务详情
diff --git a/backend/model/spider.go b/backend/model/spider.go
index 93eafbc5..6e7102d6 100644
--- a/backend/model/spider.go
+++ b/backend/model/spider.go
@@ -2,6 +2,7 @@ package model
 
 import (
 	"crawlab/database"
+	"crawlab/entity"
 	"github.com/apex/log"
 	"github.com/globalsign/mgo"
 	"github.com/globalsign/mgo/bson"
@@ -47,6 +48,7 @@ type Spider struct {
 	UpdateTs time.Time `json:"update_ts" bson:"update_ts"`
 }
 
+// 更新爬虫
 func (spider *Spider) Save() error {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -60,6 +62,7 @@ func (spider *Spider) Save() error {
 	return nil
 }
 
+// 新增爬虫
 func (spider *Spider) Add() error {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -74,6 +77,7 @@ func (spider *Spider) Add() error {
 	return nil
 }
 
+// 获取爬虫的任务
 func (spider *Spider) GetTasks() ([]Task, error) {
 	tasks, err := GetTaskList(bson.M{"spider_id": spider.Id}, 0, 10, "-create_ts")
 	if err != nil {
@@ -82,6 +86,7 @@ func (spider *Spider) GetTasks() ([]Task, error) {
 	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 {
@@ -93,6 +98,7 @@ func (spider *Spider) GetLastTask() (Task, error) {
 	return tasks[0], nil
 }
 
+// 爬虫列表
 func GetSpiderList(filter interface{}, skip int, limit int) ([]Spider, int, error) {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -124,6 +130,7 @@ func GetSpiderList(filter interface{}, skip int, limit int) ([]Spider, int, erro
 	return spiders, count, nil
 }
 
+// 获取爬虫
 func GetSpiderByFileId(fileId bson.ObjectId) *Spider {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -137,6 +144,7 @@ func GetSpiderByFileId(fileId bson.ObjectId) *Spider {
 	return result
 }
 
+// 获取爬虫
 func GetSpiderByName(name string) *Spider {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -150,6 +158,7 @@ func GetSpiderByName(name string) *Spider {
 	return result
 }
 
+// 获取爬虫
 func GetSpider(id bson.ObjectId) (Spider, error) {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -165,6 +174,7 @@ func GetSpider(id bson.ObjectId) (Spider, error) {
 	return result, nil
 }
 
+// 更新爬虫
 func UpdateSpider(id bson.ObjectId, item Spider) error {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -181,6 +191,7 @@ func UpdateSpider(id bson.ObjectId, item Spider) error {
 	return nil
 }
 
+// 删除爬虫
 func RemoveSpider(id bson.ObjectId) error {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -209,6 +220,7 @@ func RemoveSpider(id bson.ObjectId) error {
 	return nil
 }
 
+// 删除所有爬虫
 func RemoveAllSpider() error {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -227,6 +239,7 @@ func RemoveAllSpider() error {
 	return nil
 }
 
+// 爬虫总数
 func GetSpiderCount() (int, error) {
 	s, c := database.GetCol("spiders")
 	defer s.Close()
@@ -235,6 +248,25 @@ func GetSpiderCount() (int, error) {
 	if err != nil {
 		return 0, err
 	}
-
 	return count, nil
 }
+
+// 爬虫类型
+func GetSpiderTypes() ([]*entity.SpiderType, error) {
+	s, c := database.GetCol("spiders")
+	defer s.Close()
+
+	group := bson.M{
+		"$group": bson.M{
+			"type": bson.M{"$sum": 1},
+		},
+	}
+	var types []*entity.SpiderType
+	if err := c.Pipe([]bson.M{group}).All(&types); err != nil {
+		log.Errorf("get spider types error: %s", err.Error())
+		debug.PrintStack()
+		return nil, err
+	}
+
+	return types, nil
+}
diff --git a/backend/routes/spider.go b/backend/routes/spider.go
index efd76c0b..8f8fdd13 100644
--- a/backend/routes/spider.go
+++ b/backend/routes/spider.go
@@ -337,6 +337,20 @@ func PostSpiderFile(c *gin.Context) {
 	})
 }
 
+// 爬虫类型
+func GetSpiderTypes(c *gin.Context) {
+	types, err := model.GetSpiderTypes()
+	if err != nil {
+		HandleError(http.StatusInternalServerError, c, err)
+		return
+	}
+	c.JSON(http.StatusOK, Response{
+		Status:  "ok",
+		Message: "success",
+		Data:    types,
+	})
+}
+
 func GetSpiderStats(c *gin.Context) {
 	type Overview struct {
 		TaskCount            int     `json:"task_count" bson:"task_count"`