mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-05-20 23:59:34 +08:00
feat: add import command and batch import functionality
- Implemented the `/import` command to allow users to import files from storage to Telegram. - Added support for listing files in storage and filtering based on regex patterns. - Created a batch import task to handle multiple file uploads concurrently. - Introduced progress tracking for batch imports, providing real-time updates to users. - Enhanced storage interfaces to support file listing and reading capabilities. - Updated localization files for the new import command and its usage instructions. - Added utility functions for file size formatting and speed calculation. - Refactored Telegram storage handling to support reading from non-seekable streams.
This commit is contained in:
krau
180
client/bot/handlers/import.go
Normal file
180
client/bot/handlers/import.go
Normal file
@@ -0,0 +1,180 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/celestix/gotgproto/dispatcher"
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
storconfig "github.com/krau/SaveAny-Bot/config/storage"
|
||||
"github.com/krau/SaveAny-Bot/core"
|
||||
"github.com/krau/SaveAny-Bot/core/tasks/batchimport"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
"github.com/rs/xid"
|
||||
)
|
||||
|
||||
func handleImportCmd(ctx *ext.Context, update *ext.Update) error {
|
||||
logger := log.FromContext(ctx)
|
||||
args := strings.Split(update.EffectiveMessage.Text, " ")
|
||||
|
||||
if len(args) < 3 {
|
||||
ctx.Reply(update, ext.ReplyTextString("用法: /import <storage_name> <dir_path> [target_chat_id] [filter]\n\n"+
|
||||
"示例:\n"+
|
||||
"/import 本机1 /downloads\n"+
|
||||
"/import MyAlist /media/photos -1001234567890\n"+
|
||||
"/import MyLocal /backup \".*\\\\.mp4$\""), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
storageName := args[1]
|
||||
dirPath := args[2]
|
||||
|
||||
userID := update.GetUserChat().GetID()
|
||||
|
||||
// 1. 获取源存储端
|
||||
stor, err := storage.GetStorageByUserIDAndName(ctx, userID, storageName)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to get storage by user ID and name: %s", err)
|
||||
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("存储端 '%s' 不存在或您无权访问: %v", storageName, err)), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
// 2. 检查是否支持列举
|
||||
listable, ok := stor.(storage.StorageListable)
|
||||
if !ok {
|
||||
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("存储端 '%s' 不支持列举文件功能", storageName)), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
// 3. 检查是否支持读取
|
||||
_, ok = stor.(storage.StorageReadable)
|
||||
if !ok {
|
||||
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("存储端 '%s' 不支持读取文件功能", storageName)), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
// 4. 获取目标 Telegram 存储
|
||||
telegramStorage, err := storage.GetTelegramStorageByUserID(ctx, userID)
|
||||
if err != nil {
|
||||
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("未找到可用的 Telegram 存储: %v", err)), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
// 5. 列举目录文件
|
||||
replied, err := ctx.Reply(update, ext.ReplyTextString("正在获取文件列表..."), nil)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to reply: %s", err)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
files, err := listable.ListFiles(ctx, dirPath)
|
||||
if err != nil {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: replied.ID,
|
||||
Message: fmt.Sprintf("获取文件列表失败: %v", err),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
// 6. 过滤文件
|
||||
var filter *regexp.Regexp
|
||||
if len(args) >= 5 {
|
||||
filter, err = regexp.Compile(args[4])
|
||||
if err != nil {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: replied.ID,
|
||||
Message: fmt.Sprintf("正则表达式无效: %v", err),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
}
|
||||
|
||||
filteredFiles := make([]storagetypes.FileInfo, 0)
|
||||
for _, file := range files {
|
||||
if file.IsDir {
|
||||
continue
|
||||
}
|
||||
if filter != nil && !filter.MatchString(file.Name) {
|
||||
continue
|
||||
}
|
||||
filteredFiles = append(filteredFiles, file)
|
||||
}
|
||||
|
||||
if len(filteredFiles) == 0 {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: replied.ID,
|
||||
Message: "目录中没有可导入的文件",
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
// 7. 解析目标 Chat ID
|
||||
// Get default chat_id from Telegram storage config
|
||||
targetChatID := int64(0)
|
||||
if telegramCfg := config.C().GetStorageByName(telegramStorage.Name()); telegramCfg != nil {
|
||||
if tgCfg, ok := telegramCfg.(*storconfig.TelegramStorageConfig); ok {
|
||||
targetChatID = tgCfg.ChatID
|
||||
}
|
||||
}
|
||||
|
||||
if len(args) >= 4 {
|
||||
parsedChatID, err := tgutil.ParseChatID(ctx, args[3])
|
||||
if err != nil {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: replied.ID,
|
||||
Message: fmt.Sprintf("无效的 Chat ID: %v", err),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
targetChatID = parsedChatID
|
||||
}
|
||||
|
||||
if targetChatID == 0 {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: replied.ID,
|
||||
Message: "未指定目标频道 ID，且 Telegram 存储未配置默认 chat_id",
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
// 8. 创建任务元素
|
||||
elems := make([]batchimport.TaskElement, 0, len(filteredFiles))
|
||||
var totalSize int64
|
||||
for _, file := range filteredFiles {
|
||||
elem := batchimport.NewTaskElement(stor, file, telegramStorage, targetChatID)
|
||||
elems = append(elems, *elem)
|
||||
totalSize += file.Size
|
||||
}
|
||||
|
||||
// 9. 创建并添加任务
|
||||
taskID := xid.New().String()
|
||||
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
|
||||
task := batchimport.NewBatchImportTask(
|
||||
taskID,
|
||||
injectCtx,
|
||||
elems,
|
||||
batchimport.NewProgressTracker(replied.ID, userID),
|
||||
true, // IgnoreErrors
|
||||
)
|
||||
|
||||
if err := core.AddTask(injectCtx, task); err != nil {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: replied.ID,
|
||||
Message: fmt.Sprintf("添加任务失败: %v", err),
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
ID: replied.ID,
|
||||
Message: fmt.Sprintf("✅ 已添加 %d 个文件到导入队列\n总大小: %.2f MB\n任务 ID: %s", len(elems), float64(totalSize)/(1024*1024), taskID),
|
||||
})
|
||||
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
@@ -31,6 +31,7 @@ var CommandHandlers = []DescCommandHandler{
|
||||
{"dl", i18nk.BotMsgCmdDl, handleDlCmd},
|
||||
{"aria2dl", i18nk.BotMsgCmdAria2dl, handleAria2DlCmd},
|
||||
{"ytdlp", i18nk.BotMsgCmdYtdlp, handleYtdlpCmd},
|
||||
{"import", i18nk.BotMsgCmdImport, handleImportCmd},
|
||||
{"task", i18nk.BotMsgCmdTask, handleTaskCmd},
|
||||
{"cancel", i18nk.BotMsgCmdCancel, handleCancelCmd},
|
||||
{"config", i18nk.BotMsgCmdConfig, handleConfigCmd},
|
||||
|
||||
Reference in New Issue
Block a user