mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-08-28 19:56:41 +08:00
refactor: refactor task logic for better scalability (#76)
* refactor: a big refactor. wip * refactor: port handle file * refactor: place all handlers * fix: task info nil pointer * feat: enhance task progress tracking and context management * feat: cancel task * feat: stream mode * feat: silent mode * feat: dir cmd * refactor: remove unused old file * feat: rule cmd * feat: handle silent mode * feat: batch task * fix: batch task progress and temp file cleanup * refactor: update file creation and cleanup methods for better resource management * feat: add save command with silent mode handling * feat: message link * feat: update message prompts to include file count in storage selection * feat: slient save links * refactor: reduce dup code * feat: rule type * feat: chose dir * feat: refactor file handling and storage rules, improve error handling and logging * feat: rule mode * feat: telegraph pics * fix: tphpics nil pointer and inaccurate dirpath * feat: silent save telegraph * feat: add suffix to avoid file overwrite * feat: new storage telegram * chore: tidy go mod
This commit is contained in:
Krau
@@ -0,0 +1,193 @@
|
||||
// Some shortcuts for duplicate code in handlers, they should return dispatcher errors
|
||||
package shortcut
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/celestix/gotgproto/dispatcher"
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/celestix/gotgproto/types"
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/mediautil"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/re"
|
||||
"github.com/krau/SaveAny-Bot/common/cache"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tphutil"
|
||||
"github.com/krau/SaveAny-Bot/pkg/telegraph"
|
||||
"github.com/krau/SaveAny-Bot/pkg/tfile"
|
||||
)
|
||||
|
||||
// 获取消息中的文件并回复等待消息, 返回等待消息, 获取到的文件
|
||||
func GetFileFromMessageWithReply(ctx *ext.Context, update *ext.Update, message *tg.Message, tfileopts ...tfile.TGFileOptions) (replied *types.Message,
|
||||
file tfile.TGFileMessage, err error,
|
||||
) {
|
||||
logger := log.FromContext(ctx)
|
||||
media := message.Media
|
||||
supported := mediautil.IsSupported(media)
|
||||
if !supported {
|
||||
ctx.Reply(update, ext.ReplyTextString("不支持的消息类型"), nil)
|
||||
return nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
|
||||
replied, err = ctx.Reply(update, ext.ReplyTextString("正在获取文件信息..."), nil)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to reply: %s", err)
|
||||
return nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
options := []tfile.TGFileOptions{
|
||||
tfile.WithMessage(message),
|
||||
}
|
||||
if len(tfileopts) > 0 {
|
||||
options = append(options, tfileopts...)
|
||||
} else {
|
||||
options = append(options, tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*message)))
|
||||
}
|
||||
file, err = tfile.FromMediaMessage(media, message, options...)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to get file from media: %s", err)
|
||||
ctx.Reply(update, ext.ReplyTextString("获取文件失败: "+err.Error()), nil)
|
||||
return nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
return replied, file, nil
|
||||
}
|
||||
|
||||
type EditMessageFunc func(text string, markup tg.ReplyMarkupClass)
|
||||
|
||||
// 获取链接中的文件并回复等待消息
|
||||
func GetFilesFromUpdateLinkMessageWithReplyEdit(ctx *ext.Context, update *ext.Update) (replied *types.Message, files []tfile.TGFileMessage, editReplied EditMessageFunc, err error) {
|
||||
logger := log.FromContext(ctx)
|
||||
msgLinks := re.TgMessageLinkRegexp.FindAllString(update.EffectiveMessage.GetMessage(), -1)
|
||||
if len(msgLinks) == 0 {
|
||||
logger.Warn("no matched message links but called handleMessageLink")
|
||||
return nil, nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
replied, err = ctx.Reply(update, ext.ReplyTextString("正在获取消息..."), nil)
|
||||
if err != nil {
|
||||
logger.Errorf("failed to reply: %s", err)
|
||||
return nil, nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
editReplied = func(text string, markup tg.ReplyMarkupClass) {
|
||||
if _, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: replied.ID,
|
||||
Message: text,
|
||||
ReplyMarkup: markup,
|
||||
}); err != nil {
|
||||
logger.Errorf("failed to edit message: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
files = make([]tfile.TGFileMessage, 0, len(msgLinks))
|
||||
for _, link := range msgLinks {
|
||||
chatId, msgId, err := tgutil.ParseMessageLink(ctx, link)
|
||||
if err != nil {
|
||||
logger.Errorf("failed to parse message link %s: %s", link, err)
|
||||
continue
|
||||
}
|
||||
msg, err := tgutil.GetMessageByID(ctx, chatId, msgId)
|
||||
if err != nil {
|
||||
logger.Errorf("failed to get message by ID: %s", err)
|
||||
continue
|
||||
}
|
||||
media, ok := msg.GetMedia()
|
||||
if !ok {
|
||||
logger.Debugf("message %d has no media", msg.GetID())
|
||||
continue
|
||||
}
|
||||
file, err := tfile.FromMediaMessage(media, msg, tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*msg)))
|
||||
if err != nil {
|
||||
logger.Errorf("failed to create file from media: %s", err)
|
||||
continue
|
||||
}
|
||||
files = append(files, file)
|
||||
}
|
||||
if len(files) == 0 {
|
||||
editReplied("没有找到可保存的文件", nil)
|
||||
return nil, nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
return replied, files, editReplied, nil
|
||||
}
|
||||
|
||||
func GetCallbackDataWithAnswer[DataType any](ctx *ext.Context, update *ext.Update, dataid string) (DataType, error) {
|
||||
data, ok := cache.Get[DataType](dataid)
|
||||
if !ok {
|
||||
log.FromContext(ctx).Warnf("Invalid data ID: %s", dataid)
|
||||
queryID := update.CallbackQuery.GetQueryID()
|
||||
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(queryID, "数据已过期或无效"))
|
||||
var zero DataType
|
||||
return zero, dispatcher.EndGroups
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
type TelegraphResult struct {
|
||||
Pics []string `json:"pics"` // image urls
|
||||
TphDir string `json:"tph_dir"` // telegraph path, unescaped
|
||||
Page *telegraph.Page `json:"page"` // telegraph page node
|
||||
}
|
||||
|
||||
// return replied message, image urls, telegraph path(unescaped), error
|
||||
func GetTphPicsFromMessageWithReply(ctx *ext.Context, update *ext.Update) (*types.Message, *TelegraphResult, error) {
|
||||
logger := log.FromContext(ctx)
|
||||
tphurl := re.TelegraphUrlRegexp.FindString(update.EffectiveMessage.GetMessage()) // TODO: batch urls
|
||||
if tphurl == "" {
|
||||
logger.Warnf("No telegraph url found but called handleTelegraph")
|
||||
return nil, nil, dispatcher.ContinueGroups
|
||||
}
|
||||
pagepath := strings.Split(tphurl, "/")[len(strings.Split(tphurl, "/"))-1]
|
||||
tphdir, err := url.PathUnescape(pagepath)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to unescape telegraph path: %s", err)
|
||||
ctx.Reply(update, ext.ReplyTextString("解析 telegraph 路径失败: "+err.Error()), nil)
|
||||
return nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
msg, err := ctx.Reply(update, ext.ReplyTextString("正在获取 telegraph 页面..."), nil)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to reply to update: %s", err)
|
||||
return nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
page, err := tphutil.DefaultClient().GetPage(ctx, pagepath)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to get telegraph page: %s", err)
|
||||
ctx.Reply(update, ext.ReplyTextString("获取 telegraph 页面失败: "+err.Error()), nil)
|
||||
return nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
imgs := make([]string, 0)
|
||||
for _, elem := range page.Content {
|
||||
var node telegraph.NodeElement
|
||||
data, err := json.Marshal(elem)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to marshal element: %s", err)
|
||||
continue
|
||||
}
|
||||
err = json.Unmarshal(data, &node)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to unmarshal element: %s", err)
|
||||
continue
|
||||
}
|
||||
|
||||
if len(node.Children) != 0 {
|
||||
for _, child := range node.Children {
|
||||
imgs = append(imgs, tphutil.GetNodeImages(child)...)
|
||||
}
|
||||
}
|
||||
if node.Tag == "img" {
|
||||
if src, ok := node.Attrs["src"]; ok {
|
||||
imgs = append(imgs, src)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(imgs) == 0 {
|
||||
logger.Warn("No images found in telegraph page")
|
||||
ctx.Reply(update, ext.ReplyTextString("在 telegraph 页面中未找到图片"), nil)
|
||||
return nil, nil, dispatcher.EndGroups
|
||||
}
|
||||
return msg, &TelegraphResult{
|
||||
Pics: imgs,
|
||||
TphDir: tphdir,
|
||||
Page: page,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
package shortcut
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path"
|
||||
|
||||
"github.com/celestix/gotgproto/dispatcher"
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/ruleutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
"github.com/krau/SaveAny-Bot/core"
|
||||
"github.com/krau/SaveAny-Bot/core/batchtftask"
|
||||
"github.com/krau/SaveAny-Bot/core/tftask"
|
||||
"github.com/krau/SaveAny-Bot/database"
|
||||
"github.com/krau/SaveAny-Bot/pkg/tfile"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
"github.com/rs/xid"
|
||||
)
|
||||
|
||||
// 创建一个 tftask.TGFileTask 并添加到任务队列中, 以编辑消息的方式反馈结果
|
||||
func CreateAndAddTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor storage.Storage, dirPath string, file tfile.TGFileMessage, trackMsgID int) error {
|
||||
logger := log.FromContext(ctx)
|
||||
user, err := database.GetUserByChatID(ctx, userID)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to get user by chat ID: %s", err)
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: "获取用户失败: " + err.Error(),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
if user.ApplyRule && user.Rules != nil {
|
||||
matchedStorageName, matchedDirPath := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
|
||||
dirPath = matchedDirPath
|
||||
if matchedStorageName.IsValid() {
|
||||
stor, err = storage.GetStorageByUserIDAndName(ctx, user.ChatID, matchedStorageName.String())
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to get storage by user ID and name: %s", err)
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: "获取存储失败: " + err.Error(),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
storagePath := stor.JoinStoragePath(path.Join(dirPath, file.Name()))
|
||||
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
|
||||
taskid := xid.New().String()
|
||||
task, err := tftask.NewTGFileTask(taskid, injectCtx, file, ctx.Raw, stor, storagePath,
|
||||
tftask.NewProgressTrack(
|
||||
trackMsgID,
|
||||
userID))
|
||||
if err != nil {
|
||||
logger.Errorf("create task failed: %s", err)
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: "创建任务失败: " + err.Error(),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
if err := core.AddTask(injectCtx, task); err != nil {
|
||||
logger.Errorf("add task failed: %s", err)
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: "添加任务失败: " + err.Error(),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
text, entities := msgelem.BuildTaskAddedEntities(ctx, file.Name(), core.GetLength(injectCtx))
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: text,
|
||||
Entities: entities,
|
||||
})
|
||||
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
// 创建一个 batchtftask.BatchTGFileTask 并添加到任务队列中, 以编辑消息的方式反馈结果
|
||||
func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor storage.Storage, dirPath string, files []tfile.TGFileMessage, trackMsgID int) error {
|
||||
logger := log.FromContext(ctx)
|
||||
user, err := database.GetUserByChatID(ctx, userID)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to get user by chat ID: %s", err)
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: "获取用户失败: " + err.Error(),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
useRule := user.ApplyRule && user.Rules != nil
|
||||
applyRule := func(file tfile.TGFileMessage) (string, string) {
|
||||
if !useRule {
|
||||
return stor.Name(), dirPath
|
||||
}
|
||||
storName, dirP := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
|
||||
if !storName.IsValid() {
|
||||
return stor.Name(), dirP
|
||||
}
|
||||
return storName.String(), dirP
|
||||
}
|
||||
|
||||
elems := make([]batchtftask.TaskElement, 0, len(files))
|
||||
for _, file := range files {
|
||||
storName, dirPath := applyRule(file)
|
||||
fileStor := stor
|
||||
if storName != stor.Name() && storName != "" {
|
||||
fileStor, err = storage.GetStorageByUserIDAndName(ctx, user.ChatID, storName)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to get storage by user ID and name: %s", err)
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: "获取存储失败: " + err.Error(),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
}
|
||||
storPath := fileStor.JoinStoragePath(path.Join(dirPath, file.Name()))
|
||||
elem, err := batchtftask.NewTaskElement(fileStor, storPath, file)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to create task element: %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)
|
||||
taskid := xid.New().String()
|
||||
task := batchtftask.NewBatchTGFileTask(taskid, injectCtx, elems, ctx.Raw, batchtftask.NewProgressTracker(trackMsgID, userID), true)
|
||||
if err := core.AddTask(injectCtx, task); err != nil {
|
||||
logger.Errorf("Failed to add batch task: %s", err)
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: "批量任务添加失败: " + err.Error(),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: fmt.Sprintf("已添加批量任务, 共 %d 个文件", len(files)),
|
||||
ReplyMarkup: nil,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package shortcut
|
||||
|
||||
import (
|
||||
"github.com/celestix/gotgproto/dispatcher"
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tphutil"
|
||||
"github.com/krau/SaveAny-Bot/core"
|
||||
"github.com/krau/SaveAny-Bot/core/tphtask"
|
||||
"github.com/krau/SaveAny-Bot/pkg/telegraph"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
"github.com/rs/xid"
|
||||
)
|
||||
|
||||
func CreateAndAddTphTaskWithEdit(ctx *ext.Context,
|
||||
userID int64,
|
||||
tphpage *telegraph.Page,
|
||||
dirPath string, // unescaped ph path for file storage
|
||||
pics []string,
|
||||
stor storage.Storage,
|
||||
trackMsgID int) error {
|
||||
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
|
||||
task := tphtask.NewTask(xid.New().String(),
|
||||
injectCtx,
|
||||
tphpage.Path,
|
||||
pics,
|
||||
stor,
|
||||
stor.JoinStoragePath(dirPath),
|
||||
tphutil.DefaultClient(),
|
||||
tphtask.NewProgress(trackMsgID, userID),
|
||||
)
|
||||
if err := core.AddTask(injectCtx, task); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to add task: %s", err)
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: "任务添加失败: " + err.Error(),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
text, entities := msgelem.BuildTaskAddedEntities(ctx, tphpage.Title, core.GetLength(ctx))
|
||||
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
|
||||
ID: trackMsgID,
|
||||
Message: text,
|
||||
Entities: entities,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
Reference in New Issue
Block a user