diff --git a/controllers/export.go b/controllers/export.go index d25ff29..1ba35f5 100644 --- a/controllers/export.go +++ b/controllers/export.go @@ -35,7 +35,8 @@ func getExportActions() []Action { } type exportContext struct { - csvSvc interfaces.ExportService + csvSvc interfaces.ExportService + jsonSvc interfaces.ExportService } func (ctx *exportContext) postExport(c *gin.Context) { @@ -48,6 +49,8 @@ func (ctx *exportContext) postExport(c *gin.Context) { switch exportType { case constants.ExportTypeCsv: exportId, err = ctx.csvSvc.Export(exportType, exportTarget, exportFilter) + case constants.ExportTypeJson: + exportId, err = ctx.jsonSvc.Export(exportType, exportTarget, exportFilter) default: HandleErrorBadRequest(c, errors.New(fmt.Sprintf("invalid export type: %s", exportType))) return @@ -69,6 +72,8 @@ func (ctx *exportContext) getExport(c *gin.Context) { switch exportType { case constants.ExportTypeCsv: exp, err = ctx.csvSvc.GetExport(exportId) + case constants.ExportTypeJson: + exp, err = ctx.jsonSvc.GetExport(exportId) default: HandleErrorBadRequest(c, errors.New(fmt.Sprintf("invalid export type: %s", exportType))) } @@ -89,6 +94,8 @@ func (ctx *exportContext) getExportDownload(c *gin.Context) { switch exportType { case constants.ExportTypeCsv: exp, err = ctx.csvSvc.GetExport(exportId) + case constants.ExportTypeJson: + exp, err = ctx.jsonSvc.GetExport(exportId) default: HandleErrorBadRequest(c, errors.New(fmt.Sprintf("invalid export type: %s", exportType))) } @@ -97,7 +104,18 @@ func (ctx *exportContext) getExportDownload(c *gin.Context) { return } - c.Header("Content-Type", "text/csv") + switch exportType { + case constants.ExportTypeCsv: + c.Header("Content-Type", "text/csv") + case constants.ExportTypeJson: + c.Header("Content-Type", "text/plain") + default: + HandleErrorBadRequest(c, errors.New(fmt.Sprintf("invalid export type: %s", exportType))) + } + if err != nil { + HandleErrorInternalServerError(c, err) + return + } c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=%s", exp.GetDownloadPath())) c.Header("Content-Length", strconv.Itoa(len(exp.GetDownloadPath()))) c.File(exp.GetDownloadPath()) @@ -105,6 +123,7 @@ func (ctx *exportContext) getExportDownload(c *gin.Context) { func newExportContext() *exportContext { return &exportContext{ - csvSvc: export.GetCsvService(), + csvSvc: export.GetCsvService(), + jsonSvc: export.GetJsonService(), } } diff --git a/export/json_service.go b/export/json_service.go new file mode 100644 index 0000000..110bb18 --- /dev/null +++ b/export/json_service.go @@ -0,0 +1,223 @@ +package export + +import ( + "context" + "encoding/json" + "errors" + "github.com/ReneKroon/ttlcache" + "github.com/apex/log" + "github.com/crawlab-team/crawlab-core/constants" + "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-db/mongo" + "github.com/crawlab-team/go-trace" + "github.com/hashicorp/go-uuid" + mongo2 "go.mongodb.org/mongo-driver/mongo" + "os" + "path" + "time" +) + +type JsonService struct { + cache *ttlcache.Cache +} + +func (svc *JsonService) GenerateId() (exportId string, err error) { + exportId, err = uuid.GenerateUUID() + if err != nil { + return "", trace.TraceError(err) + } + return exportId, nil +} + +func (svc *JsonService) Export(exportType, target string, filter interfaces.Filter) (exportId string, err error) { + // generate export id + exportId, err = svc.GenerateId() + if err != nil { + return "", err + } + + // export + export := &entity.Export{ + Id: exportId, + Type: exportType, + Target: target, + Filter: filter, + Status: constants.TaskStatusRunning, + StartTs: time.Now(), + FileName: svc.getFileName(exportId), + DownloadPath: svc.getDownloadPath(exportId), + Limit: 100, + } + + // save to cache + svc.cache.Set(exportId, export) + + // execute export + go svc.export(export) + + return exportId, nil +} + +func (svc *JsonService) GetExport(exportId string) (export interfaces.Export, err error) { + // get export from cache + res, ok := svc.cache.Get(exportId) + if !ok { + return nil, trace.TraceError(errors.New("export not found")) + } + export = res.(interfaces.Export) + return export, nil +} + +func (svc *JsonService) export(export *entity.Export) { + // check empty + if export.Target == "" { + err := errors.New("empty target") + export.Status = constants.TaskStatusError + export.EndTs = time.Now() + log.Errorf("export error (id: %s): %v", export.Id, err) + trace.PrintError(err) + svc.cache.Set(export.Id, export) + return + } + + // mongo collection + col := mongo.GetMongoCol(export.Target) + + // mongo query + query, err := utils.FilterToQuery(export.Filter) + if err != nil { + export.Status = constants.TaskStatusError + export.EndTs = time.Now() + log.Errorf("export error (id: %s): %v", export.Id, err) + trace.PrintError(err) + svc.cache.Set(export.Id, export) + return + } + + // mongo cursor + cur := col.Find(query, nil).GetCursor() + + // data + var jsonData []interface{} + + // iterate cursor + i := 0 + for { + // increment counter + i++ + + // check error + err := cur.Err() + if err != nil { + if err != mongo2.ErrNoDocuments { + // error + export.Status = constants.TaskStatusError + export.EndTs = time.Now() + log.Errorf("export error (id: %s): %v", export.Id, err) + trace.PrintError(err) + } else { + // no more data + export.Status = constants.TaskStatusFinished + export.EndTs = time.Now() + log.Infof("export finished (id: %s)", export.Id) + } + svc.cache.Set(export.Id, export) + return + } + + // has data + if !cur.Next(context.Background()) { + // no more data + export.Status = constants.TaskStatusFinished + export.EndTs = time.Now() + log.Infof("export finished (id: %s)", export.Id) + svc.cache.Set(export.Id, export) + break + } + + // convert raw data to entity + var data map[string]interface{} + err = cur.Decode(&data) + if err != nil { + // error + export.Status = constants.TaskStatusError + export.EndTs = time.Now() + log.Errorf("export error (id: %s): %v", export.Id, err) + trace.PrintError(err) + svc.cache.Set(export.Id, export) + return + } + + jsonData = append(jsonData, data) + } + + jsonBytes, err := json.Marshal(jsonData) + if err != nil { + // error + export.Status = constants.TaskStatusError + export.EndTs = time.Now() + log.Errorf("export error (id: %s): %v", export.Id, err) + trace.PrintError(err) + svc.cache.Set(export.Id, export) + return + } + jsonString := string(jsonBytes) + f := utils.OpenFile(export.DownloadPath) + _, err = f.WriteString(jsonString) + if err != nil { + // error + export.Status = constants.TaskStatusError + export.EndTs = time.Now() + log.Errorf("export error (id: %s): %v", export.Id, err) + trace.PrintError(err) + svc.cache.Set(export.Id, export) + return + } +} + +func (svc *JsonService) getExportDir() (dir string, err error) { + tempDir := os.TempDir() + exportDir := path.Join(tempDir, "export", "json") + if !utils.Exists(exportDir) { + err := os.MkdirAll(exportDir, 0755) + if err != nil { + return "", err + } + } + return exportDir, nil +} + +func (svc *JsonService) getFileName(exportId string) (fileName string) { + return exportId + "_" + time.Now().Format("20060102150405") + ".json" +} + +// getDownloadPath returns the download path for the export +// format: /export//_.csv +func (svc *JsonService) getDownloadPath(exportId string) (downloadPath string) { + exportDir, err := svc.getExportDir() + if err != nil { + return "" + } + downloadPath = path.Join(exportDir, svc.getFileName(exportId)) + return downloadPath +} + +func NewJsonService() (svc2 interfaces.ExportService) { + cache := ttlcache.NewCache() + cache.SetTTL(time.Minute * 5) + svc := &JsonService{ + cache: cache, + } + return svc +} + +var _jsonService interfaces.ExportService + +func GetJsonService() (svc interfaces.ExportService) { + if _jsonService == nil { + _jsonService = NewJsonService() + } + return _jsonService +}