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 }