From 178e46dd93c7a88ddb400c23cb07fffdf0885842 Mon Sep 17 00:00:00 2001
From: marvzhang <tikazyq@163.com>
Date: Mon, 25 Nov 2019 13:36:14 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E5=8F=AF=E9=85=8D=E7=BD=AE?=
 =?UTF-8?q?=E7=88=AC=E8=99=AB=EF=BC=9A=E6=B7=BB=E5=8A=A0=E4=B8=8A=E4=BC=A0?=
 =?UTF-8?q?GridFS?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/routes/config_spider.go | 42 +++++++++++++++++++++++++++++++--
 backend/routes/spider.go        |  1 +
 backend/utils/file.go           | 19 +++++++++++++++
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/backend/routes/config_spider.go b/backend/routes/config_spider.go
index bdae068c..6f4a2893 100644
--- a/backend/routes/config_spider.go
+++ b/backend/routes/config_spider.go
@@ -2,13 +2,16 @@ package routes
 
 import (
 	"crawlab/constants"
+	"crawlab/database"
 	"crawlab/entity"
 	"crawlab/model"
 	"crawlab/services"
 	"crawlab/utils"
 	"fmt"
+	"github.com/apex/log"
 	"github.com/gin-gonic/gin"
 	"github.com/globalsign/mgo/bson"
+	uuid "github.com/satori/go.uuid"
 	"github.com/spf13/viper"
 	"gopkg.in/yaml.v2"
 	"io"
@@ -16,6 +19,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"runtime/debug"
 )
 
 // 添加可配置爬虫
@@ -190,11 +194,45 @@ func UploadConfigSpider(c *gin.Context) {
 		return
 	}
 
-	// TODO: 上传到GridFS
+	// 打包为 zip 文件
+	files, err := utils.GetFilesFromDir(spiderDir)
+	if err != nil {
+		HandleError(http.StatusInternalServerError, c, err)
+		return
+	}
+	randomId := uuid.NewV4()
+	tmpFilePath := filepath.Join(viper.GetString("other.tmppath"), spider.Name+"."+randomId.String()+".zip")
+	spiderZipFileName := spider.Name + ".zip"
+	if err := utils.Compress(files, tmpFilePath); err != nil {
+		HandleError(http.StatusInternalServerError, c, err)
+		return
+	}
+
+	// 获取 GridFS 实例
+	s, gf := database.GetGridFs("files")
+	defer s.Close()
+
+	// 判断文件是否已经存在
+	var gfFile model.GridFs
+	if err := gf.Find(bson.M{"filename": spiderZipFileName}).One(&gfFile); err == nil {
+		// 已经存在文件，则删除
+		_ = gf.RemoveId(gfFile.Id)
+	}
+
+	// 上传到GridFs
+	fid, err := services.UploadToGridFs(spiderZipFileName, tmpFilePath)
+	if err != nil {
+		log.Errorf("upload to grid fs error: %s", err.Error())
+		debug.PrintStack()
+		return
+	}
+
+	// 保存爬虫 FileId
+	spider.FileId = fid
+	_ = spider.Save()
 
 	c.JSON(http.StatusOK, Response{
 		Status:  "ok",
 		Message: "success",
-		Data:    configData,
 	})
 }
diff --git a/backend/routes/spider.go b/backend/routes/spider.go
index 4c26fcee..d351f1bb 100644
--- a/backend/routes/spider.go
+++ b/backend/routes/spider.go
@@ -153,6 +153,7 @@ func PutSpider(c *gin.Context) {
 		return
 	}
 
+	// 获取 GridFS 实例
 	s, gf := database.GetGridFs("files")
 	defer s.Close()
 
diff --git a/backend/utils/file.go b/backend/utils/file.go
index 82ed01fd..681d129a 100644
--- a/backend/utils/file.go
+++ b/backend/utils/file.go
@@ -254,6 +254,25 @@ func _Compress(file *os.File, prefix string, zw *zip.Writer) error {
 	return nil
 }
 
+func GetFilesFromDir(dirPath string) ([]*os.File, error) {
+	var res []*os.File
+	if err := filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error {
+		if !IsDir(path) {
+			f, err2 := os.Open(path)
+			if err2 != nil {
+				return err
+			}
+			res = append(res, f)
+		}
+		return nil
+	}); err != nil {
+		log.Error(err.Error())
+		debug.PrintStack()
+		return res, err
+	}
+	return res, nil
+}
+
 // File copies a single file from src to dst
 func CopyFile(src, dst string) error {
 	var err error