mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-08-06 13:03:19 +08:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e4144e73e6 | ||
|
|
1d794a7b9c | ||
|
|
52f880f0f2 |
@@ -2,6 +2,7 @@ package tgutil
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"sort"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"unicode"
|
"unicode"
|
||||||
@@ -359,9 +360,16 @@ func GetGroupedMessages(ctx *ext.Context, chatID int64, msg *tg.Message) ([]*tg.
|
|||||||
groupedMessages = append(groupedMessages, m)
|
groupedMessages = append(groupedMessages, m)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
sortMessagesByID(groupedMessages)
|
||||||
return groupedMessages, nil
|
return groupedMessages, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func sortMessagesByID(messages []*tg.Message) {
|
||||||
|
sort.Slice(messages, func(i, j int) bool {
|
||||||
|
return messages[i].GetID() < messages[j].GetID()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
func ExtractMessageEntityUrls(msg *tg.Message) []string {
|
func ExtractMessageEntityUrls(msg *tg.Message) []string {
|
||||||
if len(msg.Entities) == 0 {
|
if len(msg.Entities) == 0 {
|
||||||
return nil
|
return nil
|
||||||
|
|||||||
18
common/utils/tgutil/message_test.go
Normal file
18
common/utils/tgutil/message_test.go
Normal file
@@ -0,0 +1,18 @@
|
|||||||
|
package tgutil
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/gotd/td/tg"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSortMessagesByID(t *testing.T) {
|
||||||
|
messages := []*tg.Message{{ID: 9}, {ID: 3}, {ID: 7}}
|
||||||
|
sortMessagesByID(messages)
|
||||||
|
want := []int{3, 7, 9}
|
||||||
|
for i := range messages {
|
||||||
|
if messages[i].GetID() != want[i] {
|
||||||
|
t.Fatalf("message %d has ID %d, want %d", i, messages[i].GetID(), want[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -8,14 +8,16 @@ import (
|
|||||||
|
|
||||||
type TelegramStorageConfig struct {
|
type TelegramStorageConfig struct {
|
||||||
BaseConfig
|
BaseConfig
|
||||||
ChatID int64 `toml:"chat_id" mapstructure:"chat_id" json:"chat_id"`
|
ChatID int64 `toml:"chat_id" mapstructure:"chat_id" json:"chat_id"`
|
||||||
ForceFile bool `toml:"force_file" mapstructure:"force_file" json:"force_file"`
|
ForceFile bool `toml:"force_file" mapstructure:"force_file" json:"force_file"`
|
||||||
RateLimit int `toml:"rate_limit" mapstructure:"rate_limit" json:"rate_limit"`
|
RateLimit int `toml:"rate_limit" mapstructure:"rate_limit" json:"rate_limit"`
|
||||||
RateBurst int `toml:"rate_burst" mapstructure:"rate_burst" json:"rate_burst"`
|
RateBurst int `toml:"rate_burst" mapstructure:"rate_burst" json:"rate_burst"`
|
||||||
SkipLarge bool `toml:"skip_large" mapstructure:"skip_large" json:"skip_large"` // skip files larger than Telegram limit(2GB)
|
SkipLarge bool `toml:"skip_large" mapstructure:"skip_large" json:"skip_large"` // skip files larger than Telegram limit(2GB)
|
||||||
// split files larger than Telegram limit(2GB) into parts of specified size, in MB, leave 0 to set default(2000MB)
|
SplitLargeVideo bool `toml:"split_large_video" mapstructure:"split_large_video" json:"split_large_video"`
|
||||||
|
// split files larger than the uploader account limit into parts of specified size, in MB
|
||||||
|
// leave 0 to use the account limit (2000MB for bots/regular users, 4000MB for Premium users)
|
||||||
// only effective when SkipLarge is false
|
// only effective when SkipLarge is false
|
||||||
// use zip when splitting
|
// use zip when splitting non-video files or when lossless video splitting is disabled/unavailable
|
||||||
SplitSizeMB int64 `toml:"split_size_mb" mapstructure:"split_size_mb" json:"split_size_mb"`
|
SplitSizeMB int64 `toml:"split_size_mb" mapstructure:"split_size_mb" json:"split_size_mb"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -14,36 +14,47 @@ import (
|
|||||||
"github.com/krau/SaveAny-Bot/common/utils/ioutil"
|
"github.com/krau/SaveAny-Bot/common/utils/ioutil"
|
||||||
"github.com/krau/SaveAny-Bot/config"
|
"github.com/krau/SaveAny-Bot/config"
|
||||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||||
|
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||||
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
||||||
|
"github.com/krau/SaveAny-Bot/storage"
|
||||||
"golang.org/x/sync/errgroup"
|
"golang.org/x/sync/errgroup"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
type executionGroup struct {
|
||||||
|
elems []*TaskElement
|
||||||
|
batchSaver storage.StorageBatchSaver
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g executionGroup) usesBatchSaver() bool {
|
||||||
|
return g.batchSaver != nil
|
||||||
|
}
|
||||||
|
|
||||||
func (t *Task) Execute(ctx context.Context) error {
|
func (t *Task) Execute(ctx context.Context) error {
|
||||||
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("batch_file[%s]", t.ID))
|
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("batch_file[%s]", t.ID))
|
||||||
logger.Info("Starting batch file task")
|
logger.Info("Starting batch file task")
|
||||||
t.Progress.OnStart(ctx, t)
|
t.Progress.OnStart(ctx, t)
|
||||||
workers := config.C().Workers
|
groups := t.executionGroups()
|
||||||
eg, gctx := errgroup.WithContext(ctx)
|
var err error
|
||||||
eg.SetLimit(workers)
|
for i := 0; i < len(groups); {
|
||||||
for _, elem := range t.elems {
|
if groups[i].usesBatchSaver() {
|
||||||
eg.Go(func() error {
|
err = t.processBatch(ctx, groups[i])
|
||||||
t.processingMu.RLock()
|
i++
|
||||||
if t.processing[elem.ID] != nil {
|
} else {
|
||||||
return fmt.Errorf("element with ID %s is already being processed", elem.ID)
|
end := i + 1
|
||||||
|
for end < len(groups) && !groups[end].usesBatchSaver() {
|
||||||
|
end++
|
||||||
}
|
}
|
||||||
t.processingMu.RUnlock()
|
elems := make([]*TaskElement, 0, end-i)
|
||||||
t.processingMu.Lock()
|
for _, group := range groups[i:end] {
|
||||||
t.processing[elem.ID] = &elem
|
elems = append(elems, group.elems...)
|
||||||
t.processingMu.Unlock()
|
}
|
||||||
defer func() {
|
err = t.processElements(ctx, elems)
|
||||||
t.processingMu.Lock()
|
i = end
|
||||||
delete(t.processing, elem.ID)
|
}
|
||||||
t.processingMu.Unlock()
|
if err != nil {
|
||||||
}()
|
break
|
||||||
return t.processElement(gctx, elem)
|
}
|
||||||
})
|
|
||||||
}
|
}
|
||||||
err := eg.Wait()
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
logger.Errorf("Error during batch file processing: %v", err)
|
logger.Errorf("Error during batch file processing: %v", err)
|
||||||
} else {
|
} else {
|
||||||
@@ -53,6 +64,159 @@ func (t *Task) Execute(ctx context.Context) error {
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (t *Task) executionGroups() []executionGroup {
|
||||||
|
groups := make([]executionGroup, 0, len(t.elems))
|
||||||
|
for i := 0; i < len(t.elems); {
|
||||||
|
elem := &t.elems[i]
|
||||||
|
batchSaver, batchCapable := elem.Storage.(storage.StorageBatchSaver)
|
||||||
|
if !batchCapable || elem.sourceGroupKey == "" {
|
||||||
|
groups = append(groups, executionGroup{elems: []*TaskElement{elem}})
|
||||||
|
i++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
end := i + 1
|
||||||
|
for end < len(t.elems) {
|
||||||
|
next := &t.elems[end]
|
||||||
|
if next.Storage != elem.Storage || next.sourceGroupKey != elem.sourceGroupKey {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
end++
|
||||||
|
}
|
||||||
|
elems := make([]*TaskElement, 0, end-i)
|
||||||
|
for j := i; j < end; j++ {
|
||||||
|
elems = append(elems, &t.elems[j])
|
||||||
|
}
|
||||||
|
groups = append(groups, executionGroup{elems: elems, batchSaver: batchSaver})
|
||||||
|
i = end
|
||||||
|
}
|
||||||
|
return groups
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Task) processElements(ctx context.Context, elems []*TaskElement) error {
|
||||||
|
eg, gctx := errgroup.WithContext(ctx)
|
||||||
|
eg.SetLimit(config.C().Workers)
|
||||||
|
for _, elem := range elems {
|
||||||
|
eg.Go(func() error {
|
||||||
|
if err := t.markProcessing(elem); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer t.unmarkProcessing(elem.ID)
|
||||||
|
return t.processElement(gctx, *elem)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return eg.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
|
||||||
|
defer func() {
|
||||||
|
for _, elem := range group.elems {
|
||||||
|
if err := os.Remove(elem.localPath); err != nil && !os.IsNotExist(err) {
|
||||||
|
log.FromContext(ctx).Warnf("Failed to cleanup batch cache file %s: %v", elem.localPath, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
eg, gctx := errgroup.WithContext(ctx)
|
||||||
|
eg.SetLimit(config.C().Workers)
|
||||||
|
for _, elem := range group.elems {
|
||||||
|
eg.Go(func() error {
|
||||||
|
if err := t.markProcessing(elem); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer t.unmarkProcessing(elem.ID)
|
||||||
|
return t.downloadElement(gctx, elem)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
if err := eg.Wait(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
items := make([]storagetypes.BatchItem, 0, len(group.elems))
|
||||||
|
openFiles := make([]*os.File, 0, len(group.elems))
|
||||||
|
defer func() {
|
||||||
|
for _, file := range openFiles {
|
||||||
|
if err := file.Close(); err != nil {
|
||||||
|
log.FromContext(ctx).Warnf("Failed to close batch cache file %s: %v", file.Name(), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
for _, elem := range group.elems {
|
||||||
|
file, err := os.Open(elem.localPath)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to open cache file: %w", err)
|
||||||
|
}
|
||||||
|
stat, err := file.Stat()
|
||||||
|
if err != nil {
|
||||||
|
file.Close()
|
||||||
|
return fmt.Errorf("failed to get cache file stat: %w", err)
|
||||||
|
}
|
||||||
|
openFiles = append(openFiles, file)
|
||||||
|
items = append(items, storagetypes.BatchItem{
|
||||||
|
Reader: file,
|
||||||
|
StoragePath: elem.Path,
|
||||||
|
Size: stat.Size(),
|
||||||
|
SourceGroupKey: elem.sourceGroupKey,
|
||||||
|
Caption: elem.sourceCaption,
|
||||||
|
PreserveCaption: elem.preserveCaption,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
if err := group.batchSaver.SaveBatch(ctx, items); err != nil {
|
||||||
|
return fmt.Errorf("failed to save batch: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Task) markProcessing(elem *TaskElement) error {
|
||||||
|
t.processingMu.Lock()
|
||||||
|
defer t.processingMu.Unlock()
|
||||||
|
if t.processing[elem.ID] != nil {
|
||||||
|
return fmt.Errorf("element with ID %s is already being processed", elem.ID)
|
||||||
|
}
|
||||||
|
t.processing[elem.ID] = elem
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Task) unmarkProcessing(id string) {
|
||||||
|
t.processingMu.Lock()
|
||||||
|
delete(t.processing, id)
|
||||||
|
t.processingMu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Task) downloadElement(ctx context.Context, elem *TaskElement) error {
|
||||||
|
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", elem.File.Name()))
|
||||||
|
logger.Info("Starting file download")
|
||||||
|
localFile, err := fsutil.CreateFile(elem.localPath)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to create local file: %w", err)
|
||||||
|
}
|
||||||
|
wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) {
|
||||||
|
downloaded := t.downloaded.Add(int64(n))
|
||||||
|
t.Progress.OnProgress(ctx, t)
|
||||||
|
taskevent.Emit(ctx, taskevent.Event{
|
||||||
|
TaskID: t.ID,
|
||||||
|
Phase: taskevent.PhaseProgress,
|
||||||
|
TotalBytes: t.totalSize,
|
||||||
|
DownloadedBytes: downloaded,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
_, downloadErr := tdler.NewDownloader(elem.File).Parallel(ctx, wrAt)
|
||||||
|
closeErr := localFile.Close()
|
||||||
|
if downloadErr != nil {
|
||||||
|
return fmt.Errorf("failed to download file: %w", downloadErr)
|
||||||
|
}
|
||||||
|
if closeErr != nil {
|
||||||
|
return fmt.Errorf("failed to close cache file: %w", closeErr)
|
||||||
|
}
|
||||||
|
logger.Info("File downloaded successfully")
|
||||||
|
if path.Ext(elem.FileName()) == "" {
|
||||||
|
if ext := fsutil.DetectFileExt(elem.localPath); ext != "" {
|
||||||
|
elem.Path += ext
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||||
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", elem.File.Name()))
|
logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", elem.File.Name()))
|
||||||
if elem.stream {
|
if elem.stream {
|
||||||
|
|||||||
57
core/tasks/batchtfile/execute_group_test.go
Normal file
57
core/tasks/batchtfile/execute_group_test.go
Normal file
@@ -0,0 +1,57 @@
|
|||||||
|
package batchtfile
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/gotd/td/tg"
|
||||||
|
"github.com/krau/SaveAny-Bot/pkg/tfile"
|
||||||
|
tgstorage "github.com/krau/SaveAny-Bot/storage/telegram"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestExecutionGroupsPreserveSourceAlbums(t *testing.T) {
|
||||||
|
stor := new(tgstorage.Telegram)
|
||||||
|
otherStor := new(tgstorage.Telegram)
|
||||||
|
task := Task{elems: []TaskElement{
|
||||||
|
{Storage: stor, sourceGroupKey: "album-1"},
|
||||||
|
{Storage: stor, sourceGroupKey: "album-1"},
|
||||||
|
{Storage: stor},
|
||||||
|
{Storage: stor, sourceGroupKey: "album-2"},
|
||||||
|
{Storage: stor, sourceGroupKey: "album-2"},
|
||||||
|
{Storage: otherStor, sourceGroupKey: "album-2"},
|
||||||
|
}}
|
||||||
|
|
||||||
|
groups := task.executionGroups()
|
||||||
|
wantSizes := []int{2, 1, 2, 1}
|
||||||
|
wantBatch := []bool{true, false, true, true}
|
||||||
|
if len(groups) != len(wantSizes) {
|
||||||
|
t.Fatalf("got %d groups, want %d", len(groups), len(wantSizes))
|
||||||
|
}
|
||||||
|
for i := range groups {
|
||||||
|
if got := len(groups[i].elems); got != wantSizes[i] {
|
||||||
|
t.Errorf("group %d has %d elements, want %d", i, got, wantSizes[i])
|
||||||
|
}
|
||||||
|
if got := groups[i].usesBatchSaver(); got != wantBatch[i] {
|
||||||
|
t.Errorf("group %d batch=%v, want %v", i, got, wantBatch[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSourceMetadataPreservesAlbumIdentityAndCaption(t *testing.T) {
|
||||||
|
msg := &tg.Message{
|
||||||
|
PeerID: &tg.PeerChannel{ChannelID: 77},
|
||||||
|
Message: "original caption",
|
||||||
|
}
|
||||||
|
msg.SetGroupedID(42)
|
||||||
|
file := tfile.NewTGFile(nil, nil, 0, "photo.jpg", tfile.WithMessage(msg))
|
||||||
|
|
||||||
|
groupKey, caption, preserveCaption := sourceMetadata(file)
|
||||||
|
if groupKey != "*tg.PeerChannel:77:42" {
|
||||||
|
t.Fatalf("group key = %q, want %q", groupKey, "*tg.PeerChannel:77:42")
|
||||||
|
}
|
||||||
|
if caption != "original caption" {
|
||||||
|
t.Fatalf("caption = %q, want original caption", caption)
|
||||||
|
}
|
||||||
|
if !preserveCaption {
|
||||||
|
t.Fatal("preserveCaption = false, want true")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"sync"
|
"sync"
|
||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
|
|
||||||
|
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||||
"github.com/krau/SaveAny-Bot/config"
|
"github.com/krau/SaveAny-Bot/config"
|
||||||
"github.com/krau/SaveAny-Bot/core"
|
"github.com/krau/SaveAny-Bot/core"
|
||||||
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
|
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
|
||||||
@@ -18,12 +19,15 @@ import (
|
|||||||
var _ core.Executable = (*Task)(nil)
|
var _ core.Executable = (*Task)(nil)
|
||||||
|
|
||||||
type TaskElement struct {
|
type TaskElement struct {
|
||||||
ID string
|
ID string
|
||||||
Storage storage.Storage
|
Storage storage.Storage
|
||||||
Path string
|
Path string
|
||||||
File tfile.TGFile
|
File tfile.TGFile
|
||||||
localPath string
|
localPath string
|
||||||
stream bool
|
stream bool
|
||||||
|
sourceGroupKey string
|
||||||
|
sourceCaption string
|
||||||
|
preserveCaption bool
|
||||||
}
|
}
|
||||||
|
|
||||||
type Task struct {
|
type Task struct {
|
||||||
@@ -54,6 +58,7 @@ func NewTaskElement(
|
|||||||
file tfile.TGFile,
|
file tfile.TGFile,
|
||||||
) (*TaskElement, error) {
|
) (*TaskElement, error) {
|
||||||
id := xid.New().String()
|
id := xid.New().String()
|
||||||
|
groupKey, caption, preserveCaption := sourceMetadata(file)
|
||||||
_, ok := stor.(storage.StorageCannotStream)
|
_, ok := stor.(storage.StorageCannotStream)
|
||||||
if !config.C().Stream || ok {
|
if !config.C().Stream || ok {
|
||||||
cachePath, err := filepath.Abs(filepath.Join(config.C().Temp.BasePath, fmt.Sprintf("%s_%s", id, file.Name())))
|
cachePath, err := filepath.Abs(filepath.Join(config.C().Temp.BasePath, fmt.Sprintf("%s_%s", id, file.Name())))
|
||||||
@@ -61,22 +66,42 @@ func NewTaskElement(
|
|||||||
return nil, fmt.Errorf("failed to get absolute path for cache: %w", err)
|
return nil, fmt.Errorf("failed to get absolute path for cache: %w", err)
|
||||||
}
|
}
|
||||||
return &TaskElement{
|
return &TaskElement{
|
||||||
ID: id,
|
ID: id,
|
||||||
Storage: stor,
|
Storage: stor,
|
||||||
Path: path,
|
Path: path,
|
||||||
File: file,
|
File: file,
|
||||||
localPath: cachePath,
|
localPath: cachePath,
|
||||||
|
sourceGroupKey: groupKey,
|
||||||
|
sourceCaption: caption,
|
||||||
|
preserveCaption: preserveCaption,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
return &TaskElement{
|
return &TaskElement{
|
||||||
ID: id,
|
ID: id,
|
||||||
Storage: stor,
|
Storage: stor,
|
||||||
Path: path,
|
Path: path,
|
||||||
File: file,
|
File: file,
|
||||||
stream: true,
|
stream: true,
|
||||||
|
sourceGroupKey: groupKey,
|
||||||
|
sourceCaption: caption,
|
||||||
|
preserveCaption: preserveCaption,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func sourceMetadata(file tfile.TGFile) (groupKey, caption string, preserveCaption bool) {
|
||||||
|
messageFile, ok := file.(tfile.TGFileMessage)
|
||||||
|
if !ok || messageFile.Message() == nil {
|
||||||
|
return "", "", false
|
||||||
|
}
|
||||||
|
msg := messageFile.Message()
|
||||||
|
groupID, grouped := msg.GetGroupedID()
|
||||||
|
if !grouped || groupID == 0 {
|
||||||
|
return "", "", false
|
||||||
|
}
|
||||||
|
chatID := tgutil.ChatIdFromPeer(msg.GetPeerID())
|
||||||
|
return fmt.Sprintf("%T:%d:%d", msg.GetPeerID(), chatID, groupID), msg.GetMessage(), true
|
||||||
|
}
|
||||||
|
|
||||||
func NewBatchTGFileTask(
|
func NewBatchTGFileTask(
|
||||||
id string,
|
id string,
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
|
|||||||
42
core/tasks/tfile/caption_test.go
Normal file
42
core/tasks/tfile/caption_test.go
Normal file
@@ -0,0 +1,42 @@
|
|||||||
|
package tfile
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/gotd/td/tg"
|
||||||
|
tfilepkg "github.com/krau/SaveAny-Bot/pkg/tfile"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSourceCaption(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
file tfilepkg.TGFile
|
||||||
|
want string
|
||||||
|
ok bool
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "original caption",
|
||||||
|
file: tfilepkg.NewTGFile(nil, nil, 0, "video.mov", tfilepkg.WithMessage(&tg.Message{Message: "original caption"})),
|
||||||
|
want: "original caption",
|
||||||
|
ok: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty caption suppresses storage fallback",
|
||||||
|
file: tfilepkg.NewTGFile(nil, nil, 0, "video.mov", tfilepkg.WithMessage(&tg.Message{})),
|
||||||
|
ok: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "file without source message",
|
||||||
|
file: tfilepkg.NewTGFile(nil, nil, 0, "video.mov"),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
got, ok := sourceCaption(tt.file)
|
||||||
|
if ok != tt.ok || got != tt.want {
|
||||||
|
t.Fatalf("sourceCaption() = (%q, %v), want (%q, %v)", got, ok, tt.want, tt.ok)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -12,6 +12,8 @@ import (
|
|||||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||||
"github.com/krau/SaveAny-Bot/config"
|
"github.com/krau/SaveAny-Bot/config"
|
||||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||||
|
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||||
|
tfilepkg "github.com/krau/SaveAny-Bot/pkg/tfile"
|
||||||
)
|
)
|
||||||
|
|
||||||
func (t *Task) Execute(ctx context.Context) error {
|
func (t *Task) Execute(ctx context.Context) error {
|
||||||
@@ -57,6 +59,9 @@ func (t *Task) Execute(ctx context.Context) error {
|
|||||||
return fmt.Errorf("failed to get file stat: %w", err)
|
return fmt.Errorf("failed to get file stat: %w", err)
|
||||||
}
|
}
|
||||||
vctx := context.WithValue(ctx, ctxkey.ContentLength, fileStat.Size())
|
vctx := context.WithValue(ctx, ctxkey.ContentLength, fileStat.Size())
|
||||||
|
if caption, ok := sourceCaption(t.File); ok {
|
||||||
|
vctx = storagetypes.WithSourceCaption(vctx, caption)
|
||||||
|
}
|
||||||
err = retry.Retry(func() error {
|
err = retry.Retry(func() error {
|
||||||
file, err := os.Open(t.localPath)
|
file, err := os.Open(t.localPath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -73,3 +78,11 @@ func (t *Task) Execute(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func sourceCaption(file tfilepkg.TGFile) (string, bool) {
|
||||||
|
messageFile, ok := file.(tfilepkg.TGFileMessage)
|
||||||
|
if !ok || messageFile.Message() == nil {
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
return messageFile.Message().GetMessage(), true
|
||||||
|
}
|
||||||
|
|||||||
@@ -79,7 +79,8 @@ Stream mode is not supported.
|
|||||||
chat_id = "123456789" # Telegram chat ID, the bot will send files to this chat
|
chat_id = "123456789" # Telegram chat ID, the bot will send files to this chat
|
||||||
force_file = false # Force sending as file, default is false
|
force_file = false # Force sending as file, default is false
|
||||||
skip_large = false # Skip large files, default is false. If enabled, files exceeding Telegram's limit will not be uploaded.
|
skip_large = false # Skip large files, default is false. If enabled, files exceeding Telegram's limit will not be uploaded.
|
||||||
spilt_size_mb = 2000 # Split size in MB, default is 2000 MB (2 GB). Files larger than this will be split into multiple parts (zip format). Ignored when skip_large is true.
|
split_large_video = false # Losslessly split oversized videos into one album of up to 10 independently playable parts. Falls back to ZIP parts on failure or when more than 10 parts are required.
|
||||||
|
split_size_mb = 0 # Split size in MB. 0 uses the uploader account limit: 2000 MB for bots/regular users and 4000 MB for Premium users. Oversized non-video files use ZIP parts. Ignored when skip_large is true.
|
||||||
```
|
```
|
||||||
|
|
||||||
## Rclone
|
## Rclone
|
||||||
|
|||||||
@@ -82,10 +82,12 @@ chat_id = "123456789"
|
|||||||
force_file = false
|
force_file = false
|
||||||
# 是否跳过大文件, 默认为 false. 如果启用, 超过 Telegram 限制的文件将不会上传.
|
# 是否跳过大文件, 默认为 false. 如果启用, 超过 Telegram 限制的文件将不会上传.
|
||||||
skip_large = false
|
skip_large = false
|
||||||
# 分卷大小, 单位 MB, 默认为 2000 MB (2 GB).
|
# 超限视频是否使用 FFmpeg 无损分割成一个媒体组(最多 10 个)内可独立播放的小视频;失败或需要超过 10 段时回退到 ZIP 分卷.
|
||||||
# 超过该大小的文件将被分割成多个部分上传.(使用 zip 格式)
|
split_large_video = false
|
||||||
|
# 分卷大小, 单位 MB. 设为 0 时使用实际上传账号的限制:Bot/普通用户为 2000 MB,Premium 用户为 4000 MB.
|
||||||
|
# 超过该大小的文件将被分割成多个部分上传;非视频使用 ZIP 格式.
|
||||||
# 当 skip_large 启用时, 该选项无效.
|
# 当 skip_large 启用时, 该选项无效.
|
||||||
spilt_size_mb = 2000
|
split_size_mb = 0
|
||||||
```
|
```
|
||||||
|
|
||||||
## Rclone
|
## Rclone
|
||||||
|
|||||||
17
pkg/storagetypes/batch.go
Normal file
17
pkg/storagetypes/batch.go
Normal file
@@ -0,0 +1,17 @@
|
|||||||
|
package storagetypes
|
||||||
|
|
||||||
|
import "io"
|
||||||
|
|
||||||
|
// BatchItem describes one seekable file in a logical batch storage operation.
|
||||||
|
type BatchItem struct {
|
||||||
|
Reader io.ReadSeeker
|
||||||
|
StoragePath string
|
||||||
|
Size int64
|
||||||
|
|
||||||
|
// SourceGroupKey is empty for standalone source messages.
|
||||||
|
SourceGroupKey string
|
||||||
|
Caption string
|
||||||
|
// PreserveCaption distinguishes an intentionally empty source caption from
|
||||||
|
// the storage backend's default caption.
|
||||||
|
PreserveCaption bool
|
||||||
|
}
|
||||||
19
pkg/storagetypes/context.go
Normal file
19
pkg/storagetypes/context.go
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
package storagetypes
|
||||||
|
|
||||||
|
import "context"
|
||||||
|
|
||||||
|
type sourceCaptionContextKey struct{}
|
||||||
|
|
||||||
|
// WithSourceCaption records the source Telegram message caption for storage
|
||||||
|
// backends that can preserve it. Calling this function with an empty caption
|
||||||
|
// intentionally suppresses a backend-generated fallback caption.
|
||||||
|
func WithSourceCaption(ctx context.Context, caption string) context.Context {
|
||||||
|
return context.WithValue(ctx, sourceCaptionContextKey{}, caption)
|
||||||
|
}
|
||||||
|
|
||||||
|
// SourceCaptionFromContext returns the source caption and whether the caller
|
||||||
|
// explicitly supplied one.
|
||||||
|
func SourceCaptionFromContext(ctx context.Context) (string, bool) {
|
||||||
|
caption, ok := ctx.Value(sourceCaptionContextKey{}).(string)
|
||||||
|
return caption, ok
|
||||||
|
}
|
||||||
24
pkg/storagetypes/context_test.go
Normal file
24
pkg/storagetypes/context_test.go
Normal file
@@ -0,0 +1,24 @@
|
|||||||
|
package storagetypes
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSourceCaptionContext(t *testing.T) {
|
||||||
|
if _, ok := SourceCaptionFromContext(context.Background()); ok {
|
||||||
|
t.Fatal("caption unexpectedly present on empty context")
|
||||||
|
}
|
||||||
|
|
||||||
|
tests := []string{"original caption", ""}
|
||||||
|
for _, want := range tests {
|
||||||
|
ctx := WithSourceCaption(context.Background(), want)
|
||||||
|
got, ok := SourceCaptionFromContext(ctx)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("caption %q was not marked as present", want)
|
||||||
|
}
|
||||||
|
if got != want {
|
||||||
|
t.Fatalf("caption = %q, want %q", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -31,6 +31,13 @@ type StorageCannotStream interface {
|
|||||||
CannotStream() string
|
CannotStream() string
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// StorageBatchSaver can preserve relationships between files when saving a
|
||||||
|
// logical batch, such as a Telegram media album.
|
||||||
|
type StorageBatchSaver interface {
|
||||||
|
Storage
|
||||||
|
SaveBatch(ctx context.Context, items []storagetypes.BatchItem) error
|
||||||
|
}
|
||||||
|
|
||||||
// StorageListable 表示支持列举目录内容的存储
|
// StorageListable 表示支持列举目录内容的存储
|
||||||
type StorageListable interface {
|
type StorageListable interface {
|
||||||
Storage
|
Storage
|
||||||
|
|||||||
140
storage/telegram/media_group_test.go
Normal file
140
storage/telegram/media_group_test.go
Normal file
@@ -0,0 +1,140 @@
|
|||||||
|
package telegram
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"io"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestPlanMediaGroups(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
items []batchMediaItem
|
||||||
|
wantSizes []int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "same source album",
|
||||||
|
items: []batchMediaItem{
|
||||||
|
albumItem("a", 1, true),
|
||||||
|
albumItem("a", 1, true),
|
||||||
|
},
|
||||||
|
wantSizes: []int{2},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "different source albums",
|
||||||
|
items: []batchMediaItem{
|
||||||
|
albumItem("a", 1, true),
|
||||||
|
albumItem("b", 1, true),
|
||||||
|
},
|
||||||
|
wantSizes: []int{1, 1},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "ungrouped messages",
|
||||||
|
items: []batchMediaItem{
|
||||||
|
albumItem("", 1, true),
|
||||||
|
albumItem("", 1, true),
|
||||||
|
},
|
||||||
|
wantSizes: []int{1, 1},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "different target chats",
|
||||||
|
items: []batchMediaItem{
|
||||||
|
albumItem("a", 1, true),
|
||||||
|
albumItem("a", 2, true),
|
||||||
|
},
|
||||||
|
wantSizes: []int{1, 1},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "ineligible media does not bridge albums",
|
||||||
|
items: []batchMediaItem{
|
||||||
|
albumItem("a", 1, true),
|
||||||
|
albumItem("a", 1, false),
|
||||||
|
albumItem("a", 1, true),
|
||||||
|
},
|
||||||
|
wantSizes: []int{1, 1, 1},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "maximum album size",
|
||||||
|
items: repeatedAlbumItems(11),
|
||||||
|
wantSizes: []int{10, 1},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
groups := planMediaGroups(tt.items)
|
||||||
|
if len(groups) != len(tt.wantSizes) {
|
||||||
|
t.Fatalf("got %d groups, want %d", len(groups), len(tt.wantSizes))
|
||||||
|
}
|
||||||
|
for i, want := range tt.wantSizes {
|
||||||
|
if got := len(groups[i]); got != want {
|
||||||
|
t.Errorf("group %d has %d items, want %d", i, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMediaCaption(t *testing.T) {
|
||||||
|
empty := ""
|
||||||
|
original := "original caption"
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
override *string
|
||||||
|
wantLen int
|
||||||
|
}{
|
||||||
|
{name: "filename fallback", wantLen: 1},
|
||||||
|
{name: "preserve empty source caption", override: &empty, wantLen: 0},
|
||||||
|
{name: "preserve source caption", override: &original, wantLen: 1},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
if got := len(mediaCaption("file.jpg", tt.override)); got != tt.wantLen {
|
||||||
|
t.Fatalf("got %d caption options, want %d", got, tt.wantLen)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInspectBatchItemRewindsBeforeMimetypeDetection(t *testing.T) {
|
||||||
|
data := []byte("\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR")
|
||||||
|
reader := bytes.NewReader(data)
|
||||||
|
if _, err := reader.Seek(4, io.SeekStart); err != nil {
|
||||||
|
t.Fatalf("failed to set initial reader offset: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
mediaItem, err := new(Telegram).inspectBatchItem(nil, storagetypes.BatchItem{
|
||||||
|
Reader: reader,
|
||||||
|
StoragePath: "photo.png",
|
||||||
|
Size: int64(len(data)),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("inspectBatchItem returned an error: %v", err)
|
||||||
|
}
|
||||||
|
if !mediaItem.albumEligible {
|
||||||
|
t.Fatal("albumEligible = false, want true for PNG input")
|
||||||
|
}
|
||||||
|
if offset, err := reader.Seek(0, io.SeekCurrent); err != nil {
|
||||||
|
t.Fatalf("failed to get final reader offset: %v", err)
|
||||||
|
} else if offset != 0 {
|
||||||
|
t.Fatalf("reader offset = %d, want 0", offset)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func albumItem(group string, chatID int64, eligible bool) batchMediaItem {
|
||||||
|
return batchMediaItem{
|
||||||
|
item: storagetypes.BatchItem{SourceGroupKey: group},
|
||||||
|
chatID: chatID,
|
||||||
|
albumEligible: eligible,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func repeatedAlbumItems(count int) []batchMediaItem {
|
||||||
|
items := make([]batchMediaItem, count)
|
||||||
|
for i := range items {
|
||||||
|
items[i] = albumItem("a", 1, true)
|
||||||
|
}
|
||||||
|
return items
|
||||||
|
}
|
||||||
@@ -12,6 +12,7 @@ import (
|
|||||||
|
|
||||||
"github.com/celestix/gotgproto/ext"
|
"github.com/celestix/gotgproto/ext"
|
||||||
"github.com/charmbracelet/log"
|
"github.com/charmbracelet/log"
|
||||||
|
"github.com/duke-git/lancet/v2/retry"
|
||||||
"github.com/duke-git/lancet/v2/slice"
|
"github.com/duke-git/lancet/v2/slice"
|
||||||
"github.com/duke-git/lancet/v2/validator"
|
"github.com/duke-git/lancet/v2/validator"
|
||||||
"github.com/gabriel-vasile/mimetype"
|
"github.com/gabriel-vasile/mimetype"
|
||||||
@@ -26,14 +27,16 @@ import (
|
|||||||
"github.com/krau/SaveAny-Bot/pkg/consts/tglimit"
|
"github.com/krau/SaveAny-Bot/pkg/consts/tglimit"
|
||||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||||
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
|
||||||
|
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||||
"github.com/rs/xid"
|
"github.com/rs/xid"
|
||||||
"golang.org/x/time/rate"
|
"golang.org/x/time/rate"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
// https://core.telegram.org/api/config#upload-max-fileparts-default
|
// https://core.telegram.org/api/config#upload-max-fileparts-default
|
||||||
DefaultSplitSize = 4000 * 524288 // 4000 * 512 KB
|
DefaultSplitSize = 4000 * 524288 // 4000 * 512 KB
|
||||||
MaxUploadFileSize = 4000 * 524288 // 4000 * 512 KB
|
MaxUploadFileSize = 4000 * 524288 // 4000 * 512 KB
|
||||||
|
PremiumMaxUploadFileSize = 8000 * 524288 // 8000 * 512 KB
|
||||||
)
|
)
|
||||||
|
|
||||||
type Telegram struct {
|
type Telegram struct {
|
||||||
@@ -41,6 +44,19 @@ type Telegram struct {
|
|||||||
limiter *rate.Limiter
|
limiter *rate.Limiter
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type preparedMedia struct {
|
||||||
|
peer tg.InputPeerClass
|
||||||
|
uploader *uploader.Uploader
|
||||||
|
media message.MultiMediaOption
|
||||||
|
}
|
||||||
|
|
||||||
|
type batchMediaItem struct {
|
||||||
|
item storagetypes.BatchItem
|
||||||
|
chatID int64
|
||||||
|
albumEligible bool
|
||||||
|
useSingleSave bool
|
||||||
|
}
|
||||||
|
|
||||||
func (t *Telegram) Init(ctx context.Context, cfg storconfig.StorageConfig) error {
|
func (t *Telegram) Init(ctx context.Context, cfg storconfig.StorageConfig) error {
|
||||||
telegramConfig, ok := cfg.(*storconfig.TelegramStorageConfig)
|
telegramConfig, ok := cfg.(*storconfig.TelegramStorageConfig)
|
||||||
if !ok {
|
if !ok {
|
||||||
@@ -72,36 +88,130 @@ func (t *Telegram) Exists(ctx context.Context, storagePath string) bool {
|
|||||||
|
|
||||||
func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) error {
|
func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) error {
|
||||||
storagePath = path.Clean(storagePath)
|
storagePath = path.Clean(storagePath)
|
||||||
|
captionOverride := sourceCaptionOverride(ctx)
|
||||||
tctx := tgutil.ExtFromContext(ctx)
|
tctx := tgutil.ExtFromContext(ctx)
|
||||||
if tctx == nil {
|
if tctx == nil {
|
||||||
return fmt.Errorf("failed to get telegram context")
|
return fmt.Errorf("failed to get telegram context")
|
||||||
}
|
}
|
||||||
size := func() int64 {
|
size := contentLength(ctx)
|
||||||
if length := ctx.Value(ctxkey.ContentLength); length != nil {
|
maxUploadSize := maxUploadFileSize(tctx)
|
||||||
if l, ok := length.(int64); ok {
|
if t.config.SkipLarge && size > maxUploadSize {
|
||||||
return l
|
log.FromContext(ctx).Warnf("Skipping file larger than Telegram limit (%d bytes): %d bytes", maxUploadSize, size)
|
||||||
}
|
|
||||||
}
|
|
||||||
return -1 // unknown size
|
|
||||||
}()
|
|
||||||
if t.config.SkipLarge && size > MaxUploadFileSize {
|
|
||||||
log.FromContext(ctx).Warnf("Skipping file larger than Telegram limit (%d bytes): %d bytes", MaxUploadFileSize, size)
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
rs, seekable := r.(io.ReadSeeker)
|
splitSize := t.splitSize(maxUploadSize)
|
||||||
splitSize := t.config.SplitSizeMB * 1024 * 1024
|
if size > splitSize {
|
||||||
if splitSize <= 0 {
|
filename, chatID := t.target(tctx, storagePath)
|
||||||
splitSize = DefaultSplitSize
|
if filename == "" {
|
||||||
|
if rs, ok := r.(io.ReadSeeker); ok {
|
||||||
|
mtype, err := mimetype.DetectReader(rs)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to detect mimetype: %w", err)
|
||||||
|
}
|
||||||
|
filename = xid.New().String() + mtype.Extension()
|
||||||
|
if _, err := rs.Seek(0, io.SeekStart); err != nil {
|
||||||
|
return fmt.Errorf("failed to seek reader: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
upler := t.newUploader(tctx, size)
|
||||||
|
peer := tryGetInputPeer(tctx, chatID)
|
||||||
|
if peer == nil || peer.Zero() {
|
||||||
|
return fmt.Errorf("failed to get input peer for chat ID %d", chatID)
|
||||||
|
}
|
||||||
|
if err := t.limiter.Wait(ctx); err != nil {
|
||||||
|
return fmt.Errorf("rate limit failed: %w", err)
|
||||||
|
}
|
||||||
|
if t.config.SplitLargeVideo {
|
||||||
|
rs, ok := r.(io.ReadSeeker)
|
||||||
|
if ok {
|
||||||
|
mtype, detectErr := mimetype.DetectReader(rs)
|
||||||
|
if _, seekErr := rs.Seek(0, io.SeekStart); seekErr != nil {
|
||||||
|
return fmt.Errorf("failed to seek large file after mimetype detection: %w", seekErr)
|
||||||
|
}
|
||||||
|
if detectErr != nil {
|
||||||
|
log.FromContext(ctx).Warnf("Failed to detect large file type, falling back to ZIP split: %s", detectErr)
|
||||||
|
} else if strings.HasPrefix(mtype.String(), "video/") {
|
||||||
|
parts, cleanup, splitErr := createLosslessVideoParts(
|
||||||
|
ctx,
|
||||||
|
rs,
|
||||||
|
filename,
|
||||||
|
size,
|
||||||
|
splitSize,
|
||||||
|
)
|
||||||
|
if splitErr != nil {
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
return ctx.Err()
|
||||||
|
}
|
||||||
|
log.FromContext(ctx).Warnf("Lossless video split failed, falling back to ZIP split: %s", splitErr)
|
||||||
|
} else {
|
||||||
|
defer cleanup()
|
||||||
|
log.FromContext(ctx).Infof("Uploading oversized video as %d lossless-playable parts", len(parts))
|
||||||
|
for _, part := range parts {
|
||||||
|
log.FromContext(ctx).Infof("Prepared lossless video part %s (%d bytes)", part.Name, part.Size)
|
||||||
|
}
|
||||||
|
return t.uploadLosslessVideoParts(ctx, tctx, storagePath, parts, captionOverride)
|
||||||
|
}
|
||||||
|
if _, seekErr := rs.Seek(0, io.SeekStart); seekErr != nil {
|
||||||
|
return fmt.Errorf("failed to seek large video before ZIP fallback: %w", seekErr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return t.splitUpload(tctx, r, filename, upler, peer, size, splitSize)
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := t.limiter.Wait(ctx); err != nil {
|
if err := t.limiter.Wait(ctx); err != nil {
|
||||||
return fmt.Errorf("rate limit failed: %w", err)
|
return fmt.Errorf("rate limit failed: %w", err)
|
||||||
}
|
}
|
||||||
|
prepared, err := t.prepareMedia(ctx, tctx, r, storagePath, size, nil)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
_, err = tctx.Sender.
|
||||||
|
WithUploader(prepared.uploader).
|
||||||
|
To(prepared.peer).
|
||||||
|
Media(ctx, prepared.media)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func contentLength(ctx context.Context) int64 {
|
||||||
|
if length := ctx.Value(ctxkey.ContentLength); length != nil {
|
||||||
|
if size, ok := length.(int64); ok {
|
||||||
|
return size
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return -1
|
||||||
|
}
|
||||||
|
|
||||||
|
func sourceCaptionOverride(ctx context.Context) *string {
|
||||||
|
caption, ok := storagetypes.SourceCaptionFromContext(ctx)
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return &caption
|
||||||
|
}
|
||||||
|
|
||||||
|
func maxUploadFileSize(tctx *ext.Context) int64 {
|
||||||
|
if tctx != nil && tctx.Self != nil && !tctx.Self.GetBot() && tctx.Self.GetPremium() {
|
||||||
|
return PremiumMaxUploadFileSize
|
||||||
|
}
|
||||||
|
return MaxUploadFileSize
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Telegram) splitSize(maxUploadSize int64) int64 {
|
||||||
|
splitSize := t.config.SplitSizeMB * 1024 * 1024
|
||||||
|
if splitSize <= 0 {
|
||||||
|
return maxUploadSize
|
||||||
|
}
|
||||||
|
return min(splitSize, maxUploadSize)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Telegram) target(tctx *ext.Context, storagePath string) (string, int64) {
|
||||||
// 去除前导斜杠并分隔路径, 当 len(parts):
|
// 去除前导斜杠并分隔路径, 当 len(parts):
|
||||||
// ==0, 存储到配置文件中的 chat_id, 随机文件名
|
// ==0, 存储到配置文件中的 chat_id, 随机文件名
|
||||||
// ==1, 视作只有文件名, 存储到配置文件中的 chat_id
|
// ==1, 视作只有文件名, 存储到配置文件中的 chat_id
|
||||||
// ==2, parts[0]: 视作要存储到的 chat_id, parts[1]: filename
|
// >=2, parts[0]: 视作要存储到的 chat_id, 最后一项为 filename
|
||||||
parts := slice.Compact(strings.Split(strings.TrimPrefix(storagePath, "/"), "/"))
|
parts := slice.Compact(strings.Split(strings.TrimPrefix(storagePath, "/"), "/"))
|
||||||
filename := ""
|
filename := ""
|
||||||
chatID := t.config.ChatID
|
chatID := t.config.ChatID
|
||||||
@@ -111,38 +221,54 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er
|
|||||||
if len(parts) >= 2 && validator.IsAlphaNumeric(parts[0]) {
|
if len(parts) >= 2 && validator.IsAlphaNumeric(parts[0]) {
|
||||||
cid, err := tgutil.ParseChatID(tctx, parts[0])
|
cid, err := tgutil.ParseChatID(tctx, parts[0])
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// id不合法时使用配置文件中的 chat_id
|
log.FromContext(tctx).Warnf("Failed to parse chat ID from path, using configured chat_id: %s", err)
|
||||||
log.FromContext(ctx).Warnf("Failed to parse chat ID from path, using configured chat_id: %s", err)
|
|
||||||
cid = chatID
|
cid = chatID
|
||||||
}
|
}
|
||||||
chatID = cid
|
chatID = cid
|
||||||
}
|
}
|
||||||
upler := uploader.NewUploader(tctx.Raw).
|
return filename, chatID
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Telegram) newUploader(tctx *ext.Context, size int64) *uploader.Uploader {
|
||||||
|
return uploader.NewUploader(tctx.Raw).
|
||||||
WithPartSize(tglimit.MaxUploadPartSize).
|
WithPartSize(tglimit.MaxUploadPartSize).
|
||||||
WithThreads(dlutil.BestThreads(size, config.C().Threads))
|
WithThreads(dlutil.BestThreads(size, config.C().Threads))
|
||||||
|
}
|
||||||
|
|
||||||
|
func mediaCaption(filename string, override *string) []message.StyledTextOption {
|
||||||
|
if override == nil {
|
||||||
|
return []message.StyledTextOption{styling.Plain(filename)}
|
||||||
|
}
|
||||||
|
if *override == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return []message.StyledTextOption{styling.Plain(*override)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Telegram) prepareMedia(ctx context.Context, tctx *ext.Context, r io.Reader, storagePath string, size int64, captionOverride *string) (*preparedMedia, error) {
|
||||||
|
storagePath = path.Clean(storagePath)
|
||||||
|
filename, chatID := t.target(tctx, storagePath)
|
||||||
|
upler := t.newUploader(tctx, size)
|
||||||
peer := tryGetInputPeer(tctx, chatID)
|
peer := tryGetInputPeer(tctx, chatID)
|
||||||
if peer == nil || peer.Zero() {
|
if peer == nil || peer.Zero() {
|
||||||
return fmt.Errorf("failed to get input peer for chat ID %d", chatID)
|
return nil, fmt.Errorf("failed to get input peer for chat ID %d", chatID)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
rs, seekable := r.(io.ReadSeeker)
|
||||||
var mtype *mimetype.MIME
|
var mtype *mimetype.MIME
|
||||||
if seekable {
|
if seekable {
|
||||||
var err error
|
var err error
|
||||||
mtype, err = mimetype.DetectReader(rs)
|
mtype, err = mimetype.DetectReader(rs)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("failed to detect mimetype: %w", err)
|
return nil, fmt.Errorf("failed to detect mimetype: %w", err)
|
||||||
}
|
}
|
||||||
if filename == "" {
|
if filename == "" {
|
||||||
filename = xid.New().String() + mtype.Extension()
|
filename = xid.New().String() + mtype.Extension()
|
||||||
}
|
}
|
||||||
|
|
||||||
if _, err := rs.Seek(0, io.SeekStart); err != nil {
|
if _, err := rs.Seek(0, io.SeekStart); err != nil {
|
||||||
return fmt.Errorf("failed to seek reader: %w", err)
|
return nil, fmt.Errorf("failed to seek reader: %w", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if size > splitSize {
|
|
||||||
// large file, use split uploader
|
|
||||||
return t.splitUpload(tctx, r, filename, upler, peer, size, splitSize)
|
|
||||||
}
|
|
||||||
|
|
||||||
var file tg.InputFileClass
|
var file tg.InputFileClass
|
||||||
var err error
|
var err error
|
||||||
@@ -152,21 +278,20 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er
|
|||||||
file, err = upler.Upload(ctx, uploader.NewUpload(filename, r, size))
|
file, err = upler.Upload(ctx, uploader.NewUpload(filename, r, size))
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("failed to upload file to telegram: %w", err)
|
return nil, fmt.Errorf("failed to upload file to telegram: %w", err)
|
||||||
}
|
}
|
||||||
caption := styling.Plain(filename)
|
caption := mediaCaption(filename, captionOverride)
|
||||||
forceFile := t.config.ForceFile
|
forceFile := t.config.ForceFile
|
||||||
|
|
||||||
if mtype != nil && strings.HasPrefix(mtype.String(), "image/") && size >= tglimit.MaxPhotoSize {
|
if mtype != nil && strings.HasPrefix(mtype.String(), "image/") && size >= tglimit.MaxPhotoSize {
|
||||||
forceFile = true
|
forceFile = true
|
||||||
}
|
}
|
||||||
doc := message.UploadedDocument(file, caption).
|
doc := message.UploadedDocument(file, caption...).
|
||||||
Filename(filename).
|
Filename(filename).
|
||||||
ForceFile(forceFile)
|
ForceFile(forceFile)
|
||||||
if mtype != nil {
|
if mtype != nil {
|
||||||
doc = doc.MIME(mtype.String())
|
doc = doc.MIME(mtype.String())
|
||||||
}
|
}
|
||||||
var media message.MediaOption = doc
|
var media message.MultiMediaOption = doc
|
||||||
if mtype != nil && rs != nil {
|
if mtype != nil && rs != nil {
|
||||||
switch mtypeStr := mtype.String(); {
|
switch mtypeStr := mtype.String(); {
|
||||||
case strings.HasPrefix(mtypeStr, "video/"):
|
case strings.HasPrefix(mtypeStr, "video/"):
|
||||||
@@ -205,12 +330,136 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er
|
|||||||
case strings.HasPrefix(mtypeStr, "audio/"):
|
case strings.HasPrefix(mtypeStr, "audio/"):
|
||||||
media = doc.Audio().Title(filename)
|
media = doc.Audio().Title(filename)
|
||||||
case strings.HasPrefix(mtypeStr, "image/") && !strings.HasSuffix(mtypeStr, "webp"):
|
case strings.HasPrefix(mtypeStr, "image/") && !strings.HasSuffix(mtypeStr, "webp"):
|
||||||
media = message.UploadedPhoto(file, caption)
|
media = message.UploadedPhoto(file, caption...)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
sender := tctx.Sender
|
return &preparedMedia{
|
||||||
_, err = sender.WithUploader(upler).To(peer).Media(ctx, media)
|
peer: peer,
|
||||||
return err
|
uploader: upler,
|
||||||
|
media: media,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// SaveBatch preserves each source photo/video group as a Telegram album.
|
||||||
|
func (t *Telegram) SaveBatch(ctx context.Context, items []storagetypes.BatchItem) error {
|
||||||
|
tctx := tgutil.ExtFromContext(ctx)
|
||||||
|
if tctx == nil {
|
||||||
|
return fmt.Errorf("failed to get telegram context")
|
||||||
|
}
|
||||||
|
|
||||||
|
inspected := make([]batchMediaItem, 0, len(items))
|
||||||
|
for _, item := range items {
|
||||||
|
mediaItem, err := t.inspectBatchItem(tctx, item)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
inspected = append(inspected, mediaItem)
|
||||||
|
}
|
||||||
|
for _, group := range planMediaGroups(inspected) {
|
||||||
|
if err := t.saveMediaGroup(ctx, tctx, group); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Telegram) inspectBatchItem(tctx *ext.Context, item storagetypes.BatchItem) (batchMediaItem, error) {
|
||||||
|
_, chatID := t.target(tctx, path.Clean(item.StoragePath))
|
||||||
|
result := batchMediaItem{item: item, chatID: chatID}
|
||||||
|
maxUploadSize := maxUploadFileSize(tctx)
|
||||||
|
if (t.config.SkipLarge && item.Size > maxUploadSize) ||
|
||||||
|
item.Size > t.splitSize(maxUploadSize) {
|
||||||
|
result.useSingleSave = true
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
if _, err := item.Reader.Seek(0, io.SeekStart); err != nil {
|
||||||
|
return result, fmt.Errorf("failed to seek batch item before mimetype detection: %w", err)
|
||||||
|
}
|
||||||
|
mtype, err := mimetype.DetectReader(item.Reader)
|
||||||
|
if err != nil {
|
||||||
|
return result, fmt.Errorf("failed to detect batch item mimetype: %w", err)
|
||||||
|
}
|
||||||
|
if _, err := item.Reader.Seek(0, io.SeekStart); err != nil {
|
||||||
|
return result, fmt.Errorf("failed to seek batch item: %w", err)
|
||||||
|
}
|
||||||
|
mtypeStr := mtype.String()
|
||||||
|
forceFile := t.config.ForceFile || strings.HasPrefix(mtypeStr, "image/") && item.Size >= tglimit.MaxPhotoSize
|
||||||
|
result.albumEligible = !forceFile && (strings.HasPrefix(mtypeStr, "video/") ||
|
||||||
|
strings.HasPrefix(mtypeStr, "image/") && mtypeStr != "image/webp" && mtypeStr != "image/gif")
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func planMediaGroups(items []batchMediaItem) [][]batchMediaItem {
|
||||||
|
groups := make([][]batchMediaItem, 0, len(items))
|
||||||
|
for i := 0; i < len(items); {
|
||||||
|
item := items[i]
|
||||||
|
if item.useSingleSave || !item.albumEligible || item.item.SourceGroupKey == "" {
|
||||||
|
groups = append(groups, items[i:i+1])
|
||||||
|
i++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
end := i + 1
|
||||||
|
for end < len(items) && end-i < 10 {
|
||||||
|
next := items[end]
|
||||||
|
if next.useSingleSave || !next.albumEligible || next.chatID != item.chatID || next.item.SourceGroupKey != item.item.SourceGroupKey {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
end++
|
||||||
|
}
|
||||||
|
groups = append(groups, items[i:end])
|
||||||
|
i = end
|
||||||
|
}
|
||||||
|
return groups
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Telegram) saveMediaGroup(ctx context.Context, tctx *ext.Context, group []batchMediaItem) error {
|
||||||
|
return retry.Retry(func() error {
|
||||||
|
if len(group) == 1 && group[0].useSingleSave {
|
||||||
|
item := group[0].item
|
||||||
|
if _, err := item.Reader.Seek(0, io.SeekStart); err != nil {
|
||||||
|
return fmt.Errorf("failed to seek batch item: %w", err)
|
||||||
|
}
|
||||||
|
itemCtx := context.WithValue(ctx, ctxkey.ContentLength, item.Size)
|
||||||
|
if item.PreserveCaption {
|
||||||
|
itemCtx = storagetypes.WithSourceCaption(itemCtx, item.Caption)
|
||||||
|
}
|
||||||
|
return t.Save(itemCtx, item.Reader, item.StoragePath)
|
||||||
|
}
|
||||||
|
if err := t.limiter.Wait(ctx); err != nil {
|
||||||
|
return fmt.Errorf("rate limit failed: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
prepared := make([]preparedMedia, 0, len(group))
|
||||||
|
for _, mediaItem := range group {
|
||||||
|
item := mediaItem.item
|
||||||
|
if _, err := item.Reader.Seek(0, io.SeekStart); err != nil {
|
||||||
|
return fmt.Errorf("failed to seek batch item: %w", err)
|
||||||
|
}
|
||||||
|
var captionOverride *string
|
||||||
|
if item.PreserveCaption {
|
||||||
|
captionOverride = &item.Caption
|
||||||
|
}
|
||||||
|
media, err := t.prepareMedia(ctx, tctx, item.Reader, item.StoragePath, item.Size, captionOverride)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
prepared = append(prepared, *media)
|
||||||
|
}
|
||||||
|
|
||||||
|
builder := tctx.Sender.WithUploader(prepared[0].uploader).To(prepared[0].peer)
|
||||||
|
if len(prepared) == 1 {
|
||||||
|
_, err := builder.Media(ctx, prepared[0].media)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
media := make([]message.MultiMediaOption, len(prepared))
|
||||||
|
for i := range prepared {
|
||||||
|
media[i] = prepared[i].media
|
||||||
|
}
|
||||||
|
if _, err := builder.Album(ctx, media[0], media[1:]...); err != nil {
|
||||||
|
return fmt.Errorf("failed to send media album: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}, retry.Context(ctx), retry.RetryTimes(uint(config.C().Retry)))
|
||||||
}
|
}
|
||||||
|
|
||||||
func (t *Telegram) CannotStream() string {
|
func (t *Telegram) CannotStream() string {
|
||||||
|
|||||||
119
storage/telegram/upload_limit_test.go
Normal file
119
storage/telegram/upload_limit_test.go
Normal file
@@ -0,0 +1,119 @@
|
|||||||
|
package telegram
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/celestix/gotgproto/ext"
|
||||||
|
"github.com/gotd/td/tg"
|
||||||
|
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestMaxUploadFileSize(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
ctx *ext.Context
|
||||||
|
want int64
|
||||||
|
}{
|
||||||
|
{name: "missing context", want: MaxUploadFileSize},
|
||||||
|
{name: "missing self", ctx: new(ext.Context), want: MaxUploadFileSize},
|
||||||
|
{
|
||||||
|
name: "bot",
|
||||||
|
ctx: uploadAccountContext(true, false),
|
||||||
|
want: MaxUploadFileSize,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "premium bot still uses bot limit",
|
||||||
|
ctx: uploadAccountContext(true, true),
|
||||||
|
want: MaxUploadFileSize,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "regular user",
|
||||||
|
ctx: uploadAccountContext(false, false),
|
||||||
|
want: MaxUploadFileSize,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "premium user",
|
||||||
|
ctx: uploadAccountContext(false, true),
|
||||||
|
want: PremiumMaxUploadFileSize,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
if got := maxUploadFileSize(tt.ctx); got != tt.want {
|
||||||
|
t.Fatalf("maxUploadFileSize() = %d, want %d", got, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func uploadAccountContext(bot, premium bool) *ext.Context {
|
||||||
|
self := new(tg.User)
|
||||||
|
self.SetBot(bot)
|
||||||
|
self.SetPremium(premium)
|
||||||
|
return &ext.Context{Self: self}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSplitSizeUsesUploaderLimit(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
splitSizeMB int64
|
||||||
|
accountLimit int64
|
||||||
|
want int64
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "bot default",
|
||||||
|
accountLimit: MaxUploadFileSize,
|
||||||
|
want: MaxUploadFileSize,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "premium default",
|
||||||
|
accountLimit: PremiumMaxUploadFileSize,
|
||||||
|
want: PremiumMaxUploadFileSize,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "explicit lower limit",
|
||||||
|
splitSizeMB: 1500,
|
||||||
|
accountLimit: PremiumMaxUploadFileSize,
|
||||||
|
want: 1500 * 1024 * 1024,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "explicit limit is capped by account",
|
||||||
|
splitSizeMB: 5000,
|
||||||
|
accountLimit: PremiumMaxUploadFileSize,
|
||||||
|
want: PremiumMaxUploadFileSize,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
telegramStorage := Telegram{}
|
||||||
|
telegramStorage.config.SplitSizeMB = tt.splitSizeMB
|
||||||
|
if got := telegramStorage.splitSize(tt.accountLimit); got != tt.want {
|
||||||
|
t.Fatalf("splitSize() = %d, want %d", got, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBatchUploadLimitAppliesPerFile(t *testing.T) {
|
||||||
|
telegramStorage := Telegram{}
|
||||||
|
tctx := uploadAccountContext(true, false)
|
||||||
|
data := []byte("\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isom")
|
||||||
|
itemSize := int64(MaxUploadFileSize/2 + 1)
|
||||||
|
|
||||||
|
for index := 0; index < 2; index++ {
|
||||||
|
item, err := telegramStorage.inspectBatchItem(tctx, storagetypes.BatchItem{
|
||||||
|
Reader: bytes.NewReader(data),
|
||||||
|
StoragePath: "video.mp4",
|
||||||
|
Size: itemSize,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("inspectBatchItem() failed for item %d: %v", index, err)
|
||||||
|
}
|
||||||
|
if item.useSingleSave {
|
||||||
|
t.Fatalf("item %d was treated as oversized even though only the batch total exceeds the limit", index)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -5,6 +5,7 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"os"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/celestix/gotgproto/ext"
|
"github.com/celestix/gotgproto/ext"
|
||||||
@@ -14,6 +15,34 @@ import (
|
|||||||
"github.com/yapingcat/gomedia/go-mp4"
|
"github.com/yapingcat/gomedia/go-mp4"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// sourceFile returns a filesystem path to the media for ffmpeg/ffprobe. Those
|
||||||
|
// tools need a SEEKABLE input: feeding them a pipe (pipe:0) fails for
|
||||||
|
// non-faststart MP4s whose moov atom is at the END of the file, because ffmpeg
|
||||||
|
// cannot seek backwards on a stream to read it. That silently broke thumbnail
|
||||||
|
// and metadata extraction for a large share of downloaded videos. If rs is
|
||||||
|
// already an *os.File we use it directly; otherwise we spool it to a temp file
|
||||||
|
// and return a cleanup func.
|
||||||
|
func sourceFile(rs io.ReadSeeker) (path string, cleanup func(), err error) {
|
||||||
|
noop := func() {}
|
||||||
|
if f, ok := rs.(*os.File); ok {
|
||||||
|
return f.Name(), noop, nil
|
||||||
|
}
|
||||||
|
if _, err = rs.Seek(0, io.SeekStart); err != nil {
|
||||||
|
return "", noop, err
|
||||||
|
}
|
||||||
|
tf, err := os.CreateTemp("", "saveany-media-*.tmp")
|
||||||
|
if err != nil {
|
||||||
|
return "", noop, err
|
||||||
|
}
|
||||||
|
if _, err = io.Copy(tf, rs); err != nil {
|
||||||
|
tf.Close()
|
||||||
|
os.Remove(tf.Name())
|
||||||
|
return "", noop, err
|
||||||
|
}
|
||||||
|
tf.Close()
|
||||||
|
return tf.Name(), func() { os.Remove(tf.Name()) }, nil
|
||||||
|
}
|
||||||
|
|
||||||
type VideoMetadata struct {
|
type VideoMetadata struct {
|
||||||
Duration int
|
Duration int
|
||||||
Width int
|
Width int
|
||||||
@@ -52,16 +81,14 @@ func getMP4Meta(rs io.ReadSeeker) (metadata *VideoMetadata, err error) {
|
|||||||
|
|
||||||
// getVideoMetadata uses ffprobe to get video metadata
|
// getVideoMetadata uses ffprobe to get video metadata
|
||||||
func getVideoMetadata(rs io.ReadSeeker) (*VideoMetadata, error) {
|
func getVideoMetadata(rs io.ReadSeeker) (*VideoMetadata, error) {
|
||||||
pipeReader, pipeWriter := io.Pipe()
|
path, cleanup, err := sourceFile(rs)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer cleanup()
|
||||||
|
|
||||||
go func() {
|
result, err := ffmpeg.ProbeWithTimeout(
|
||||||
defer pipeWriter.Close()
|
path,
|
||||||
rs.Seek(0, io.SeekStart)
|
|
||||||
io.Copy(pipeWriter, rs)
|
|
||||||
}()
|
|
||||||
|
|
||||||
result, err := ffmpeg.ProbeReaderWithTimeout(
|
|
||||||
pipeReader,
|
|
||||||
time.Second*10,
|
time.Second*10,
|
||||||
ffmpeg.KwArgs{
|
ffmpeg.KwArgs{
|
||||||
"select_streams": "v:0",
|
"select_streams": "v:0",
|
||||||
@@ -114,25 +141,22 @@ func extractThumbFrame(rs io.ReadSeeker) ([]byte, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func extractFrameAt(rs io.ReadSeeker, timestamp float64) ([]byte, error) {
|
func extractFrameAt(rs io.ReadSeeker, timestamp float64) ([]byte, error) {
|
||||||
pipeReader, pipeWriter := io.Pipe()
|
path, cleanup, err := sourceFile(rs)
|
||||||
|
if err != nil {
|
||||||
go func() {
|
return nil, err
|
||||||
defer pipeWriter.Close()
|
}
|
||||||
rs.Seek(0, io.SeekStart)
|
defer cleanup()
|
||||||
io.Copy(pipeWriter, rs)
|
|
||||||
}()
|
|
||||||
|
|
||||||
var out bytes.Buffer
|
var out bytes.Buffer
|
||||||
|
|
||||||
err := ffmpeg.
|
err = ffmpeg.
|
||||||
Input("pipe:0", ffmpeg.KwArgs{
|
Input(path, ffmpeg.KwArgs{
|
||||||
"ss": fmt.Sprintf("%.3f", timestamp),
|
"ss": fmt.Sprintf("%.3f", timestamp),
|
||||||
}).
|
}).
|
||||||
Output("pipe:1", ffmpeg.KwArgs{
|
Output("pipe:1", ffmpeg.KwArgs{
|
||||||
"vframes": 1,
|
"vframes": 1,
|
||||||
"f": "mjpeg",
|
"f": "mjpeg",
|
||||||
}).
|
}).
|
||||||
WithInput(pipeReader).
|
|
||||||
WithOutput(&out).
|
WithOutput(&out).
|
||||||
OverWriteOutput().
|
OverWriteOutput().
|
||||||
Run()
|
Run()
|
||||||
|
|||||||
358
storage/telegram/video_split.go
Normal file
358
storage/telegram/video_split.go
Normal file
@@ -0,0 +1,358 @@
|
|||||||
|
package telegram
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"math"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"path"
|
||||||
|
"path/filepath"
|
||||||
|
"sort"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/celestix/gotgproto/ext"
|
||||||
|
"github.com/charmbracelet/log"
|
||||||
|
"github.com/gotd/td/telegram/message"
|
||||||
|
"github.com/rs/xid"
|
||||||
|
|
||||||
|
"github.com/krau/SaveAny-Bot/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
videoPartTargetRatio = 0.95
|
||||||
|
videoSplitAttempts = 4
|
||||||
|
minSegmentDuration = 1.0
|
||||||
|
maxLosslessVideoParts = 10
|
||||||
|
)
|
||||||
|
|
||||||
|
type losslessVideoPart struct {
|
||||||
|
Path string
|
||||||
|
Name string
|
||||||
|
Size int64
|
||||||
|
}
|
||||||
|
|
||||||
|
type mediaToolRunner func(ctx context.Context, name string, args ...string) ([]byte, error)
|
||||||
|
|
||||||
|
var runMediaTool mediaToolRunner = func(ctx context.Context, name string, args ...string) ([]byte, error) {
|
||||||
|
cmd := exec.CommandContext(ctx, name, args...)
|
||||||
|
output, err := cmd.CombinedOutput()
|
||||||
|
if err != nil {
|
||||||
|
message := strings.TrimSpace(string(output))
|
||||||
|
if message == "" {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("%w: %s", err, message)
|
||||||
|
}
|
||||||
|
return output, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func createLosslessVideoParts(
|
||||||
|
ctx context.Context,
|
||||||
|
r io.ReadSeeker,
|
||||||
|
filename string,
|
||||||
|
fileSize, maxPartSize int64,
|
||||||
|
) ([]losslessVideoPart, func(), error) {
|
||||||
|
inputPath, sourceCleanup, err := sourceFile(r)
|
||||||
|
if err != nil {
|
||||||
|
return nil, func() {}, fmt.Errorf("failed to prepare seekable video source: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tempBase := config.C().Temp.BasePath
|
||||||
|
if err := os.MkdirAll(tempBase, 0o755); err != nil {
|
||||||
|
sourceCleanup()
|
||||||
|
return nil, func() {}, fmt.Errorf("failed to create video split base directory: %w", err)
|
||||||
|
}
|
||||||
|
tempDir, err := os.MkdirTemp(tempBase, "telegram-video-split-*")
|
||||||
|
if err != nil {
|
||||||
|
sourceCleanup()
|
||||||
|
return nil, func() {}, fmt.Errorf("failed to create video split directory: %w", err)
|
||||||
|
}
|
||||||
|
cleanup := func() {
|
||||||
|
if err := os.RemoveAll(tempDir); err != nil {
|
||||||
|
log.FromContext(ctx).Warnf("Failed to clean lossless video parts: %s", err)
|
||||||
|
}
|
||||||
|
sourceCleanup()
|
||||||
|
}
|
||||||
|
|
||||||
|
parts, err := splitLosslessVideo(ctx, inputPath, tempDir, filename, fileSize, maxPartSize)
|
||||||
|
if err != nil {
|
||||||
|
cleanup()
|
||||||
|
return nil, func() {}, err
|
||||||
|
}
|
||||||
|
return parts, cleanup, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func splitLosslessVideo(
|
||||||
|
ctx context.Context,
|
||||||
|
inputPath, outputDir, filename string,
|
||||||
|
fileSize, maxPartSize int64,
|
||||||
|
) ([]losslessVideoPart, error) {
|
||||||
|
if fileSize <= maxPartSize {
|
||||||
|
return nil, fmt.Errorf("video size %d does not exceed part limit %d", fileSize, maxPartSize)
|
||||||
|
}
|
||||||
|
if maxPartSize <= 0 {
|
||||||
|
return nil, fmt.Errorf("invalid video part limit: %d", maxPartSize)
|
||||||
|
}
|
||||||
|
|
||||||
|
duration, err := probeMediaDuration(ctx, inputPath)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to probe source video duration: %w", err)
|
||||||
|
}
|
||||||
|
if duration < minSegmentDuration {
|
||||||
|
return nil, fmt.Errorf("invalid source video duration: %.3f", duration)
|
||||||
|
}
|
||||||
|
|
||||||
|
targetSize := int64(float64(maxPartSize) * videoPartTargetRatio)
|
||||||
|
segmentDuration := initialSegmentDuration(duration, fileSize, targetSize)
|
||||||
|
extension := videoPartExtension(filename)
|
||||||
|
outputPattern := filepath.Join(outputDir, "part-%03d"+extension)
|
||||||
|
|
||||||
|
var lastOversize int64
|
||||||
|
for attempt := 0; attempt < videoSplitAttempts; attempt++ {
|
||||||
|
if err := clearVideoParts(outputDir); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if err := runFFmpegSegment(ctx, inputPath, outputPattern, extension, segmentDuration); err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to losslessly split video: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
parts, largest, err := collectVideoParts(ctx, outputDir, filename)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if len(parts) < 2 {
|
||||||
|
return nil, fmt.Errorf("video split produced %d part(s), expected at least 2", len(parts))
|
||||||
|
}
|
||||||
|
if len(parts) > maxLosslessVideoParts {
|
||||||
|
return nil, fmt.Errorf(
|
||||||
|
"video split produced %d parts, exceeding the single-album limit of %d",
|
||||||
|
len(parts),
|
||||||
|
maxLosslessVideoParts,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
if largest <= maxPartSize {
|
||||||
|
return parts, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
lastOversize = largest
|
||||||
|
segmentDuration *= float64(targetSize) / float64(largest)
|
||||||
|
if segmentDuration < minSegmentDuration {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil, fmt.Errorf(
|
||||||
|
"unable to keep lossless video parts below %d bytes; largest part was %d bytes",
|
||||||
|
maxPartSize,
|
||||||
|
lastOversize,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
func initialSegmentDuration(duration float64, fileSize, targetSize int64) float64 {
|
||||||
|
partCount := math.Ceil(float64(fileSize) / float64(targetSize))
|
||||||
|
if partCount < 2 {
|
||||||
|
partCount = 2
|
||||||
|
}
|
||||||
|
segmentDuration := duration / partCount
|
||||||
|
if segmentDuration < minSegmentDuration {
|
||||||
|
return minSegmentDuration
|
||||||
|
}
|
||||||
|
return segmentDuration
|
||||||
|
}
|
||||||
|
|
||||||
|
func videoPartExtension(filename string) string {
|
||||||
|
extension := strings.ToLower(filepath.Ext(filepath.Base(filename)))
|
||||||
|
switch extension {
|
||||||
|
case ".avi", ".m4v", ".mkv", ".mov", ".mp4", ".ts", ".webm":
|
||||||
|
return extension
|
||||||
|
default:
|
||||||
|
return ".mp4"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func runFFmpegSegment(
|
||||||
|
ctx context.Context,
|
||||||
|
inputPath, outputPattern, extension string,
|
||||||
|
segmentDuration float64,
|
||||||
|
) error {
|
||||||
|
args := []string{
|
||||||
|
"-hide_banner",
|
||||||
|
"-loglevel", "error",
|
||||||
|
"-nostdin",
|
||||||
|
"-y",
|
||||||
|
"-i", inputPath,
|
||||||
|
"-map", "0",
|
||||||
|
"-map_metadata", "0",
|
||||||
|
"-c", "copy",
|
||||||
|
"-f", "segment",
|
||||||
|
"-segment_time", strconv.FormatFloat(segmentDuration, 'f', 3, 64),
|
||||||
|
"-segment_start_number", "1",
|
||||||
|
"-reset_timestamps", "1",
|
||||||
|
"-avoid_negative_ts", "make_zero",
|
||||||
|
}
|
||||||
|
if extension == ".mp4" || extension == ".m4v" || extension == ".mov" {
|
||||||
|
args = append(args, "-segment_format_options", "movflags=+faststart")
|
||||||
|
}
|
||||||
|
args = append(args, outputPattern)
|
||||||
|
_, err := runMediaTool(ctx, "ffmpeg", args...)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func probeMediaDuration(ctx context.Context, filePath string) (float64, error) {
|
||||||
|
output, err := runMediaTool(
|
||||||
|
ctx,
|
||||||
|
"ffprobe",
|
||||||
|
"-v", "error",
|
||||||
|
"-show_entries", "format=duration",
|
||||||
|
"-of", "default=noprint_wrappers=1:nokey=1",
|
||||||
|
filePath,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
duration, err := strconv.ParseFloat(strings.TrimSpace(string(output)), 64)
|
||||||
|
if err != nil {
|
||||||
|
return 0, fmt.Errorf("invalid ffprobe duration %q: %w", strings.TrimSpace(string(output)), err)
|
||||||
|
}
|
||||||
|
return duration, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func clearVideoParts(outputDir string) error {
|
||||||
|
entries, err := os.ReadDir(outputDir)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to read video split directory: %w", err)
|
||||||
|
}
|
||||||
|
for _, entry := range entries {
|
||||||
|
if entry.IsDir() {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := os.Remove(filepath.Join(outputDir, entry.Name())); err != nil {
|
||||||
|
return fmt.Errorf("failed to clear old video part %s: %w", entry.Name(), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func collectVideoParts(
|
||||||
|
ctx context.Context,
|
||||||
|
outputDir, filename string,
|
||||||
|
) ([]losslessVideoPart, int64, error) {
|
||||||
|
matches, err := filepath.Glob(filepath.Join(outputDir, "part-*"))
|
||||||
|
if err != nil {
|
||||||
|
return nil, 0, fmt.Errorf("failed to list video parts: %w", err)
|
||||||
|
}
|
||||||
|
sort.Strings(matches)
|
||||||
|
if len(matches) == 0 {
|
||||||
|
return nil, 0, fmt.Errorf("ffmpeg did not produce any video parts")
|
||||||
|
}
|
||||||
|
|
||||||
|
extension := videoPartExtension(filename)
|
||||||
|
base := strings.TrimSuffix(filepath.Base(filename), filepath.Ext(filepath.Base(filename)))
|
||||||
|
if base == "" || base == "." {
|
||||||
|
base = xid.New().String()
|
||||||
|
}
|
||||||
|
parts := make([]losslessVideoPart, 0, len(matches))
|
||||||
|
var largest int64
|
||||||
|
for index, match := range matches {
|
||||||
|
info, err := os.Stat(match)
|
||||||
|
if err != nil {
|
||||||
|
return nil, 0, fmt.Errorf("failed to stat video part %s: %w", match, err)
|
||||||
|
}
|
||||||
|
if info.Size() <= 0 {
|
||||||
|
return nil, 0, fmt.Errorf("video part %s is empty", match)
|
||||||
|
}
|
||||||
|
if _, err := probeMediaDuration(ctx, match); err != nil {
|
||||||
|
return nil, 0, fmt.Errorf("failed to validate video part %s: %w", match, err)
|
||||||
|
}
|
||||||
|
if info.Size() > largest {
|
||||||
|
largest = info.Size()
|
||||||
|
}
|
||||||
|
parts = append(parts, losslessVideoPart{
|
||||||
|
Path: match,
|
||||||
|
Name: fmt.Sprintf("%s.part%03d%s", base, index+1, extension),
|
||||||
|
Size: info.Size(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return parts, largest, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func partStoragePath(storagePath, partName string) string {
|
||||||
|
directory := path.Dir(path.Clean(storagePath))
|
||||||
|
if directory == "." || directory == "/" {
|
||||||
|
return partName
|
||||||
|
}
|
||||||
|
return path.Join(directory, partName)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Telegram) uploadLosslessVideoParts(
|
||||||
|
ctx context.Context,
|
||||||
|
tctx *ext.Context,
|
||||||
|
storagePath string,
|
||||||
|
parts []losslessVideoPart,
|
||||||
|
sourceCaption *string,
|
||||||
|
) error {
|
||||||
|
if len(parts) == 0 {
|
||||||
|
return fmt.Errorf("no lossless video parts to upload")
|
||||||
|
}
|
||||||
|
if len(parts) > maxLosslessVideoParts {
|
||||||
|
return fmt.Errorf(
|
||||||
|
"refusing to upload %d lossless video parts as multiple albums; maximum is %d",
|
||||||
|
len(parts),
|
||||||
|
maxLosslessVideoParts,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
prepared := make([]preparedMedia, 0, len(parts))
|
||||||
|
for index, part := range parts {
|
||||||
|
partFile, err := os.Open(part.Path)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to open video part %s: %w", part.Name, err)
|
||||||
|
}
|
||||||
|
media, prepareErr := t.prepareMedia(
|
||||||
|
ctx,
|
||||||
|
tctx,
|
||||||
|
partFile,
|
||||||
|
partStoragePath(storagePath, part.Name),
|
||||||
|
part.Size,
|
||||||
|
videoPartCaption(sourceCaption, index),
|
||||||
|
)
|
||||||
|
closeErr := partFile.Close()
|
||||||
|
if prepareErr != nil {
|
||||||
|
return fmt.Errorf("failed to prepare video part %s: %w", part.Name, prepareErr)
|
||||||
|
}
|
||||||
|
if closeErr != nil {
|
||||||
|
return fmt.Errorf("failed to close video part %s: %w", part.Name, closeErr)
|
||||||
|
}
|
||||||
|
prepared = append(prepared, *media)
|
||||||
|
}
|
||||||
|
|
||||||
|
builder := tctx.Sender.WithUploader(prepared[0].uploader).To(prepared[0].peer)
|
||||||
|
if len(prepared) == 1 {
|
||||||
|
if _, err := builder.Media(ctx, prepared[0].media); err != nil {
|
||||||
|
return fmt.Errorf("failed to send video part: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
media := make([]message.MultiMediaOption, len(prepared))
|
||||||
|
for index := range prepared {
|
||||||
|
media[index] = prepared[index].media
|
||||||
|
}
|
||||||
|
if _, err := builder.Album(ctx, media[0], media[1:]...); err != nil {
|
||||||
|
return fmt.Errorf("failed to send video parts as album: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func videoPartCaption(sourceCaption *string, index int) *string {
|
||||||
|
if sourceCaption == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if index == 0 {
|
||||||
|
return sourceCaption
|
||||||
|
}
|
||||||
|
empty := ""
|
||||||
|
return &empty
|
||||||
|
}
|
||||||
174
storage/telegram/video_split_test.go
Normal file
174
storage/telegram/video_split_test.go
Normal file
@@ -0,0 +1,174 @@
|
|||||||
|
package telegram
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestInitialSegmentDuration(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
duration float64
|
||||||
|
fileSize int64
|
||||||
|
targetSize int64
|
||||||
|
want float64
|
||||||
|
}{
|
||||||
|
{name: "two balanced parts", duration: 120, fileSize: 2200, targetSize: 1900, want: 60},
|
||||||
|
{name: "three balanced parts", duration: 120, fileSize: 3900, targetSize: 1900, want: 40},
|
||||||
|
{name: "minimum one second", duration: 1.5, fileSize: 3900, targetSize: 1900, want: 1},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
got := initialSegmentDuration(tt.duration, tt.fileSize, tt.targetSize)
|
||||||
|
if got != tt.want {
|
||||||
|
t.Fatalf("initialSegmentDuration() = %v, want %v", got, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVideoPartExtension(t *testing.T) {
|
||||||
|
tests := map[string]string{
|
||||||
|
"clip.MOV": ".mov",
|
||||||
|
"clip.mkv": ".mkv",
|
||||||
|
"clip.exe": ".mp4",
|
||||||
|
"clip": ".mp4",
|
||||||
|
}
|
||||||
|
for input, want := range tests {
|
||||||
|
if got := videoPartExtension(input); got != want {
|
||||||
|
t.Fatalf("videoPartExtension(%q) = %q, want %q", input, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPartStoragePath(t *testing.T) {
|
||||||
|
tests := map[string]string{
|
||||||
|
"clip.mov": "clip.part001.mov",
|
||||||
|
"123456/clip.mov": "123456/clip.part001.mov",
|
||||||
|
"/123456/clip.mov": "/123456/clip.part001.mov",
|
||||||
|
"folder/sub/clip.mov": "folder/sub/clip.part001.mov",
|
||||||
|
}
|
||||||
|
for input, want := range tests {
|
||||||
|
if got := partStoragePath(input, "clip.part001.mov"); got != want {
|
||||||
|
t.Fatalf("partStoragePath(%q) = %q, want %q", input, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVideoPartCaption(t *testing.T) {
|
||||||
|
if got := videoPartCaption(nil, 0); got != nil {
|
||||||
|
t.Fatalf("caption without source = %q, want nil", *got)
|
||||||
|
}
|
||||||
|
|
||||||
|
source := "original caption"
|
||||||
|
first := videoPartCaption(&source, 0)
|
||||||
|
if first == nil || *first != source {
|
||||||
|
t.Fatalf("first part caption = %v, want %q", first, source)
|
||||||
|
}
|
||||||
|
second := videoPartCaption(&source, 1)
|
||||||
|
if second == nil || *second != "" {
|
||||||
|
t.Fatalf("second part caption = %v, want an explicit empty caption", second)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSplitLosslessVideoRetriesOversizedPart(t *testing.T) {
|
||||||
|
tempDir := t.TempDir()
|
||||||
|
inputPath := filepath.Join(tempDir, "source.mov")
|
||||||
|
if err := os.WriteFile(inputPath, []byte("source"), 0o600); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
outputDir := filepath.Join(tempDir, "parts")
|
||||||
|
if err := os.Mkdir(outputDir, 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
originalRunner := runMediaTool
|
||||||
|
t.Cleanup(func() { runMediaTool = originalRunner })
|
||||||
|
ffmpegCalls := 0
|
||||||
|
runMediaTool = func(ctx context.Context, name string, args ...string) ([]byte, error) {
|
||||||
|
switch name {
|
||||||
|
case "ffprobe":
|
||||||
|
return []byte("120.0\n"), nil
|
||||||
|
case "ffmpeg":
|
||||||
|
ffmpegCalls++
|
||||||
|
pattern := args[len(args)-1]
|
||||||
|
sizes := []int{2001, 199}
|
||||||
|
if ffmpegCalls > 1 {
|
||||||
|
sizes = []int{1100, 1100}
|
||||||
|
}
|
||||||
|
for index, size := range sizes {
|
||||||
|
partPath := strings.Replace(pattern, "%03d", fmt.Sprintf("%03d", index+1), 1)
|
||||||
|
if err := os.WriteFile(partPath, make([]byte, size), 0o600); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil, nil
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("unexpected media tool %q", name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
parts, err := splitLosslessVideo(t.Context(), inputPath, outputDir, "example.mov", 2200, 2000)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("splitLosslessVideo() failed: %v", err)
|
||||||
|
}
|
||||||
|
if ffmpegCalls != 2 {
|
||||||
|
t.Fatalf("ffmpeg calls = %d, want 2", ffmpegCalls)
|
||||||
|
}
|
||||||
|
if len(parts) != 2 {
|
||||||
|
t.Fatalf("parts = %d, want 2", len(parts))
|
||||||
|
}
|
||||||
|
if parts[0].Name != "example.part001.mov" || parts[1].Name != "example.part002.mov" {
|
||||||
|
t.Fatalf("unexpected part names: %#v", parts)
|
||||||
|
}
|
||||||
|
for _, part := range parts {
|
||||||
|
if part.Size > 2000 {
|
||||||
|
t.Fatalf("part %s exceeds limit: %d", part.Name, part.Size)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSplitLosslessVideoRejectsMultipleAlbums(t *testing.T) {
|
||||||
|
tempDir := t.TempDir()
|
||||||
|
inputPath := filepath.Join(tempDir, "source.mov")
|
||||||
|
if err := os.WriteFile(inputPath, []byte("source"), 0o600); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
outputDir := filepath.Join(tempDir, "parts")
|
||||||
|
if err := os.Mkdir(outputDir, 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
originalRunner := runMediaTool
|
||||||
|
t.Cleanup(func() { runMediaTool = originalRunner })
|
||||||
|
runMediaTool = func(ctx context.Context, name string, args ...string) ([]byte, error) {
|
||||||
|
switch name {
|
||||||
|
case "ffprobe":
|
||||||
|
return []byte("120.0\n"), nil
|
||||||
|
case "ffmpeg":
|
||||||
|
pattern := args[len(args)-1]
|
||||||
|
for index := 0; index < maxLosslessVideoParts+1; index++ {
|
||||||
|
partPath := strings.Replace(pattern, "%03d", fmt.Sprintf("%03d", index+1), 1)
|
||||||
|
if err := os.WriteFile(partPath, make([]byte, 100), 0o600); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil, nil
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("unexpected media tool %q", name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := splitLosslessVideo(t.Context(), inputPath, outputDir, "example.mov", 2200, 2000)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("splitLosslessVideo() succeeded with more than one album of parts")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "single-album limit") {
|
||||||
|
t.Fatalf("splitLosslessVideo() error = %q, want single-album limit", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user