283 lines
7.0 KiB
Go
283 lines
7.0 KiB
Go
package export_task
|
|
|
|
import (
|
|
"91porn-server/common/file"
|
|
"91porn-server/common/log"
|
|
commonWorker "91porn-server/common/worker"
|
|
"91porn-server/models/v/export_task_mod"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"reflect"
|
|
"time"
|
|
|
|
tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
|
|
"github.com/pkg/errors"
|
|
"github.com/xuri/excelize/v2"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
)
|
|
|
|
const (
|
|
proName = "91PORN"
|
|
tgToken = "5852134617:AAEO0_GPZiUJS_LeR5dM7byEfOpnRg4dGqU"
|
|
maxDataNum = 20000 // 每个excel存储的数据量
|
|
maxWorkerNum = 3
|
|
maxSize = 1024 * 1024 * 20
|
|
)
|
|
|
|
var Bot *tgbotapi.BotAPI
|
|
var worker *commonWorker.Worker
|
|
|
|
func init() {
|
|
bot, err := tgbotapi.NewBotAPI(tgToken)
|
|
if err != nil || bot == nil {
|
|
log.Fatal("初始化tgbot失败", log.E(err))
|
|
}
|
|
Bot = bot
|
|
worker = commonWorker.NewWorker(maxWorkerNum)
|
|
}
|
|
|
|
var execExportTaskRuining bool
|
|
|
|
func ExecExportTask() {
|
|
//防止重入
|
|
if execExportTaskRuining {
|
|
return
|
|
}
|
|
execExportTaskRuining = true
|
|
defer func() {
|
|
execExportTaskRuining = false
|
|
}()
|
|
// 获取任务列表
|
|
list, err := export_task_mod.GetTaskList()
|
|
if err != nil {
|
|
log.Error("ExecExportTask export_task_mod.GetTaskList fail", log.E(err))
|
|
return
|
|
}
|
|
for _, task := range list {
|
|
switch task.Type {
|
|
case export_task_mod.ExportUserTask:
|
|
CallExportFunc(task, ExecExportUserTask)
|
|
case export_task_mod.ExportVidTask:
|
|
CallExportFunc(task, ExecExportVideoTask)
|
|
case export_task_mod.ExportAdvanceOrderTask:
|
|
CallExportFunc(task, ExecExportAdvanceOrderTask)
|
|
case export_task_mod.ExportProductHistoryTask:
|
|
CallExportFunc(task, ExecExportProductHistoryTask)
|
|
}
|
|
}
|
|
}
|
|
|
|
func GetProName() string {
|
|
//if skdg.Conf.Base.Env == "test" {
|
|
// return "测试环境_" + proName
|
|
//}
|
|
return proName
|
|
}
|
|
|
|
func CallExportFunc(task *export_task_mod.ExportTask, f func(t *export_task_mod.ExportTask) (int64, error)) {
|
|
if task == nil {
|
|
return
|
|
}
|
|
worker.Exec(func() {
|
|
// 已经执行过了
|
|
if task.Status != 0 {
|
|
return
|
|
}
|
|
//key := fmt.Sprintf("ExecExportTask:%v", task.ID.Hex())
|
|
//ok, err := skdg.Redis.SetNX(key, 1, time.Minute*30)
|
|
//if err != nil {
|
|
// log.Error("ExecExportTask fail", log.E(err))
|
|
// return
|
|
//}
|
|
//if !ok {
|
|
// log.Info("ExportTask already executed")
|
|
// return
|
|
//}
|
|
//defer skdg.Redis.Del(key)
|
|
start := time.Now().Unix()
|
|
total, err := f(task)
|
|
if err != nil {
|
|
log.Error("CallExportFunc fail", log.Any("task", task), log.E(err))
|
|
// 修改任务信息
|
|
export_task_mod.UpdateTask(task.ID, bson.M{
|
|
"status": 2,
|
|
"reason": err.Error(),
|
|
})
|
|
return
|
|
}
|
|
// 修改任务信息
|
|
export_task_mod.UpdateTask(task.ID, bson.M{
|
|
"status": 1,
|
|
"time": time.Now().Unix() - start, // 耗时
|
|
"total": total,
|
|
})
|
|
})
|
|
return
|
|
}
|
|
|
|
// SendTg 发送到tg
|
|
func SendTg(admin string, excelList []string, zipFileName string) (err error) {
|
|
exportPath, err := filepath.Abs("temp")
|
|
if err != nil {
|
|
return
|
|
}
|
|
// 尝试创建
|
|
file.MakeDir(exportPath)
|
|
zipList, err := Zip(excelList, zipFileName)
|
|
if err != nil {
|
|
log.Error("SendTg Zip fail", log.Any("excelList", excelList), log.Any("zipFileName", zipFileName), log.E(err))
|
|
return
|
|
}
|
|
// 测试环境不需要发送到tg群
|
|
//if skdg.Conf.Base.Env == "test" {
|
|
// return
|
|
//}
|
|
if len(zipList) == 0 {
|
|
return
|
|
}
|
|
var chatId int64 = -1003399433452
|
|
|
|
content := fmt.Sprintf("%v总共有%v个压缩包，请查收", admin, len(zipList))
|
|
_, err = Bot.Send(tgbotapi.NewMessage(chatId, content))
|
|
if err != nil {
|
|
log.Error("tgbot send fail", log.E(err))
|
|
return errors.Wrap(err, "发送TG消息失败")
|
|
}
|
|
for _, zipName := range zipList {
|
|
fileBytes, err := os.ReadFile("./temp/" + zipName)
|
|
if err != nil {
|
|
log.Error("ReadFile fail", log.Any("excelList", excelList), log.Any("zipFileName", zipFileName), log.E(err))
|
|
return errors.Wrap(err, "打开压缩文件失败")
|
|
}
|
|
|
|
// 创建要发送的文件
|
|
tFile := tgbotapi.FileBytes{Name: zipName, Bytes: fileBytes}
|
|
|
|
//生成csv发送到tg群
|
|
newDocument := tgbotapi.NewDocumentUpload(chatId, tFile)
|
|
_, err = Bot.Send(newDocument)
|
|
if err != nil {
|
|
log.Error("tgbot send fail", log.E(err))
|
|
return errors.Wrap(err, "发送TG消息失败")
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func Zip(filePathList []string, zipFileName string) (zipList []string, err error) {
|
|
if len(filePathList) == 0 {
|
|
return
|
|
}
|
|
var currentSize int64
|
|
zipIndex := 1
|
|
currentZipName := fmt.Sprintf("%v_%v.zip", zipFileName, zipIndex)
|
|
args := []string{currentZipName}
|
|
for k, filePath := range filePathList {
|
|
// 获取文件大小
|
|
stat, err := os.Stat(filePath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
currentSize = currentSize + stat.Size()
|
|
args = append(args, filePath)
|
|
// 判断当前所有文件加起来是否已经超过20M，超过就直接压缩,否则就接着等待压缩
|
|
if currentSize < maxSize && k != (len(filePathList)-1) {
|
|
continue
|
|
}
|
|
cmd1 := exec.Command("zip", args...)
|
|
out1, err := cmd1.CombinedOutput()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
fmt.Println(string(out1))
|
|
// 移动文件到目录里
|
|
cmd2 := exec.Command("mv", currentZipName, "./temp")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out2, err := cmd2.CombinedOutput()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
fmt.Println(string(out2))
|
|
zipList = append(zipList, currentZipName)
|
|
// 重置
|
|
zipIndex = zipIndex + 1
|
|
currentZipName = fmt.Sprintf("%v_%v.zip", zipFileName, zipIndex)
|
|
args = []string{currentZipName}
|
|
currentSize = 0
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// SaveExcel 保存数据到excel中
|
|
func SaveExcel(records interface{}, path string, fileName string) (filePath string, err error) {
|
|
|
|
sheet := "sheet1"
|
|
xlsx := excelize.NewFile() // new file
|
|
index, _ := xlsx.NewSheet(sheet) // new sheet
|
|
xlsx.SetActiveSheet(index) // set active (default) sheet
|
|
t := reflect.TypeOf(records)
|
|
if t.Kind() != reflect.Slice {
|
|
panic("records must be slice")
|
|
}
|
|
|
|
s := reflect.ValueOf(records)
|
|
for i := 0; i < s.Len(); i++ {
|
|
elem := s.Index(i).Interface()
|
|
elemType := reflect.TypeOf(elem)
|
|
elemValue := reflect.ValueOf(elem)
|
|
if elemType.Kind() == reflect.Ptr {
|
|
elemType = elemType.Elem()
|
|
elemValue = elemValue.Elem()
|
|
}
|
|
if elemType.Kind() != reflect.Struct {
|
|
panic("record in slice must be a struct")
|
|
}
|
|
k := 0
|
|
for j := 0; j < elemType.NumField(); j++ {
|
|
field := elemType.Field(j)
|
|
tag := field.Tag.Get("xlsx")
|
|
if tag == "-" {
|
|
continue
|
|
}
|
|
if tag == "" {
|
|
tag = field.Name
|
|
}
|
|
column, _ := excelize.ColumnNumberToName(k + 1)
|
|
k++
|
|
name := tag
|
|
// 设置表头
|
|
if i == 0 {
|
|
_ = xlsx.SetCellValue(sheet, fmt.Sprintf("%s%d", column, i+1), name)
|
|
}
|
|
// 判断是否是时间类型
|
|
|
|
colValue := elemValue.Field(j).Interface()
|
|
if colTime, ok := colValue.(time.Time); ok {
|
|
colValue = colTime.Format("2006-01-02 15:04:05")
|
|
}
|
|
// 设置内容
|
|
_ = xlsx.SetCellValue(sheet, fmt.Sprintf("%s%d", column, i+2), colValue)
|
|
}
|
|
}
|
|
// 保存的目录看是否存在不存在就创建
|
|
_, err = os.Stat(path)
|
|
if os.IsNotExist(err) {
|
|
err = os.MkdirAll(path, os.ModePerm)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
filePath = path + "/" + fileName
|
|
err = xlsx.SaveAs(filePath)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return filePath, nil
|
|
}
|