mirror of https://github.com/stashapp/stash.git
136 lines
2.7 KiB
Go
136 lines
2.7 KiB
Go
package task
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"github.com/stashapp/stash/pkg/job"
|
|
"github.com/stashapp/stash/pkg/logger"
|
|
"github.com/stashapp/stash/pkg/scene"
|
|
"github.com/stashapp/stash/pkg/txn"
|
|
)
|
|
|
|
type MigrateSceneScreenshotsJob struct {
|
|
ScreenshotsPath string
|
|
Input scene.MigrateSceneScreenshotsInput
|
|
SceneRepo scene.HashFinderCoverUpdater
|
|
TxnManager txn.Manager
|
|
}
|
|
|
|
func (j *MigrateSceneScreenshotsJob) Execute(ctx context.Context, progress *job.Progress) {
|
|
var err error
|
|
progress.ExecuteTask("Counting files", func() {
|
|
var count int
|
|
count, err = j.countFiles(ctx)
|
|
progress.SetTotal(count)
|
|
})
|
|
|
|
if err != nil {
|
|
logger.Errorf("Error counting files: %s", err.Error())
|
|
return
|
|
}
|
|
|
|
progress.ExecuteTask("Migrating files", func() {
|
|
err = j.migrateFiles(ctx, progress)
|
|
})
|
|
|
|
if job.IsCancelled(ctx) {
|
|
logger.Info("Cancelled migrating scene screenshots")
|
|
return
|
|
}
|
|
|
|
if err != nil {
|
|
logger.Errorf("Error migrating scene screenshots: %v", err)
|
|
return
|
|
}
|
|
|
|
logger.Infof("Finished migrating scene screenshots")
|
|
}
|
|
|
|
func (j *MigrateSceneScreenshotsJob) countFiles(ctx context.Context) (int, error) {
|
|
f, err := os.Open(j.ScreenshotsPath)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer f.Close()
|
|
|
|
const batchSize = 1000
|
|
ret := 0
|
|
files, err := f.ReadDir(batchSize)
|
|
for err == nil && ctx.Err() == nil {
|
|
ret += len(files)
|
|
|
|
files, err = f.ReadDir(batchSize)
|
|
}
|
|
|
|
if errors.Is(err, io.EOF) {
|
|
// end of directory
|
|
return ret, nil
|
|
}
|
|
|
|
return 0, err
|
|
}
|
|
|
|
func (j *MigrateSceneScreenshotsJob) migrateFiles(ctx context.Context, progress *job.Progress) error {
|
|
f, err := os.Open(j.ScreenshotsPath)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer f.Close()
|
|
|
|
m := scene.ScreenshotMigrator{
|
|
Options: j.Input,
|
|
SceneUpdater: j.SceneRepo,
|
|
TxnManager: j.TxnManager,
|
|
}
|
|
|
|
const batchSize = 1000
|
|
files, err := f.ReadDir(batchSize)
|
|
for err == nil && ctx.Err() == nil {
|
|
for _, f := range files {
|
|
if ctx.Err() != nil {
|
|
return nil
|
|
}
|
|
|
|
progress.ExecuteTask("Migrating file "+f.Name(), func() {
|
|
defer progress.Increment()
|
|
|
|
path := filepath.Join(j.ScreenshotsPath, f.Name())
|
|
|
|
// sanity check - only process files
|
|
if f.IsDir() {
|
|
logger.Warnf("Skipping directory %s", path)
|
|
return
|
|
}
|
|
|
|
// ignore non-jpg files
|
|
if !strings.HasSuffix(f.Name(), ".jpg") {
|
|
return
|
|
}
|
|
|
|
// ignore .thumb files
|
|
if strings.HasSuffix(f.Name(), ".thumb.jpg") {
|
|
return
|
|
}
|
|
|
|
if err := m.MigrateScreenshots(ctx, path); err != nil {
|
|
logger.Errorf("Error migrating screenshots for %s: %v", path, err)
|
|
}
|
|
})
|
|
}
|
|
|
|
files, err = f.ReadDir(batchSize)
|
|
}
|
|
|
|
if errors.Is(err, io.EOF) {
|
|
// end of directory
|
|
return nil
|
|
}
|
|
|
|
return err
|
|
}
|