mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-08-05 20:46:42 +08:00
feat: implement album handling rules and refactor related logic
This commit is contained in:
@@ -3,6 +3,8 @@ package ruleutil
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
|
||||||
|
"github.com/duke-git/lancet/v2/convertor"
|
||||||
|
|
||||||
"github.com/charmbracelet/log"
|
"github.com/charmbracelet/log"
|
||||||
"github.com/krau/SaveAny-Bot/database"
|
"github.com/krau/SaveAny-Bot/database"
|
||||||
"github.com/krau/SaveAny-Bot/pkg/consts"
|
"github.com/krau/SaveAny-Bot/pkg/consts"
|
||||||
@@ -33,11 +35,22 @@ func (m matchedStorName) String() string {
|
|||||||
return string(m)
|
return string(m)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m matchedStorName) IsValid() bool {
|
// can we use this storage name directly?
|
||||||
|
func (m matchedStorName) IsUsable() bool {
|
||||||
return m != "" && m != consts.RuleStorNameChosen
|
return m != "" && m != consts.RuleStorNameChosen
|
||||||
}
|
}
|
||||||
|
|
||||||
func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (matchedStorageName matchedStorName, dirPath string) {
|
type MatchedDirPath string
|
||||||
|
|
||||||
|
func (m MatchedDirPath) String() string {
|
||||||
|
return string(m)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m MatchedDirPath) NeedNewForAlbum() bool {
|
||||||
|
return m != "" && m == consts.RuleDirPathNewForAlbum
|
||||||
|
}
|
||||||
|
|
||||||
|
func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (matchedStorageName matchedStorName, dirPath MatchedDirPath) {
|
||||||
if inputs == nil || len(rules) == 0 {
|
if inputs == nil || len(rules) == 0 {
|
||||||
return "", ""
|
return "", ""
|
||||||
}
|
}
|
||||||
@@ -56,7 +69,7 @@ func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (m
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if ok {
|
if ok {
|
||||||
dirPath = ru.StoragePath()
|
dirPath = MatchedDirPath(ru.StoragePath())
|
||||||
matchedStorageName = matchedStorName(ru.StorageName())
|
matchedStorageName = matchedStorName(ru.StorageName())
|
||||||
}
|
}
|
||||||
case ruleenum.MessageRegex.String():
|
case ruleenum.MessageRegex.String():
|
||||||
@@ -71,7 +84,26 @@ func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (m
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if ok {
|
if ok {
|
||||||
dirPath = ru.StoragePath()
|
dirPath = MatchedDirPath(ru.StoragePath())
|
||||||
|
matchedStorageName = matchedStorName(ru.StorageName())
|
||||||
|
}
|
||||||
|
case ruleenum.IsAlbum.String():
|
||||||
|
matchAlbum, err := convertor.ToBool(ur.Data)
|
||||||
|
if err != nil {
|
||||||
|
matchAlbum = false
|
||||||
|
}
|
||||||
|
ru, err := rule.NewRuleMediaType(ur.StorageName, ur.DirPath, matchAlbum)
|
||||||
|
if err != nil {
|
||||||
|
logger.Errorf("Failed to create rule: %s", err)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
ok, err := ru.Match(inputs.File.Message().GroupedID != 0)
|
||||||
|
if err != nil {
|
||||||
|
logger.Errorf("Failed to match rule: %s", err)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if ok {
|
||||||
|
dirPath = MatchedDirPath(ru.StoragePath())
|
||||||
matchedStorageName = matchedStorName(ru.StorageName())
|
matchedStorageName = matchedStorName(ru.StorageName())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package shortcut
|
|||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"path"
|
"path"
|
||||||
|
"strings"
|
||||||
|
|
||||||
"github.com/celestix/gotgproto/dispatcher"
|
"github.com/celestix/gotgproto/dispatcher"
|
||||||
"github.com/celestix/gotgproto/ext"
|
"github.com/celestix/gotgproto/ext"
|
||||||
@@ -34,8 +35,8 @@ func CreateAndAddTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor storage
|
|||||||
}
|
}
|
||||||
if user.ApplyRule && user.Rules != nil {
|
if user.ApplyRule && user.Rules != nil {
|
||||||
matchedStorageName, matchedDirPath := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
|
matchedStorageName, matchedDirPath := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
|
||||||
dirPath = matchedDirPath
|
dirPath = matchedDirPath.String()
|
||||||
if matchedStorageName.IsValid() {
|
if matchedStorageName.IsUsable() {
|
||||||
stor, err = storage.GetStorageByUserIDAndName(ctx, user.ChatID, matchedStorageName.String())
|
stor, err = storage.GetStorageByUserIDAndName(ctx, user.ChatID, matchedStorageName.String())
|
||||||
if err != nil {
|
if err != nil {
|
||||||
logger.Errorf("Failed to get storage by user ID and name: %s", err)
|
logger.Errorf("Failed to get storage by user ID and name: %s", err)
|
||||||
@@ -93,19 +94,28 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st
|
|||||||
})
|
})
|
||||||
return dispatcher.EndGroups
|
return dispatcher.EndGroups
|
||||||
}
|
}
|
||||||
|
|
||||||
useRule := user.ApplyRule && user.Rules != nil
|
useRule := user.ApplyRule && user.Rules != nil
|
||||||
applyRule := func(file tfile.TGFileMessage) (string, string) {
|
|
||||||
|
applyRule := func(file tfile.TGFileMessage) (string, ruleutil.MatchedDirPath) {
|
||||||
if !useRule {
|
if !useRule {
|
||||||
return stor.Name(), dirPath
|
return stor.Name(), ruleutil.MatchedDirPath(dirPath)
|
||||||
}
|
}
|
||||||
storName, dirP := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
|
storName, dirP := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
|
||||||
if !storName.IsValid() {
|
|
||||||
return stor.Name(), dirP
|
storname := storName.String()
|
||||||
|
if !storName.IsUsable() {
|
||||||
|
storname = stor.Name()
|
||||||
}
|
}
|
||||||
return storName.String(), dirP
|
return storname, dirP
|
||||||
}
|
}
|
||||||
|
|
||||||
elems := make([]batchtftask.TaskElement, 0, len(files))
|
elems := make([]batchtftask.TaskElement, 0, len(files))
|
||||||
|
type albumFile struct {
|
||||||
|
file tfile.TGFileMessage
|
||||||
|
storage storage.Storage
|
||||||
|
}
|
||||||
|
albumFiles := make(map[int64][]albumFile, 0)
|
||||||
for _, file := range files {
|
for _, file := range files {
|
||||||
storName, dirPath := applyRule(file)
|
storName, dirPath := applyRule(file)
|
||||||
fileStor := stor
|
fileStor := stor
|
||||||
@@ -120,18 +130,56 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st
|
|||||||
return dispatcher.EndGroups
|
return dispatcher.EndGroups
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
storPath := fileStor.JoinStoragePath(path.Join(dirPath, file.Name()))
|
if !dirPath.NeedNewForAlbum() {
|
||||||
elem, err := batchtftask.NewTaskElement(fileStor, storPath, file)
|
storPath := fileStor.JoinStoragePath(path.Join(dirPath.String(), file.Name()))
|
||||||
if err != nil {
|
elem, err := batchtftask.NewTaskElement(fileStor, storPath, file)
|
||||||
logger.Errorf("Failed to create task element: %s", err)
|
if err != nil {
|
||||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
logger.Errorf("Failed to create task element: %s", err)
|
||||||
ID: trackMsgID,
|
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||||
Message: "任务创建失败: " + err.Error(),
|
ID: trackMsgID,
|
||||||
|
Message: "任务创建失败: " + err.Error(),
|
||||||
|
})
|
||||||
|
return dispatcher.EndGroups
|
||||||
|
}
|
||||||
|
elems = append(elems, *elem)
|
||||||
|
} else {
|
||||||
|
groupId, isGroup := file.Message().GetGroupedID()
|
||||||
|
if !isGroup || groupId == 0 {
|
||||||
|
logger.Warnf("File %s is not in a group, skipping album handling", file.Name())
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := albumFiles[groupId]; !ok {
|
||||||
|
albumFiles[groupId] = make([]albumFile, 0)
|
||||||
|
}
|
||||||
|
albumFiles[groupId] = append(albumFiles[groupId], albumFile{
|
||||||
|
file: file,
|
||||||
|
storage: fileStor,
|
||||||
})
|
})
|
||||||
return dispatcher.EndGroups
|
|
||||||
}
|
}
|
||||||
elems = append(elems, *elem)
|
|
||||||
}
|
}
|
||||||
|
for _, afiles := range albumFiles {
|
||||||
|
if len(afiles) <= 1 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
// 对于需要新建目录的文件, 将第一个文件的文件名(去除扩展名)作为目录名
|
||||||
|
// 存储以第一个文件的存储为准
|
||||||
|
albumDir := strings.TrimSuffix(path.Base(afiles[0].file.Name()), path.Ext(afiles[0].file.Name()))
|
||||||
|
albumStor := afiles[0].storage
|
||||||
|
for _, af := range afiles {
|
||||||
|
afstorPath := af.storage.JoinStoragePath(path.Join(dirPath, albumDir, af.file.Name()))
|
||||||
|
elem, err := batchtftask.NewTaskElement(albumStor, afstorPath, af.file)
|
||||||
|
if err != nil {
|
||||||
|
logger.Errorf("Failed to create task element for album file: %s", err)
|
||||||
|
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||||
|
ID: trackMsgID,
|
||||||
|
Message: "任务创建失败: " + err.Error(),
|
||||||
|
})
|
||||||
|
return dispatcher.EndGroups
|
||||||
|
}
|
||||||
|
elems = append(elems, *elem)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
|
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
|
||||||
taskid := xid.New().String()
|
taskid := xid.New().String()
|
||||||
task := batchtftask.NewBatchTGFileTask(taskid, injectCtx, elems, batchtftask.NewProgressTracker(trackMsgID, userID), true)
|
task := batchtftask.NewBatchTGFileTask(taskid, injectCtx, elems, batchtftask.NewProgressTracker(trackMsgID, userID), true)
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
package consts
|
package consts
|
||||||
|
|
||||||
const (
|
const (
|
||||||
RuleStorNameChosen = "CHOSEN"
|
RuleStorNameChosen = "CHOSEN"
|
||||||
|
RuleDirPathNewForAlbum = "NEW-FOR-ALBUM" // create a new directory for album files
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ type RuleType string
|
|||||||
const (
|
const (
|
||||||
FileNameRegex RuleType = "FILENAME-REGEX"
|
FileNameRegex RuleType = "FILENAME-REGEX"
|
||||||
MessageRegex RuleType = "MESSAGE-REGEX"
|
MessageRegex RuleType = "MESSAGE-REGEX"
|
||||||
|
IsAlbum RuleType = "IS-ALBUM"
|
||||||
)
|
)
|
||||||
|
|
||||||
func (r RuleType) String() string {
|
func (r RuleType) String() string {
|
||||||
@@ -12,5 +13,5 @@ func (r RuleType) String() string {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func Values() []RuleType {
|
func Values() []RuleType {
|
||||||
return []RuleType{FileNameRegex, MessageRegex}
|
return []RuleType{FileNameRegex, MessageRegex, IsAlbum}
|
||||||
}
|
}
|
||||||
|
|||||||
38
pkg/rule/is_album.go
Normal file
38
pkg/rule/is_album.go
Normal file
@@ -0,0 +1,38 @@
|
|||||||
|
package rule
|
||||||
|
|
||||||
|
import (
|
||||||
|
ruleenum "github.com/krau/SaveAny-Bot/pkg/enums/rule"
|
||||||
|
)
|
||||||
|
|
||||||
|
var _ RuleClass[bool] = (*RuleMediaType)(nil)
|
||||||
|
|
||||||
|
type RuleMediaType struct {
|
||||||
|
storInfo
|
||||||
|
matchAlbum bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r RuleMediaType) Type() ruleenum.RuleType {
|
||||||
|
return ruleenum.IsAlbum
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r RuleMediaType) Match(input bool) (bool, error) {
|
||||||
|
return r.matchAlbum == input, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r RuleMediaType) StorageName() string {
|
||||||
|
return r.storName
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r RuleMediaType) StoragePath() string {
|
||||||
|
return r.storPath
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewRuleMediaType(storName, storPath string, matchAlbum bool) (*RuleMediaType, error) {
|
||||||
|
return &RuleMediaType{
|
||||||
|
storInfo: storInfo{
|
||||||
|
storName: storName,
|
||||||
|
storPath: storPath,
|
||||||
|
},
|
||||||
|
matchAlbum: matchAlbum,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user