优化

2026-10-02 20:26:36 +08:00 · 2026-08-27 23:40:12 +08:00
parent d3051eaffe
commit 6e8eac9887
26 changed files with 2018 additions and 95 deletions
@@ -15,6 +15,11 @@ type LibraryRepository struct{ db *gorm.DB }

 // Create persists a new library row.
 func (r *LibraryRepository) Create(ctx context.Context, l *model.Library) error {
+	if l != nil && l.SortOrder == 0 {
+		var maxSort int
+		_ = r.db.WithContext(ctx).Model(&model.Library{}).Select("COALESCE(MAX(sort_order), -1)").Scan(&maxSort)
+		l.SortOrder = maxSort + 1
+	}
 	return r.db.WithContext(ctx).Create(l).Error
 }

@@ -23,6 +28,11 @@ func (r *LibraryRepository) CreateWithRoots(ctx context.Context, l *model.Librar
 		return r.Create(ctx, l)
 	}
 	return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		if l != nil && l.SortOrder == 0 {
+			var maxSort int
+			_ = tx.Model(&model.Library{}).Select("COALESCE(MAX(sort_order), -1)").Scan(&maxSort)
+			l.SortOrder = maxSort + 1
+		}
 		if err := tx.Create(l).Error; err != nil {
 			return err
 		}
@@ -30,9 +30,10 @@ type Container struct {
 	StrmSyncPath   *StrmSyncPathRepository
 	StrmSyncRecord *StrmSyncRecordRepository
 	StrmDownload   *StrmDownloadTaskRepository
-	StrmUpload     *StrmUploadTaskRepository
-	StrmDirCache   *StrmDirCacheRepository
-}
+		StrmUpload     *StrmUploadTaskRepository
+		StrmDirCache   *StrmDirCacheRepository
+		ScrapeTask     *ScrapeTaskRepository
+	}

 // New 将每个 repository 连接到单个 *gorm.DB。
 func New(db *gorm.DB) *Container {
@@ -60,5 +61,6 @@ func New(db *gorm.DB) *Container {
 		StrmDownload:   &StrmDownloadTaskRepository{db: db},
 		StrmUpload:     &StrmUploadTaskRepository{db: db},
 		StrmDirCache:   &StrmDirCacheRepository{db: db},
+		ScrapeTask:     &ScrapeTaskRepository{db: db},
 	}
 }
@@ -0,0 +1,274 @@
+package repository
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"sync"
+	"time"
+
+	"gorm.io/gorm"
+
+	"github.com/ShukeBta/MMTL/internal/model"
+)
+
+var scrapeClaimMu sync.Mutex
+
+// ScrapeTaskRepository persists model.ScrapeTask.
+type ScrapeTaskRepository struct{ db *gorm.DB }
+
+func (r *ScrapeTaskRepository) Create(ctx context.Context, t *model.ScrapeTask) error {
+	return withSQLiteBusyRetry(ctx, func() error {
+		return r.db.WithContext(ctx).Create(t).Error
+	})
+}
+
+func (r *ScrapeTaskRepository) CreateBatch(ctx context.Context, tasks []model.ScrapeTask) error {
+	if len(tasks) == 0 {
+		return nil
+	}
+	return withSQLiteBusyRetry(ctx, func() error {
+		return r.db.WithContext(ctx).CreateInBatches(tasks, 100).Error
+	})
+}
+
+func (r *ScrapeTaskRepository) FindByID(ctx context.Context, id string) (*model.ScrapeTask, error) {
+	var t model.ScrapeTask
+	err := r.db.WithContext(ctx).Where("id = ?", id).First(&t).Error
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		return nil, nil
+	}
+	return &t, err
+}
+
+func (r *ScrapeTaskRepository) FindActiveByMediaID(ctx context.Context, mediaID string) (*model.ScrapeTask, error) {
+	var t model.ScrapeTask
+	err := r.db.WithContext(ctx).
+		Where("media_id = ? AND status IN ?", mediaID, []string{model.ScrapeTaskPending, model.ScrapeTaskRunning}).
+		First(&t).Error
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		return nil, nil
+	}
+	return &t, err
+}
+
+func (r *ScrapeTaskRepository) List(ctx context.Context, status string, page, pageSize int) ([]model.ScrapeTask, int64, error) {
+	if page < 1 {
+		page = 1
+	}
+	if pageSize < 1 || pageSize > 200 {
+		pageSize = 50
+	}
+	q := r.db.WithContext(ctx).Model(&model.ScrapeTask{})
+	if strings.TrimSpace(status) != "" && status != "all" {
+		q = q.Where("status = ?", strings.TrimSpace(status))
+	}
+	var total int64
+	if err := q.Count(&total).Error; err != nil {
+		return nil, 0, err
+	}
+	var rows []model.ScrapeTask
+	err := q.Order("created_at desc").
+		Offset((page - 1) * pageSize).
+		Limit(pageSize).
+		Find(&rows).Error
+	return rows, total, err
+}
+
+func (r *ScrapeTaskRepository) CountByStatus(ctx context.Context) (map[string]int64, error) {
+	var rows []struct {
+		Status string
+		Count  int64
+	}
+	err := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
+		Select("status, count(*) as count").
+		Group("status").Scan(&rows).Error
+	if err != nil {
+		return nil, err
+	}
+	out := map[string]int64{}
+	for _, row := range rows {
+		out[row.Status] = row.Count
+	}
+	return out, nil
+}
+
+// ClaimPending picks pending scrape tasks and marks them running.
+func (r *ScrapeTaskRepository) ClaimPending(ctx context.Context, limit int) ([]model.ScrapeTask, error) {
+	scrapeClaimMu.Lock()
+	defer scrapeClaimMu.Unlock()
+
+	var rows []model.ScrapeTask
+	err := withSQLiteBusyRetry(ctx, func() error {
+		return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+			if err := tx.Where("status = ?", model.ScrapeTaskPending).
+				Order("created_at asc").Limit(limit).Find(&rows).Error; err != nil {
+				return err
+			}
+			if len(rows) == 0 {
+				return nil
+			}
+			ids := make([]string, 0, len(rows))
+			now := time.Now()
+			for i := range rows {
+				ids = append(ids, rows[i].ID)
+				rows[i].Status = model.ScrapeTaskRunning
+				rows[i].StartedAt = &now
+			}
+			return tx.Model(&model.ScrapeTask{}).Where("id IN ?", ids).
+				Updates(map[string]any{"status": model.ScrapeTaskRunning, "started_at": now}).Error
+		})
+	})
+	if err != nil {
+		return nil, err
+	}
+	return rows, nil
+}
+
+func (r *ScrapeTaskRepository) Update(ctx context.Context, t *model.ScrapeTask) error {
+	return withSQLiteBusyRetry(ctx, func() error {
+		return r.db.WithContext(ctx).Model(&model.ScrapeTask{}).Where("id = ?", t.ID).Updates(map[string]any{
+			"status":        t.Status,
+			"error":         t.Error,
+			"provider":      t.Provider,
+			"matched_title": t.MatchedTitle,
+			"matched_year":  t.MatchedYear,
+			"poster_url":    t.PosterURL,
+			"backdrop_url":  t.BackdropURL,
+			"retry_count":   t.RetryCount,
+			"started_at":    t.StartedAt,
+			"finished_at":   t.FinishedAt,
+			"updated_at":    time.Now(),
+		}).Error
+	})
+}
+
+func (r *ScrapeTaskRepository) Delete(ctx context.Context, id string) error {
+	return withSQLiteBusyRetry(ctx, func() error {
+		return r.db.WithContext(ctx).Unscoped().Where("id = ?", id).Delete(&model.ScrapeTask{}).Error
+	})
+}
+
+func (r *ScrapeTaskRepository) DeleteBatch(ctx context.Context, ids []string) (int64, error) {
+	if len(ids) == 0 {
+		return 0, nil
+	}
+	var count int64
+	err := withSQLiteBusyRetry(ctx, func() error {
+		res := r.db.WithContext(ctx).Unscoped().Where("id IN ?", ids).Delete(&model.ScrapeTask{})
+		count = res.RowsAffected
+		return res.Error
+	})
+	return count, err
+}
+
+func (r *ScrapeTaskRepository) RetryBatch(ctx context.Context, ids []string) (int64, error) {
+	if len(ids) == 0 {
+		return 0, nil
+	}
+	var count int64
+	err := withSQLiteBusyRetry(ctx, func() error {
+		res := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
+			Where("id IN ? AND status IN ?", ids, []string{model.ScrapeTaskFailed, model.ScrapeTaskCanceled}).
+			Updates(map[string]any{
+				"status":      model.ScrapeTaskPending,
+				"error":       "",
+				"retry_count": 0,
+				"started_at":  nil,
+				"finished_at": nil,
+				"updated_at":  time.Now(),
+			})
+		count = res.RowsAffected
+		return res.Error
+	})
+	return count, err
+}
+
+func (r *ScrapeTaskRepository) CancelBatch(ctx context.Context, ids []string) (int64, error) {
+	if len(ids) == 0 {
+		return 0, nil
+	}
+	now := time.Now()
+	var count int64
+	err := withSQLiteBusyRetry(ctx, func() error {
+		res := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
+			Where("id IN ? AND status IN ?", ids, []string{model.ScrapeTaskPending, model.ScrapeTaskRunning}).
+			Updates(map[string]any{
+				"status":      model.ScrapeTaskCanceled,
+				"error":       "已批量取消",
+				"finished_at": now,
+				"updated_at":  now,
+			})
+		count = res.RowsAffected
+		return res.Error
+	})
+	return count, err
+}
+
+func (r *ScrapeTaskRepository) ClearDone(ctx context.Context) (int64, error) {
+	var count int64
+	err := withSQLiteBusyRetry(ctx, func() error {
+		res := r.db.WithContext(ctx).Unscoped().Where("status = ?", model.ScrapeTaskDone).Delete(&model.ScrapeTask{})
+		count = res.RowsAffected
+		return res.Error
+	})
+	return count, err
+}
+
+func (r *ScrapeTaskRepository) ClearFinished(ctx context.Context) (int64, error) {
+	var count int64
+	err := withSQLiteBusyRetry(ctx, func() error {
+		res := r.db.WithContext(ctx).Unscoped().Where("status IN ?", []string{model.ScrapeTaskDone, model.ScrapeTaskFailed, model.ScrapeTaskCanceled}).
+			Delete(&model.ScrapeTask{})
+		count = res.RowsAffected
+		return res.Error
+	})
+	return count, err
+}
+
+func (r *ScrapeTaskRepository) ClearCanceled(ctx context.Context) (int64, error) {
+	var count int64
+	err := withSQLiteBusyRetry(ctx, func() error {
+		res := r.db.WithContext(ctx).Unscoped().Where("status = ?", model.ScrapeTaskCanceled).Delete(&model.ScrapeTask{})
+		count = res.RowsAffected
+		return res.Error
+	})
+	return count, err
+}
+
+func (r *ScrapeTaskRepository) RetryAllFailed(ctx context.Context) (int64, error) {
+	var count int64
+	err := withSQLiteBusyRetry(ctx, func() error {
+		res := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
+			Where("status = ?", model.ScrapeTaskFailed).
+			Updates(map[string]any{
+				"status":      model.ScrapeTaskPending,
+				"error":       "",
+				"retry_count": 0,
+				"started_at":  nil,
+				"finished_at": nil,
+				"updated_at":  time.Now(),
+			})
+		count = res.RowsAffected
+		return res.Error
+	})
+	return count, err
+}
+
+func (r *ScrapeTaskRepository) CancelPending(ctx context.Context) (int64, error) {
+	now := time.Now()
+	var count int64
+	err := withSQLiteBusyRetry(ctx, func() error {
+		res := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
+			Where("status IN ?", []string{model.ScrapeTaskPending, model.ScrapeTaskRunning}).
+			Updates(map[string]any{
+				"status":      model.ScrapeTaskCanceled,
+				"error":       "已批量取消",
+				"finished_at": now,
+				"updated_at":  now,
+			})
+		count = res.RowsAffected
+		return res.Error
+	})
+	return count, err
+}