mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-05-22 00:30:36 +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:
@@ -6,6 +6,7 @@ import (
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
)
|
||||
|
||||
var UserStorages = make(map[int64][]Storage)
|
||||
@@ -79,3 +80,14 @@ func LoadStorages(ctx context.Context) {
|
||||
UserStorages[int64(user)] = GetUserStorages(ctx, int64(user))
|
||||
}
|
||||
}
|
||||
|
||||
// GetTelegramStorageByUserID returns the first enabled Telegram storage for the user
|
||||
func GetTelegramStorageByUserID(ctx context.Context, chatID int64) (Storage, error) {
|
||||
storages := GetUserStorages(ctx, chatID)
|
||||
for _, stor := range storages {
|
||||
if stor.Type() == storenum.Telegram {
|
||||
return stor, nil
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("no telegram storage found for user %d", chatID)
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"github.com/duke-git/lancet/v2/fileutil"
|
||||
config "github.com/krau/SaveAny-Bot/config/storage"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
)
|
||||
|
||||
type Local struct {
|
||||
@@ -81,3 +82,51 @@ func (l *Local) Exists(ctx context.Context, storagePath string) bool {
|
||||
}
|
||||
return fileutil.IsExist(absPath)
|
||||
}
|
||||
|
||||
// ListFiles implements StorageListable interface
|
||||
func (l *Local) ListFiles(ctx context.Context, dirPath string) ([]storagetypes.FileInfo, error) {
|
||||
absPath := l.JoinStoragePath(dirPath)
|
||||
|
||||
entries, err := os.ReadDir(absPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read directory %s: %w", absPath, err)
|
||||
}
|
||||
|
||||
files := make([]storagetypes.FileInfo, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
l.logger.Warnf("Failed to get file info for %s: %v", entry.Name(), err)
|
||||
continue
|
||||
}
|
||||
|
||||
filePath := filepath.Join(dirPath, entry.Name())
|
||||
files = append(files, storagetypes.FileInfo{
|
||||
Name: entry.Name(),
|
||||
Path: filePath,
|
||||
Size: info.Size(),
|
||||
IsDir: entry.IsDir(),
|
||||
ModTime: info.ModTime(),
|
||||
})
|
||||
}
|
||||
|
||||
return files, nil
|
||||
}
|
||||
|
||||
// OpenFile implements StorageReadable interface
|
||||
func (l *Local) OpenFile(ctx context.Context, filePath string) (io.ReadCloser, int64, error) {
|
||||
absPath := l.JoinStoragePath(filePath)
|
||||
|
||||
file, err := os.Open(absPath)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("failed to open file %s: %w", absPath, err)
|
||||
}
|
||||
|
||||
stat, err := file.Stat()
|
||||
if err != nil {
|
||||
file.Close()
|
||||
return nil, 0, fmt.Errorf("failed to stat file %s: %w", absPath, err)
|
||||
}
|
||||
|
||||
return file, stat.Size(), nil
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
|
||||
storcfg "github.com/krau/SaveAny-Bot/config/storage"
|
||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
"github.com/krau/SaveAny-Bot/storage/alist"
|
||||
"github.com/krau/SaveAny-Bot/storage/local"
|
||||
"github.com/krau/SaveAny-Bot/storage/minio"
|
||||
@@ -30,6 +31,18 @@ type StorageCannotStream interface {
|
||||
CannotStream() string
|
||||
}
|
||||
|
||||
// StorageListable 表示支持列举目录内容的存储
|
||||
type StorageListable interface {
|
||||
Storage
|
||||
ListFiles(ctx context.Context, dirPath string) ([]storagetypes.FileInfo, error)
|
||||
}
|
||||
|
||||
// StorageReadable 表示支持读取文件内容的存储
|
||||
type StorageReadable interface {
|
||||
Storage
|
||||
OpenFile(ctx context.Context, filePath string) (io.ReadCloser, int64, error)
|
||||
}
|
||||
|
||||
var Storages = make(map[string]Storage)
|
||||
|
||||
type StorageConstructor func() Storage
|
||||
|
||||
@@ -99,12 +99,6 @@ func (w *splitWriter) finalize() error {
|
||||
}
|
||||
|
||||
func CreateSplitZip(ctx context.Context, reader io.Reader, size int64, fileName, outputBase string, partSize int64) error {
|
||||
// seek the reader if possible
|
||||
if rs, ok := reader.(io.ReadSeeker); ok {
|
||||
if _, err := rs.Seek(0, io.SeekStart); err != nil {
|
||||
return fmt.Errorf("failed to seek reader: %w", err)
|
||||
}
|
||||
}
|
||||
outputDir := filepath.Dir(outputBase)
|
||||
if err := os.MkdirAll(outputDir, os.ModePerm); err != nil {
|
||||
return fmt.Errorf("failed to create output directory: %w", err)
|
||||
|
||||
@@ -92,9 +92,6 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er
|
||||
return nil
|
||||
}
|
||||
rs, seekable := r.(io.ReadSeeker)
|
||||
if !seekable || rs == nil {
|
||||
return fmt.Errorf("reader must implement io.ReadSeeker")
|
||||
}
|
||||
splitSize := t.config.SplitSizeMB * 1024 * 1024
|
||||
if splitSize <= 0 {
|
||||
splitSize = DefaultSplitSize
|
||||
@@ -123,88 +120,96 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er
|
||||
}
|
||||
chatID = cid
|
||||
}
|
||||
mtype, err := mimetype.DetectReader(rs)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to detect mimetype: %w", err)
|
||||
}
|
||||
if filename == "" {
|
||||
filename = xid.New().String() + mtype.Extension()
|
||||
}
|
||||
upler := uploader.NewUploader(tctx.Raw).
|
||||
WithPartSize(tglimit.MaxUploadPartSize).
|
||||
WithThreads(dlutil.BestThreads(size, config.C().Threads))
|
||||
peer := tryGetInputPeer(tctx, chatID)
|
||||
if peer == nil || peer.Zero() {
|
||||
return fmt.Errorf("failed to get input peer for chat ID %d", chatID)
|
||||
}
|
||||
var mtype *mimetype.MIME
|
||||
if seekable {
|
||||
var err error
|
||||
mtype, err = mimetype.DetectReader(r)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to detect mimetype: %w", err)
|
||||
}
|
||||
if filename == "" {
|
||||
filename = xid.New().String() + mtype.Extension()
|
||||
}
|
||||
|
||||
if _, err := rs.Seek(0, io.SeekStart); err != nil {
|
||||
return fmt.Errorf("failed to seek reader: %w", err)
|
||||
if _, err := rs.Seek(0, io.SeekStart); err != nil {
|
||||
return fmt.Errorf("failed to seek reader: %w", err)
|
||||
}
|
||||
}
|
||||
upler := uploader.NewUploader(tctx.Raw).
|
||||
WithPartSize(tglimit.MaxUploadPartSize).
|
||||
WithThreads(dlutil.BestThreads(size, config.C().Threads))
|
||||
if size > splitSize {
|
||||
// large file, use split uploader
|
||||
return t.splitUpload(tctx, rs, filename, upler, peer, size, splitSize)
|
||||
return t.splitUpload(tctx, r, filename, upler, peer, size, splitSize)
|
||||
}
|
||||
|
||||
var file tg.InputFileClass
|
||||
if size < 0 {
|
||||
file, err = upler.FromReader(ctx, filename, rs)
|
||||
var err error
|
||||
if size <= 0 {
|
||||
file, err = upler.FromReader(ctx, filename, r)
|
||||
} else {
|
||||
file, err = upler.Upload(ctx, uploader.NewUpload(filename, rs, size))
|
||||
file, err = upler.Upload(ctx, uploader.NewUpload(filename, r, size))
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to upload file to telegram: %w", err)
|
||||
}
|
||||
caption := styling.Plain(filename)
|
||||
forceFile := t.config.ForceFile
|
||||
if strings.HasPrefix(mtype.String(), "image/") && size >= tglimit.MaxPhotoSize {
|
||||
|
||||
if mtype != nil && strings.HasPrefix(mtype.String(), "image/") && size >= tglimit.MaxPhotoSize {
|
||||
forceFile = true
|
||||
}
|
||||
doc := message.UploadedDocument(file, caption).
|
||||
Filename(filename).
|
||||
ForceFile(forceFile).
|
||||
MIME(mtype.String())
|
||||
|
||||
ForceFile(forceFile)
|
||||
if mtype != nil {
|
||||
doc = doc.MIME(mtype.String())
|
||||
}
|
||||
var media message.MediaOption = doc
|
||||
|
||||
switch mtypeStr := mtype.String(); {
|
||||
case strings.HasPrefix(mtypeStr, "video/"):
|
||||
media = doc.Video().SupportsStreaming()
|
||||
thumb, err := extractThumbFrame(rs)
|
||||
if err == nil {
|
||||
thumb, err := upler.FromBytes(ctx, "thumb.jpg", thumb)
|
||||
if mtype != nil && rs != nil {
|
||||
switch mtypeStr := mtype.String(); {
|
||||
case strings.HasPrefix(mtypeStr, "video/"):
|
||||
media = doc.Video().SupportsStreaming()
|
||||
thumb, err := extractThumbFrame(rs)
|
||||
if err == nil {
|
||||
doc = doc.Thumb(thumb)
|
||||
thumb, err := upler.FromBytes(ctx, "thumb.jpg", thumb)
|
||||
if err == nil {
|
||||
doc = doc.Thumb(thumb)
|
||||
}
|
||||
}
|
||||
rs.Seek(0, io.SeekStart)
|
||||
switch mtypeStr {
|
||||
case "video/mp4":
|
||||
info, err := getMP4Meta(rs)
|
||||
if err != nil {
|
||||
// Fallback to ffprobe if gomedia fails (e.g., malformed MP4)
|
||||
rs.Seek(0, io.SeekStart)
|
||||
info, err = getVideoMetadata(rs)
|
||||
}
|
||||
if err == nil {
|
||||
media = doc.Video().
|
||||
Duration(time.Duration(info.Duration)*time.Second).
|
||||
Resolution(info.Width, info.Height).
|
||||
SupportsStreaming()
|
||||
}
|
||||
default:
|
||||
info, err := getVideoMetadata(rs)
|
||||
if err == nil {
|
||||
media = doc.Video().
|
||||
Duration(time.Duration(info.Duration)*time.Second).
|
||||
Resolution(info.Width, info.Height).
|
||||
SupportsStreaming()
|
||||
}
|
||||
}
|
||||
case strings.HasPrefix(mtypeStr, "audio/"):
|
||||
media = doc.Audio().Title(filename)
|
||||
case strings.HasPrefix(mtypeStr, "image/") && !strings.HasSuffix(mtypeStr, "webp"):
|
||||
media = message.UploadedPhoto(file, caption)
|
||||
}
|
||||
rs.Seek(0, io.SeekStart)
|
||||
switch mtypeStr {
|
||||
case "video/mp4":
|
||||
info, err := getMP4Meta(rs)
|
||||
if err != nil {
|
||||
// Fallback to ffprobe if gomedia fails (e.g., malformed MP4)
|
||||
rs.Seek(0, io.SeekStart)
|
||||
info, err = getVideoMetadata(rs)
|
||||
}
|
||||
if err == nil {
|
||||
media = doc.Video().
|
||||
Duration(time.Duration(info.Duration)*time.Second).
|
||||
Resolution(info.Width, info.Height).
|
||||
SupportsStreaming()
|
||||
}
|
||||
default:
|
||||
info, err := getVideoMetadata(rs)
|
||||
if err == nil {
|
||||
media = doc.Video().
|
||||
Duration(time.Duration(info.Duration)*time.Second).
|
||||
Resolution(info.Width, info.Height).
|
||||
SupportsStreaming()
|
||||
}
|
||||
}
|
||||
case strings.HasPrefix(mtypeStr, "audio/"):
|
||||
media = doc.Audio().Title(filename)
|
||||
case strings.HasPrefix(mtypeStr, "image/") && !strings.HasSuffix(mtypeStr, "webp"):
|
||||
media = message.UploadedPhoto(file, caption)
|
||||
}
|
||||
sender := tctx.Sender
|
||||
_, err = sender.WithUploader(upler).To(peer).Media(ctx, media)
|
||||
@@ -215,7 +220,7 @@ func (t *Telegram) CannotStream() string {
|
||||
return "Telegram storage must use a ReaderSeeker"
|
||||
}
|
||||
|
||||
func (t *Telegram) splitUpload(ctx *ext.Context, rs io.ReadSeeker, filename string, upler *uploader.Uploader, peer tg.InputPeerClass, fileSize, splitSize int64) error {
|
||||
func (t *Telegram) splitUpload(ctx *ext.Context, r io.Reader, filename string, upler *uploader.Uploader, peer tg.InputPeerClass, fileSize, splitSize int64) error {
|
||||
tempId := xid.New().String()
|
||||
outputBase := filepath.Join(config.C().Temp.BasePath, tempId, strings.Split(filename, ".")[0])
|
||||
defer func() {
|
||||
@@ -224,7 +229,7 @@ func (t *Telegram) splitUpload(ctx *ext.Context, rs io.ReadSeeker, filename stri
|
||||
log.FromContext(ctx).Warnf("Failed to cleanup temp split files: %s", err)
|
||||
}
|
||||
}()
|
||||
if err := CreateSplitZip(ctx, rs, fileSize, filename, outputBase, splitSize); err != nil {
|
||||
if err := CreateSplitZip(ctx, r, fileSize, filename, outputBase, splitSize); err != nil {
|
||||
return fmt.Errorf("failed to create split zip: %w", err)
|
||||
}
|
||||
matched, err := filepath.Glob(outputBase + ".z*")
|
||||
|
||||
Reference in New Issue
Block a user