servicebase/pkg/tools/picture/thumb_task.go

package picture

import (
	"context"
	//cache "ServiceAll/modules/company/internal/client/redis"
	"servicebase/pkg/datasource"
	"servicebase/pkg/log"
	"servicebase/pkg/repo"
	"servicebase/pkg/tools"
	"fmt"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/disintegration/imaging"
	"github.com/pkg/errors"
	"github.com/spf13/viper"
	"go.uber.org/zap"
)

const (
	CurIndexName      = "CUR_INDEX_DATA_ATTACHMENT"
	CommonTypeImage   = "IMAGE"
	CommonTypeUnknown = "UNCLASSIFIED"

	pathThumb     = "thumb"
	FileTypeThumb = "jpg"
)

var typeOfCommonType map[string]string

func init() {
	typeOfCommonType = map[string]string{
		"jpg":  CommonTypeImage,
		"png":  CommonTypeImage,
		"jfif": CommonTypeImage,
		"jpeg": CommonTypeImage,
		"bmp":  CommonTypeImage,
		"webp": CommonTypeImage,
		"icon": CommonTypeImage,
	}
}

func ParseCommonType(_type string) string {
	if v, ok := typeOfCommonType[_type]; ok {
		return v
	}
	return CommonTypeUnknown
}

func StartThumbTask() {
	const lines = 10
	const dur = time.Second * 10
	for {
		if cnt := ThumbTask(lines); cnt != lines {
			time.Sleep(dur)
		}
	}
}

// ThumbTask 图片缩略图任务
// 定期扫描附件表，将未生成压缩图的图片生成压缩图
func ThumbTask(lines int) (cnt int) {
	var (
		pid  int64
		e    error
		ctx       = context.Background()
		w, h uint = 0, 100
	)
	if pid, e = curIndexPID(); e != nil {
		log.Error("ThumbTask init error", zap.String("error", e.Error()), zap.Int64("pid", pid))
		return
	}
	var list []struct {
		Pid        int64
		ID         string
		CommonType string
		FilePath   string
		FileType   string
	}
	if e = datasource.DB.WithContext(ctx).
		Raw("select * from data_attachment where pid>=? and common_type=? and file_thumb_path=? order by pid limit ?", pid, CommonTypeImage, "", lines).
		Find(&list).Error; e != nil {
		log.Error("ThumbTask select error", zap.String("error", e.Error()), zap.Int64("pid", pid))
		return
	}
	// todo with WaitGroup
	//userCache := cache.User()
	for _, item := range list {
		if e = DoThumb(item.FilePath, item.ID, item.CommonType, item.FileType, w, h); e != nil {
			log.Error("ThumbTask thumb error",
				zap.String("error", e.Error()),
				zap.Int64("pid", pid),
				zap.String("path", item.FilePath),
				zap.String("file", item.ID),
				zap.String("type", item.FileType))
			return
		}
		//userCache.Set(ctx, CurIndexName, item.Pid+1, 0)
		cnt++
	}
	return
}

// 获取开始生成缩略图的索引
func curIndexPID() (pid int64, e error) {
	// 取当前已压缩的位置
	//ctx := context.Background()
	//userCache := cache.User()
	//if pid, _ = userCache.Get(ctx, CurIndexName).Int64(); pid == 0 {
	//	// 当前不存在，从数据库统计
	//	if e = datasource.DB.WithContext(ctx).
	//		Raw("select pid from data_attachment where common_type=? and file_thumb_path=? order by pid limit 1", CommonTypeImage, "").
	//		First(&pid).Error; e != nil {
	//		return
	//	}
	//	userCache.Set(ctx, CurIndexName, pid, 0)
	//}
	return
}

func DoThumb(path, id, commonType, _type string, w, h uint) error {
	if w == 0 && h == 0 {
		return errors.New("尺寸错误")
	}
	commonType = strings.ToLower(commonType)
	// 得到后缀和
	image, err := imaging.Open(filepath.Join(viper.GetString("attachment.path"), path))
	if err != nil {
		return err
	}
	var suffix []byte
	if w > 0 {
		suffix = append(suffix, 'w')
		suffix = append(suffix, strconv.Itoa(int(w))...)
	}
	if h > 0 {
		suffix = append(suffix, 'h')
		suffix = append(suffix, strconv.Itoa(int(h))...)
	}
	_ = tools.Mkdir(filepath.Join(viper.GetString("attachment.path"), commonType), pathThumb)
	thumbFilename := filepath.Join(commonType, pathThumb, fmt.Sprintf("%s.%s.%s", id, suffix, FileTypeThumb))
	targetImage := imaging.Resize(image, int(w), int(h), imaging.Lanczos)
	return repo.Q.Transaction(func(tx *repo.Query) error {
		if _, e := tx.DataAttach.WithContext(context.Background()).
			Select(tx.DataAttach.FileThumbPath, tx.DataAttach.FileThumbType).
			Where(tx.DataAttach.ID.Eq(id)).
			UpdateSimple(
				tx.DataAttach.FileThumbPath.Value(thumbFilename),
				tx.DataAttach.FileThumbType.Value(FileTypeThumb),
			); e != nil {
			return e
		}
		return imaging.Save(targetImage, filepath.Join(filepath.Join(viper.GetString("attachment.path"), thumbFilename)))
	})
}

func Do(path, id, commonType, _type string, w, h uint) (string, error) {
	if w == 0 && h == 0 {
		return "", errors.New("尺寸错误")
	}
	commonType = strings.ToLower(commonType)
	// 得到后缀和
	image, err := imaging.Open(filepath.Join(viper.GetString("attachment.path"), path))
	if err != nil {
		return "", err
	}
	var suffix []byte
	if w > 0 {
		suffix = append(suffix, 'w')
		suffix = append(suffix, strconv.Itoa(int(w))...)
	}
	if h > 0 {
		suffix = append(suffix, 'h')
		suffix = append(suffix, strconv.Itoa(int(h))...)
	}
	_ = tools.Mkdir(filepath.Join(viper.GetString("attachment.path"), commonType), pathThumb)
	thumbFilename := filepath.Join(commonType, pathThumb, fmt.Sprintf("%s.%s.%s", id, suffix, FileTypeThumb))
	targetImage := imaging.Resize(image, int(w), int(h), imaging.Lanczos)
	return thumbFilename, imaging.Save(targetImage, filepath.Join(filepath.Join(viper.GetString("attachment.path"), thumbFilename)))
}