308 lines
12 KiB
Go
308 lines
12 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
var ErrInvalidTransition = errors.New("invalid job status transition")
|
|
|
|
type Store struct{ pool *pgxpool.Pool }
|
|
|
|
func New(pool *pgxpool.Pool) *Store { return &Store{pool: pool} }
|
|
|
|
type Episode struct {
|
|
ID, SeriesID, SeriesTitle, Title, BroadcasterName, ThumbnailPath string
|
|
DurationSeconds int
|
|
EndAt *time.Time
|
|
}
|
|
|
|
type Job struct {
|
|
ID int64
|
|
EpisodeID, SeriesID, SeriesTitle string
|
|
Title, BroadcasterName string
|
|
DurationSeconds int
|
|
EndAt *time.Time
|
|
ThumbnailPath, Source, Status string
|
|
Phase, FormatID *string
|
|
DownloadedBytes, TotalBytes *int64
|
|
TotalIsEstimate bool
|
|
SpeedBPS, ETASeconds *float64
|
|
FragmentIndex, FragmentCount *int32
|
|
ErrorKind, ErrorMessage, FilePath *string
|
|
FileSizeBytes *int64
|
|
Attempts int
|
|
HeartbeatAt, StartedAt, FinishedAt *time.Time
|
|
CreatedAt, UpdatedAt time.Time
|
|
}
|
|
|
|
type Progress struct {
|
|
Phase string
|
|
DownloadedBytes *int64
|
|
TotalBytes *int64
|
|
TotalIsEstimate bool
|
|
SpeedBPS *float64
|
|
ETASeconds *float64
|
|
FragmentIndex *int32
|
|
FragmentCount *int32
|
|
FormatID *string
|
|
}
|
|
|
|
type Subscription struct {
|
|
SeriesID, SeriesTitle string
|
|
Enabled bool
|
|
CreatedAt time.Time
|
|
LastCheckedAt *time.Time
|
|
JobCount int64
|
|
}
|
|
|
|
const jobColumns = `id, episode_id, series_id, series_title, title, broadcaster_name,
|
|
duration_seconds, end_at, thumbnail_path, source, status, phase,
|
|
downloaded_bytes, total_bytes, total_is_estimate, speed_bps, eta_seconds,
|
|
fragment_index, fragment_count, format_id, error_kind, error_message,
|
|
file_path, file_size_bytes, attempts, heartbeat_at, created_at,
|
|
started_at, finished_at, updated_at`
|
|
|
|
var claimedJobColumns = func() string {
|
|
columns := strings.Split(jobColumns, ",")
|
|
for i, column := range columns {
|
|
columns[i] = "j." + strings.TrimSpace(column)
|
|
}
|
|
return strings.Join(columns, ", ")
|
|
}()
|
|
|
|
func scanJob(row pgx.Row) (Job, error) {
|
|
var j Job
|
|
err := row.Scan(&j.ID, &j.EpisodeID, &j.SeriesID, &j.SeriesTitle, &j.Title,
|
|
&j.BroadcasterName, &j.DurationSeconds, &j.EndAt, &j.ThumbnailPath,
|
|
&j.Source, &j.Status, &j.Phase, &j.DownloadedBytes, &j.TotalBytes,
|
|
&j.TotalIsEstimate, &j.SpeedBPS, &j.ETASeconds, &j.FragmentIndex,
|
|
&j.FragmentCount, &j.FormatID, &j.ErrorKind, &j.ErrorMessage,
|
|
&j.FilePath, &j.FileSizeBytes, &j.Attempts, &j.HeartbeatAt,
|
|
&j.CreatedAt, &j.StartedAt, &j.FinishedAt, &j.UpdatedAt)
|
|
return j, err
|
|
}
|
|
|
|
// Enqueue は同じエピソードの既存ジョブを再利用する。
|
|
func (s *Store) Enqueue(ctx context.Context, ep Episode, source string) (Job, bool, error) {
|
|
row := s.pool.QueryRow(ctx, `INSERT INTO download_jobs
|
|
(episode_id, series_id, series_title, title, broadcaster_name, duration_seconds, end_at, thumbnail_path, source)
|
|
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9) ON CONFLICT (episode_id) DO NOTHING RETURNING `+jobColumns,
|
|
ep.ID, ep.SeriesID, ep.SeriesTitle, ep.Title, ep.BroadcasterName,
|
|
ep.DurationSeconds, ep.EndAt, ep.ThumbnailPath, source)
|
|
job, err := scanJob(row)
|
|
if err == nil {
|
|
return job, true, nil
|
|
}
|
|
if !errors.Is(err, pgx.ErrNoRows) {
|
|
return Job{}, false, err
|
|
}
|
|
job, err = scanJob(s.pool.QueryRow(ctx, `SELECT `+jobColumns+` FROM download_jobs WHERE episode_id=$1`, ep.ID))
|
|
return job, false, err
|
|
}
|
|
|
|
func (s *Store) GetJob(ctx context.Context, id int64) (Job, error) {
|
|
return scanJob(s.pool.QueryRow(ctx, `SELECT `+jobColumns+` FROM download_jobs WHERE id=$1`, id))
|
|
}
|
|
|
|
func (s *Store) GetSucceededJob(ctx context.Context, id int64) (Job, error) {
|
|
return scanJob(s.pool.QueryRow(ctx, `SELECT `+jobColumns+` FROM download_jobs WHERE id=$1 AND status='succeeded'`, id))
|
|
}
|
|
|
|
func (s *Store) ListSucceededJobs(ctx context.Context, limit int) ([]Job, error) {
|
|
if limit <= 0 {
|
|
return nil, errors.New("limit must be positive")
|
|
}
|
|
rows, err := s.pool.Query(ctx, `SELECT `+jobColumns+` FROM download_jobs
|
|
WHERE status='succeeded' ORDER BY finished_at DESC NULLS LAST, id DESC LIMIT $1`, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var jobs []Job
|
|
for rows.Next() {
|
|
job, err := scanJob(rows)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
jobs = append(jobs, job)
|
|
}
|
|
return jobs, rows.Err()
|
|
}
|
|
|
|
func (s *Store) DeleteSucceededJob(ctx context.Context, id int64) error {
|
|
tag, err := s.pool.Exec(ctx, `DELETE FROM download_jobs WHERE id=$1 AND status='succeeded'`, id)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|
|
|
|
// Retry は失敗・中断されたジョブの実行時情報を消して再投入する。
|
|
func (s *Store) Retry(ctx context.Context, id int64) error {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE download_jobs SET status='queued', phase=NULL,
|
|
downloaded_bytes=NULL, total_bytes=NULL, total_is_estimate=false, speed_bps=NULL,
|
|
eta_seconds=NULL, fragment_index=NULL, fragment_count=NULL, format_id=NULL,
|
|
error_kind=NULL, error_message=NULL, heartbeat_at=NULL, started_at=NULL,
|
|
finished_at=NULL, updated_at=now()
|
|
WHERE id=$1 AND status IN ('failed','canceled')`, id)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|
|
|
|
// RetryEpisode はエピソード ID で失敗・中断されたジョブを再投入する。
|
|
func (s *Store) RetryEpisode(ctx context.Context, episodeID string) error {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE download_jobs SET status='queued', phase=NULL,
|
|
downloaded_bytes=NULL, total_bytes=NULL, total_is_estimate=false, speed_bps=NULL,
|
|
eta_seconds=NULL, fragment_index=NULL, fragment_count=NULL, format_id=NULL,
|
|
error_kind=NULL, error_message=NULL, heartbeat_at=NULL, started_at=NULL,
|
|
finished_at=NULL, updated_at=now()
|
|
WHERE episode_id=$1 AND status IN ('failed','canceled')`, episodeID)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|
|
|
|
// ClaimNext は最も古い未実行ジョブを同時 worker と重複せず取得する。
|
|
func (s *Store) ClaimNext(ctx context.Context) (Job, error) {
|
|
return scanJob(s.pool.QueryRow(ctx, `WITH next AS (
|
|
SELECT id FROM download_jobs WHERE status='queued'
|
|
ORDER BY created_at, id LIMIT 1 FOR UPDATE SKIP LOCKED
|
|
) UPDATE download_jobs AS j SET status='running', attempts=j.attempts+1,
|
|
started_at=now(), heartbeat_at=now(), updated_at=now()
|
|
FROM next WHERE j.id=next.id RETURNING `+claimedJobColumns))
|
|
}
|
|
|
|
func (s *Store) UpdateProgress(ctx context.Context, id int64, p Progress) error {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE download_jobs SET phase=$2, downloaded_bytes=$3,
|
|
total_bytes=$4, total_is_estimate=$5, speed_bps=$6, eta_seconds=$7,
|
|
fragment_index=$8, fragment_count=$9, format_id=$10,
|
|
heartbeat_at=now(), updated_at=now() WHERE id=$1 AND status='running'`,
|
|
id, p.Phase, p.DownloadedBytes, p.TotalBytes, p.TotalIsEstimate,
|
|
p.SpeedBPS, p.ETASeconds, p.FragmentIndex, p.FragmentCount, p.FormatID)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|
|
|
|
func (s *Store) Complete(ctx context.Context, id int64, filePath string, fileSizeBytes int64) error {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE download_jobs SET status='succeeded', phase=NULL,
|
|
file_path=$2, file_size_bytes=$3, heartbeat_at=NULL, finished_at=now(), updated_at=now()
|
|
WHERE id=$1 AND status='running'`, id, filePath, fileSizeBytes)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|
|
|
|
func (s *Store) Fail(ctx context.Context, id int64, kind, message string) error {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE download_jobs SET status='failed', phase=NULL,
|
|
error_kind=$2, error_message=$3, heartbeat_at=NULL, finished_at=now(), updated_at=now()
|
|
WHERE id=$1 AND status='running'`, id, kind, message)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|
|
|
|
func (s *Store) Cancel(ctx context.Context, id int64) error {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE download_jobs SET status='canceled', phase=NULL,
|
|
error_kind='canceled', heartbeat_at=NULL, finished_at=now(), updated_at=now()
|
|
WHERE id=$1 AND status IN ('queued','running')`, id)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|
|
|
|
// CancelQueued は worker 未着手のジョブだけを取り消す。
|
|
func (s *Store) CancelQueued(ctx context.Context, id int64) error {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE download_jobs SET status='canceled', phase=NULL,
|
|
error_kind='canceled', heartbeat_at=NULL, finished_at=now(), updated_at=now()
|
|
WHERE id=$1 AND status='queued'`, id)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|
|
|
|
// RequeueStale は cutoff より前に heartbeat が止まった実行中ジョブを戻す。
|
|
func (s *Store) RequeueStale(ctx context.Context, cutoff time.Time) (int64, error) {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE download_jobs SET status='queued', phase=NULL,
|
|
downloaded_bytes=NULL, total_bytes=NULL, total_is_estimate=false, speed_bps=NULL,
|
|
eta_seconds=NULL, fragment_index=NULL, fragment_count=NULL, format_id=NULL,
|
|
heartbeat_at=NULL, started_at=NULL, updated_at=now()
|
|
WHERE status='running' AND heartbeat_at < $1`, cutoff)
|
|
return tag.RowsAffected(), err
|
|
}
|
|
|
|
func transitionResult(count int64, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if count == 0 {
|
|
return ErrInvalidTransition
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ListJobs は指定状態のジョブを新しい順に返す。空の状態は全件を意味する。
|
|
func (s *Store) ListJobs(ctx context.Context, status string, limit int) ([]Job, error) {
|
|
if status != "" && status != "queued" && status != "running" && status != "succeeded" && status != "failed" && status != "canceled" {
|
|
return nil, fmt.Errorf("invalid status: %s", status)
|
|
}
|
|
if limit <= 0 {
|
|
return nil, errors.New("limit must be positive")
|
|
}
|
|
rows, err := s.pool.Query(ctx, `SELECT `+jobColumns+` FROM download_jobs
|
|
WHERE ($1='' OR status=$1) ORDER BY created_at DESC, id DESC LIMIT $2`, status, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var jobs []Job
|
|
for rows.Next() {
|
|
job, err := scanJob(rows)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
jobs = append(jobs, job)
|
|
}
|
|
return jobs, rows.Err()
|
|
}
|
|
|
|
func (s *Store) EpisodeStatuses(ctx context.Context, ids []string) (map[string]string, error) {
|
|
result := make(map[string]string)
|
|
if len(ids) == 0 {
|
|
return result, nil
|
|
}
|
|
rows, err := s.pool.Query(ctx, `SELECT episode_id, status FROM download_jobs WHERE episode_id = ANY($1)`, ids)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
var id, status string
|
|
if err := rows.Scan(&id, &status); err != nil {
|
|
return nil, err
|
|
}
|
|
result[id] = status
|
|
}
|
|
return result, rows.Err()
|
|
}
|
|
|
|
func (s *Store) DeleteSubscription(ctx context.Context, seriesID string) (bool, error) {
|
|
tag, err := s.pool.Exec(ctx, `DELETE FROM subscriptions WHERE series_id=$1`, seriesID)
|
|
return tag.RowsAffected() == 1, err
|
|
}
|
|
|
|
func (s *Store) ListSubscriptions(ctx context.Context) ([]Subscription, error) {
|
|
rows, err := s.pool.Query(ctx, `SELECT s.series_id, s.series_title, s.enabled, s.created_at, s.last_checked_at,
|
|
(SELECT count(*) FROM download_jobs j WHERE j.series_id=s.series_id)
|
|
FROM subscriptions s ORDER BY s.created_at DESC, s.series_id`)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var subs []Subscription
|
|
for rows.Next() {
|
|
var sub Subscription
|
|
if err := rows.Scan(&sub.SeriesID, &sub.SeriesTitle, &sub.Enabled, &sub.CreatedAt, &sub.LastCheckedAt, &sub.JobCount); err != nil {
|
|
return nil, err
|
|
}
|
|
subs = append(subs, sub)
|
|
}
|
|
return subs, rows.Err()
|
|
}
|
|
|
|
func (s *Store) SetSubscriptionChecked(ctx context.Context, seriesID string, checkedAt time.Time) error {
|
|
tag, err := s.pool.Exec(ctx, `UPDATE subscriptions SET last_checked_at=$2 WHERE series_id=$1`, seriesID, checkedAt)
|
|
return transitionResult(tag.RowsAffected(), err)
|
|
}
|