From 41e7edf9f9e58b7ed4511a1d63ae73fb78160604 Mon Sep 17 00:00:00 2001 From: Marvin Zhang Date: Mon, 17 Jun 2024 23:31:34 +0800 Subject: [PATCH] feat: optimized code --- core/controllers/base_file_v2.go | 305 +++++++++++++++++++++++++++++++ core/controllers/spider_v2.go | 257 +++++--------------------- core/fs/service_v2.go | 11 ++ core/interfaces/fs_service_v2.go | 1 + core/models/models/spider_v2.go | 10 +- core/spider/admin/service_v2.go | 23 --- 6 files changed, 364 insertions(+), 243 deletions(-) create mode 100644 core/controllers/base_file_v2.go diff --git a/core/controllers/base_file_v2.go b/core/controllers/base_file_v2.go new file mode 100644 index 00000000..6c398abd --- /dev/null +++ b/core/controllers/base_file_v2.go @@ -0,0 +1,305 @@ +package controllers + +import ( + "errors" + "fmt" + log2 "github.com/apex/log" + "github.com/crawlab-team/crawlab/core/fs" + "github.com/crawlab-team/crawlab/core/interfaces" + "github.com/gin-gonic/gin" + "github.com/spf13/viper" + "go.mongodb.org/mongo-driver/bson/primitive" + "io" + "os" + "path/filepath" + "sync" +) + +func GetBaseFileListDir(rootPath string, c *gin.Context) { + path := c.Query("path") + + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + + files, err := fsSvc.List(path) + if err != nil { + if !errors.Is(err, os.ErrNotExist) { + HandleErrorInternalServerError(c, err) + return + } + } + + HandleSuccessWithData(c, files) +} + +func GetBaseFileFile(rootPath string, c *gin.Context) { + path := c.Query("path") + + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + + data, err := fsSvc.GetFile(path) + if err != nil { + HandleErrorInternalServerError(c, err) + return + } + + HandleSuccessWithData(c, string(data)) +} + +func GetBaseFileFileInfo(rootPath string, c *gin.Context) { + path := c.Query("path") + + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + + info, err := fsSvc.GetFileInfo(path) + if err != nil { + HandleErrorInternalServerError(c, err) + return + } + + HandleSuccessWithData(c, info) +} + +func PostBaseFileSaveFile(rootPath string, c *gin.Context) { + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorInternalServerError(c, err) + return + } + + if c.GetHeader("Content-Type") == "application/json" { + var payload struct { + Path string `json:"path"` + Data string `json:"data"` + } + if err := c.ShouldBindJSON(&payload); err != nil { + HandleErrorBadRequest(c, err) + return + } + if err := fsSvc.Save(payload.Path, []byte(payload.Data)); err != nil { + HandleErrorInternalServerError(c, err) + return + } + } else { + path, ok := c.GetPostForm("path") + if !ok { + HandleErrorBadRequest(c, errors.New("missing required field 'path'")) + return + } + file, err := c.FormFile("file") + if err != nil { + HandleErrorBadRequest(c, err) + return + } + f, err := file.Open() + if err != nil { + HandleErrorBadRequest(c, err) + return + } + fileData, err := io.ReadAll(f) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + if err := fsSvc.Save(path, fileData); err != nil { + HandleErrorInternalServerError(c, err) + return + } + } + + HandleSuccess(c) +} + +func PostBaseFileSaveFiles(rootPath string, c *gin.Context) { + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorInternalServerError(c, err) + return + } + + form, err := c.MultipartForm() + if err != nil { + HandleErrorBadRequest(c, err) + return + } + wg := sync.WaitGroup{} + wg.Add(len(form.File)) + for path := range form.File { + go func(path string) { + file, err := c.FormFile(path) + if err != nil { + log2.Warnf("invalid file header: %s", path) + log2.Error(err.Error()) + wg.Done() + return + } + f, err := file.Open() + if err != nil { + log2.Warnf("unable to open file: %s", path) + log2.Error(err.Error()) + wg.Done() + return + } + fileData, err := io.ReadAll(f) + if err != nil { + log2.Warnf("unable to read file: %s", path) + log2.Error(err.Error()) + wg.Done() + return + } + if err := fsSvc.Save(path, fileData); err != nil { + log2.Warnf("unable to save file: %s", path) + log2.Error(err.Error()) + wg.Done() + return + } + wg.Done() + }(path) + } + wg.Wait() + + HandleSuccess(c) +} + +func PostBaseFileSaveDir(rootPath string, c *gin.Context) { + var payload struct { + Path string `json:"path"` + NewPath string `json:"new_path"` + Data string `json:"data"` + } + if err := c.ShouldBindJSON(&payload); err != nil { + HandleErrorBadRequest(c, err) + return + } + + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + + if err := fsSvc.CreateDir(payload.Path); err != nil { + HandleErrorInternalServerError(c, err) + return + } + + HandleSuccess(c) +} + +func PostBaseFileRenameFile(rootPath string, c *gin.Context) { + var payload struct { + Path string `json:"path"` + NewPath string `json:"new_path"` + } + if err := c.ShouldBindJSON(&payload); err != nil { + HandleErrorBadRequest(c, err) + return + } + + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + + if err := fsSvc.Rename(payload.Path, payload.NewPath); err != nil { + HandleErrorInternalServerError(c, err) + return + } +} + +func DeleteBaseFileFile(rootPath string, c *gin.Context) { + var payload struct { + Path string `json:"path"` + } + if err := c.ShouldBindJSON(&payload); err != nil { + HandleErrorBadRequest(c, err) + return + } + if payload.Path == "~" { + payload.Path = "." + } + + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + + if err := fsSvc.Delete(payload.Path); err != nil { + HandleErrorInternalServerError(c, err) + return + } + _, err = fsSvc.GetFileInfo(".") + if err != nil { + _ = fsSvc.CreateDir("/") + } + + HandleSuccess(c) +} + +func PostBaseFileCopyFile(rootPath string, c *gin.Context) { + var payload struct { + Path string `json:"path"` + NewPath string `json:"new_path"` + } + if err := c.ShouldBindJSON(&payload); err != nil { + HandleErrorBadRequest(c, err) + return + } + + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + + if err := fsSvc.Copy(payload.Path, payload.NewPath); err != nil { + HandleErrorInternalServerError(c, err) + return + } + + HandleSuccess(c) +} + +func PostBaseFileExport(rootPath string, c *gin.Context) { + fsSvc, err := getBaseFileFsSvc(rootPath, c) + if err != nil { + HandleErrorBadRequest(c, err) + return + } + + // zip file path + zipFilePath, err := fsSvc.Export() + if err != nil { + HandleErrorInternalServerError(c, err) + return + } + + // download + c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=%s", zipFilePath)) + c.File(zipFilePath) +} + +func getBaseFileFsSvc(rootPath string, c *gin.Context) (svc interfaces.FsServiceV2, err error) { + id, err := primitive.ObjectIDFromHex(c.Param("id")) + if err != nil { + return nil, err + } + + workspacePath := viper.GetString("workspace") + fsSvc := fs.NewFsServiceV2(filepath.Join(workspacePath, id.Hex(), rootPath)) + + return fsSvc, nil +} diff --git a/core/controllers/spider_v2.go b/core/controllers/spider_v2.go index 11f1c491..9b28b94d 100644 --- a/core/controllers/spider_v2.go +++ b/core/controllers/spider_v2.go @@ -2,7 +2,6 @@ package controllers import ( "errors" - "fmt" log2 "github.com/apex/log" "github.com/crawlab-team/crawlab/core/constants" "github.com/crawlab-team/crawlab/core/fs" @@ -10,13 +9,13 @@ import ( "github.com/crawlab-team/crawlab/core/models/models" "github.com/crawlab-team/crawlab/core/models/service" "github.com/crawlab-team/crawlab/core/spider/admin" + "github.com/crawlab-team/crawlab/core/utils" "github.com/crawlab-team/crawlab/db/mongo" "github.com/gin-gonic/gin" "github.com/spf13/viper" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/primitive" mongo2 "go.mongodb.org/mongo-driver/mongo" - "io" "math" "os" "path/filepath" @@ -432,286 +431,93 @@ func DeleteSpiderList(c *gin.Context) { } func GetSpiderListDir(c *gin.Context) { - path := c.Query("path") - - fsSvc, err := getSpiderFsSvc(c) + s, err := allowSpiderGit(c) if err != nil { - HandleErrorBadRequest(c, err) + HandleErrorInternalServerError(c, err) return } - - files, err := fsSvc.List(path) - if err != nil { - if !errors.Is(err, os.ErrNotExist) { - HandleErrorInternalServerError(c, err) - return - } - } - - HandleSuccessWithData(c, files) + GetBaseFileListDir(s.GitRootPath, c) } func GetSpiderFile(c *gin.Context) { - path := c.Query("path") - - fsSvc, err := getSpiderFsSvc(c) - if err != nil { - HandleErrorBadRequest(c, err) - return - } - - data, err := fsSvc.GetFile(path) + s, err := allowSpiderGit(c) if err != nil { HandleErrorInternalServerError(c, err) return } - - HandleSuccessWithData(c, string(data)) + GetBaseFileFile(s.GitRootPath, c) } func GetSpiderFileInfo(c *gin.Context) { - path := c.Query("path") - - fsSvc, err := getSpiderFsSvc(c) - if err != nil { - HandleErrorBadRequest(c, err) - return - } - - info, err := fsSvc.GetFileInfo(path) + s, err := allowSpiderGit(c) if err != nil { HandleErrorInternalServerError(c, err) return } - - HandleSuccessWithData(c, info) + GetBaseFileFileInfo(s.GitRootPath, c) } func PostSpiderSaveFile(c *gin.Context) { - fsSvc, err := getSpiderFsSvc(c) + s, err := allowSpiderGit(c) if err != nil { HandleErrorInternalServerError(c, err) return } - - if c.GetHeader("Content-Type") == "application/json" { - var payload struct { - Path string `json:"path"` - Data string `json:"data"` - } - if err := c.ShouldBindJSON(&payload); err != nil { - HandleErrorBadRequest(c, err) - return - } - if err := fsSvc.Save(payload.Path, []byte(payload.Data)); err != nil { - HandleErrorInternalServerError(c, err) - return - } - } else { - path, ok := c.GetPostForm("path") - if !ok { - HandleErrorBadRequest(c, errors.New("missing required field 'path'")) - return - } - file, err := c.FormFile("file") - if err != nil { - HandleErrorBadRequest(c, err) - return - } - f, err := file.Open() - if err != nil { - HandleErrorBadRequest(c, err) - return - } - fileData, err := io.ReadAll(f) - if err != nil { - HandleErrorBadRequest(c, err) - return - } - if err := fsSvc.Save(path, fileData); err != nil { - HandleErrorInternalServerError(c, err) - return - } - } - - HandleSuccess(c) + PostBaseFileSaveFile(s.GitRootPath, c) } func PostSpiderSaveFiles(c *gin.Context) { - fsSvc, err := getSpiderFsSvc(c) + s, err := allowSpiderGit(c) if err != nil { HandleErrorInternalServerError(c, err) return } - - form, err := c.MultipartForm() - if err != nil { - HandleErrorBadRequest(c, err) - return - } - wg := sync.WaitGroup{} - wg.Add(len(form.File)) - for path := range form.File { - go func(path string) { - file, err := c.FormFile(path) - if err != nil { - log2.Warnf("invalid file header: %s", path) - log2.Error(err.Error()) - wg.Done() - return - } - f, err := file.Open() - if err != nil { - log2.Warnf("unable to open file: %s", path) - log2.Error(err.Error()) - wg.Done() - return - } - fileData, err := io.ReadAll(f) - if err != nil { - log2.Warnf("unable to read file: %s", path) - log2.Error(err.Error()) - wg.Done() - return - } - if err := fsSvc.Save(path, fileData); err != nil { - log2.Warnf("unable to save file: %s", path) - log2.Error(err.Error()) - wg.Done() - return - } - wg.Done() - }(path) - } - wg.Wait() - - HandleSuccess(c) + PostBaseFileSaveFiles(s.GitRootPath, c) } func PostSpiderSaveDir(c *gin.Context) { - var payload struct { - Path string `json:"path"` - NewPath string `json:"new_path"` - Data string `json:"data"` - } - if err := c.ShouldBindJSON(&payload); err != nil { - HandleErrorBadRequest(c, err) - return - } - - fsSvc, err := getSpiderFsSvc(c) + s, err := allowSpiderGit(c) if err != nil { - HandleErrorBadRequest(c, err) - return - } - - if err := fsSvc.CreateDir(payload.Path); err != nil { HandleErrorInternalServerError(c, err) return } - - HandleSuccess(c) + PostBaseFileSaveDir(s.GitRootPath, c) } func PostSpiderRenameFile(c *gin.Context) { - var payload struct { - Path string `json:"path"` - NewPath string `json:"new_path"` - } - if err := c.ShouldBindJSON(&payload); err != nil { - HandleErrorBadRequest(c, err) - return - } - - fsSvc, err := getSpiderFsSvc(c) + s, err := allowSpiderGit(c) if err != nil { - HandleErrorBadRequest(c, err) - return - } - - if err := fsSvc.Rename(payload.Path, payload.NewPath); err != nil { HandleErrorInternalServerError(c, err) return } + PostBaseFileRenameFile(s.GitRootPath, c) } func DeleteSpiderFile(c *gin.Context) { - var payload struct { - Path string `json:"path"` - } - if err := c.ShouldBindJSON(&payload); err != nil { - HandleErrorBadRequest(c, err) - return - } - if payload.Path == "~" { - payload.Path = "." - } - - fsSvc, err := getSpiderFsSvc(c) + s, err := allowSpiderGit(c) if err != nil { - HandleErrorBadRequest(c, err) - return - } - - if err := fsSvc.Delete(payload.Path); err != nil { HandleErrorInternalServerError(c, err) return } - _, err = fsSvc.GetFileInfo(".") - if err != nil { - _ = fsSvc.CreateDir("/") - } - - HandleSuccess(c) + DeleteBaseFileFile(s.GitRootPath, c) } func PostSpiderCopyFile(c *gin.Context) { - var payload struct { - Path string `json:"path"` - NewPath string `json:"new_path"` - } - if err := c.ShouldBindJSON(&payload); err != nil { - HandleErrorBadRequest(c, err) - return - } - - fsSvc, err := getSpiderFsSvc(c) + s, err := allowSpiderGit(c) if err != nil { - HandleErrorBadRequest(c, err) - return - } - - if err := fsSvc.Copy(payload.Path, payload.NewPath); err != nil { HandleErrorInternalServerError(c, err) return } - - HandleSuccess(c) + PostBaseFileCopyFile(s.GitRootPath, c) } func PostSpiderExport(c *gin.Context) { - id, err := primitive.ObjectIDFromHex(c.Param("id")) - if err != nil { - HandleErrorBadRequest(c, err) - return - } - - adminSvc, err := admin.GetSpiderAdminServiceV2() + s, err := allowSpiderGit(c) if err != nil { HandleErrorInternalServerError(c, err) return } - - // zip file path - zipFilePath, err := adminSvc.Export(id) - if err != nil { - HandleErrorInternalServerError(c, err) - return - } - - // download - c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=%s", zipFilePath)) - c.File(zipFilePath) + PostBaseFileExport(s.GitRootPath, c) } func PostSpiderRun(c *gin.Context) { @@ -878,3 +684,22 @@ func upsertSpiderDataCollection(s *models.SpiderV2) (err error) { } return nil } + +func allowSpiderGit(c *gin.Context) (s models.SpiderV2, err error) { + if utils.IsPro() { + return s, nil + } + id, err := primitive.ObjectIDFromHex(c.Param("id")) + if err != nil { + return s, err + } + _s, err := service.NewModelServiceV2[models.SpiderV2]().GetById(id) + if err != nil { + return s, err + } + if s.GitId.IsZero() { + return s, errors.New("git is not allowed in this edition") + } + s = *_s + return s, nil +} diff --git a/core/fs/service_v2.go b/core/fs/service_v2.go index 6ca1f057..918a3593 100644 --- a/core/fs/service_v2.go +++ b/core/fs/service_v2.go @@ -4,6 +4,8 @@ import ( "github.com/crawlab-team/crawlab/core/entity" "github.com/crawlab-team/crawlab/core/interfaces" "github.com/crawlab-team/crawlab/core/utils" + "github.com/crawlab-team/crawlab/trace" + "github.com/google/uuid" "io" "os" "path/filepath" @@ -159,6 +161,15 @@ func (svc *ServiceV2) Copy(path, newPath string) (err error) { } } +func (svc *ServiceV2) Export() (resultPath string, err error) { + zipFilePath := filepath.Join(os.TempDir(), uuid.New().String()+".zip") + if err := utils.ZipDirectory(svc.rootPath, zipFilePath); err != nil { + return "", trace.TraceError(err) + } + + return zipFilePath, nil +} + func NewFsServiceV2(path string) (svc interfaces.FsServiceV2) { return &ServiceV2{ rootPath: path, diff --git a/core/interfaces/fs_service_v2.go b/core/interfaces/fs_service_v2.go index 4bc3d7df..533a9d6b 100644 --- a/core/interfaces/fs_service_v2.go +++ b/core/interfaces/fs_service_v2.go @@ -9,4 +9,5 @@ type FsServiceV2 interface { Rename(path, newPath string) (err error) Delete(path string) (err error) Copy(path, newPath string) (err error) + Export() (resultPath string, err error) } diff --git a/core/models/models/spider_v2.go b/core/models/models/spider_v2.go index 4f98f6e7..939bad51 100644 --- a/core/models/models/spider_v2.go +++ b/core/models/models/spider_v2.go @@ -17,14 +17,16 @@ type SpiderV2 struct { ProjectId primitive.ObjectID `json:"project_id" bson:"project_id"` // Project.Id Mode string `json:"mode" bson:"mode"` // default Task.Mode NodeIds []primitive.ObjectID `json:"node_ids" bson:"node_ids"` // default Task.NodeIds - Stat *SpiderStatV2 `json:"stat,omitempty" bson:"-"` + GitId primitive.ObjectID `json:"git_id" bson:"git_id"` // related Git.Id + GitRootPath string `json:"git_root_path" bson:"git_root_path"` + Git *GitV2 `json:"git,omitempty" bson:"-"` + + // stats + Stat *SpiderStatV2 `json:"stat,omitempty" bson:"-"` // execution Cmd string `json:"cmd" bson:"cmd"` // execute command Param string `json:"param" bson:"param"` // default task param Priority int `json:"priority" bson:"priority"` AutoInstall bool `json:"auto_install" bson:"auto_install"` - - // settings - IncrementalSync bool `json:"incremental_sync" bson:"incremental_sync"` // whether to incrementally sync files } diff --git a/core/spider/admin/service_v2.go b/core/spider/admin/service_v2.go index c62a5fee..628a7f48 100644 --- a/core/spider/admin/service_v2.go +++ b/core/spider/admin/service_v2.go @@ -14,13 +14,10 @@ import ( "github.com/crawlab-team/crawlab/core/utils" "github.com/crawlab-team/crawlab/trace" "github.com/crawlab-team/crawlab/vcs" - "github.com/google/uuid" "github.com/robfig/cron/v3" "github.com/spf13/viper" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/primitive" - "os" - "path" "path/filepath" "sync" "time" @@ -60,26 +57,6 @@ func (svc *ServiceV2) SyncGit() (err error) { return nil } -func (svc *ServiceV2) SyncGitOne(g *models.GitV2) (err error) { - svc.syncGitOne(g) - return nil -} - -func (svc *ServiceV2) Export(id primitive.ObjectID) (filePath string, err error) { - // spider fs - workspacePath := viper.GetString("workspace") - spiderFolderPath := filepath.Join(workspacePath, id.Hex()) - - // zip files in workspace - dirPath := spiderFolderPath - zipFilePath := path.Join(os.TempDir(), uuid.New().String()+".zip") - if err := utils.ZipDirectory(dirPath, zipFilePath); err != nil { - return "", trace.TraceError(err) - } - - return zipFilePath, nil -} - func (svc *ServiceV2) scheduleTasks(s *models.SpiderV2, opts *interfaces.SpiderRunOptions) (taskIds []primitive.ObjectID, err error) { // main task mainTask := &models.TaskV2{