259 lines
7.8 KiB
Python
259 lines
7.8 KiB
Python
"""Shared data models and the processing state machine."""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from dataclasses import dataclass, field, fields
|
||
from datetime import datetime
|
||
from enum import Enum
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
|
||
class State(str, Enum):
|
||
PENDING = "PENDING"
|
||
DOWNLOADING = "DOWNLOADING"
|
||
DOWNLOADED = "DOWNLOADED"
|
||
UPLOADING = "UPLOADING"
|
||
COMPLETED = "COMPLETED"
|
||
RETRY_WAIT = "RETRY_WAIT"
|
||
FAILED = "FAILED"
|
||
SKIPPED = "SKIPPED"
|
||
|
||
def __str__(self) -> str:
|
||
return self.value
|
||
|
||
|
||
# Allowed state changes. Anything else is a programming error and is rejected by the
|
||
# database layer, so a bug cannot silently mark an item COMPLETED.
|
||
TRANSITIONS: dict[State, frozenset[State]] = {
|
||
State.PENDING: frozenset(
|
||
{State.DOWNLOADING, State.DOWNLOADED, State.COMPLETED, State.SKIPPED, State.RETRY_WAIT, State.FAILED}
|
||
),
|
||
State.DOWNLOADING: frozenset({State.DOWNLOADED, State.PENDING, State.RETRY_WAIT, State.FAILED}),
|
||
State.DOWNLOADED: frozenset(
|
||
{State.UPLOADING, State.PENDING, State.COMPLETED, State.SKIPPED, State.RETRY_WAIT, State.FAILED}
|
||
),
|
||
State.UPLOADING: frozenset(
|
||
{State.COMPLETED, State.DOWNLOADED, State.PENDING, State.SKIPPED, State.RETRY_WAIT, State.FAILED}
|
||
),
|
||
State.RETRY_WAIT: frozenset({State.PENDING, State.DOWNLOADED, State.FAILED, State.SKIPPED}),
|
||
State.FAILED: frozenset({State.PENDING, State.DOWNLOADED}),
|
||
State.SKIPPED: frozenset({State.PENDING}),
|
||
State.COMPLETED: frozenset(),
|
||
}
|
||
|
||
IN_PROGRESS_STATES = frozenset({State.DOWNLOADING, State.UPLOADING})
|
||
|
||
|
||
def can_transition(current: State, new: State) -> bool:
|
||
return new in TRANSITIONS[current]
|
||
|
||
|
||
def source_key(channel_id: int, message_id: int) -> str:
|
||
"""Unique identity of a Telegram video, stored on the Drive file as appProperties."""
|
||
return f"{channel_id}:{message_id}"
|
||
|
||
|
||
@dataclass
|
||
class ParsedCaption:
|
||
raw: str
|
||
indexes: list[str] = field(default_factory=list)
|
||
title: str | None = None
|
||
topic: str | None = None
|
||
batch: str | None = None
|
||
extracted_by: str | None = None
|
||
extra: dict[str, str] = field(default_factory=dict)
|
||
warnings: list[str] = field(default_factory=list)
|
||
missing: list[str] = field(default_factory=list)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class MediaInfo:
|
||
kind: str # "video" or "document"
|
||
size: int | None
|
||
mime_type: str | None
|
||
file_name: str | None
|
||
media_id: int | None
|
||
duration: float | None = None
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SourceMessage:
|
||
channel_id: int
|
||
message_id: int
|
||
caption: str
|
||
media: MediaInfo | None
|
||
date: datetime | None = None
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class LocalFile:
|
||
path: Path
|
||
size: int
|
||
md5: str
|
||
sha256: str | None
|
||
|
||
|
||
@dataclass
|
||
class DriveFile:
|
||
id: str
|
||
name: str
|
||
size: int | None = None
|
||
md5: str | None = None
|
||
parents: tuple[str, ...] = ()
|
||
mime_type: str | None = None
|
||
trashed: bool = False
|
||
app_properties: dict[str, str] = field(default_factory=dict)
|
||
|
||
@classmethod
|
||
def from_api(cls, data: dict[str, Any]) -> "DriveFile":
|
||
size = data.get("size")
|
||
return cls(
|
||
id=data["id"],
|
||
name=data.get("name", ""),
|
||
size=int(size) if size is not None else None,
|
||
md5=data.get("md5Checksum"),
|
||
parents=tuple(data.get("parents") or ()),
|
||
mime_type=data.get("mimeType"),
|
||
trashed=bool(data.get("trashed", False)),
|
||
app_properties=dict(data.get("appProperties") or {}),
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class VideoRecord:
|
||
channel_id: int
|
||
message_id: int
|
||
state: State
|
||
caption: str | None = None
|
||
indexes: list[str] = field(default_factory=list)
|
||
title: str | None = None
|
||
topic: str | None = None
|
||
batch: str | None = None
|
||
extracted_by: str | None = None
|
||
warnings: list[str] = field(default_factory=list)
|
||
media_kind: str | None = None
|
||
media_id: int | None = None
|
||
mime_type: str | None = None
|
||
original_filename: str | None = None
|
||
expected_size: int | None = None
|
||
message_date: str | None = None
|
||
sequence_no: int | None = None
|
||
local_path: str | None = None
|
||
local_size: int | None = None
|
||
sha256: str | None = None
|
||
md5: str | None = None
|
||
final_name: str | None = None
|
||
drive_folder_id: str | None = None
|
||
drive_file_id: str | None = None
|
||
upload_session_uri: str | None = None
|
||
upload_session_created_at: str | None = None
|
||
download_status: str | None = None
|
||
upload_status: str | None = None
|
||
attempts: int = 0
|
||
last_error: str | None = None
|
||
error_kind: str | None = None
|
||
next_retry_at: str | None = None
|
||
skip_reason: str | None = None
|
||
discovered_at: str | None = None
|
||
updated_at: str | None = None
|
||
download_started_at: str | None = None
|
||
downloaded_at: str | None = None
|
||
upload_started_at: str | None = None
|
||
completed_at: str | None = None
|
||
task_id: int | None = None
|
||
|
||
JSON_FIELDS = ("indexes", "warnings")
|
||
|
||
@property
|
||
def key(self) -> tuple[int, int]:
|
||
return (self.channel_id, self.message_id)
|
||
|
||
@property
|
||
def source_key(self) -> str:
|
||
return source_key(self.channel_id, self.message_id)
|
||
|
||
@classmethod
|
||
def column_names(cls) -> list[str]:
|
||
return [f.name for f in fields(cls)]
|
||
|
||
@classmethod
|
||
def from_row(cls, row: Any) -> "VideoRecord":
|
||
data = {name: row[name] for name in cls.column_names()}
|
||
for name in cls.JSON_FIELDS:
|
||
data[name] = json.loads(data[name]) if data[name] else []
|
||
data["state"] = State(data["state"])
|
||
return cls(**data)
|
||
|
||
|
||
# -- interactive /upload tasks ---------------------------------------------------------
|
||
|
||
|
||
class TaskStatus(str, Enum):
|
||
AWAITING_START = "AWAITING_START"
|
||
AWAITING_END = "AWAITING_END"
|
||
AWAITING_CONFIRMATION = "AWAITING_CONFIRMATION"
|
||
QUEUED = "QUEUED"
|
||
RUNNING = "RUNNING"
|
||
PAUSED = "PAUSED"
|
||
COMPLETED = "COMPLETED"
|
||
CANCELLED = "CANCELLED"
|
||
FAILED = "FAILED"
|
||
|
||
def __str__(self) -> str:
|
||
return self.value
|
||
|
||
|
||
SETUP_TASK_STATUSES = frozenset({TaskStatus.AWAITING_START, TaskStatus.AWAITING_END, TaskStatus.AWAITING_CONFIRMATION})
|
||
ACTIVE_TASK_STATUSES = frozenset({TaskStatus.QUEUED, TaskStatus.RUNNING, TaskStatus.PAUSED})
|
||
FINAL_TASK_STATUSES = frozenset({TaskStatus.COMPLETED, TaskStatus.CANCELLED, TaskStatus.FAILED})
|
||
|
||
# Per-message outcome inside a task (task_items.result).
|
||
TASK_ITEM_RESULTS = ("pending", "uploaded", "duplicate", "skipped", "non_video", "missing", "failed", "cancelled")
|
||
|
||
|
||
@dataclass
|
||
class UploadTask:
|
||
id: int
|
||
admin_user_id: int
|
||
status: TaskStatus
|
||
channel_ref: str | None = None
|
||
channel_id: int | None = None
|
||
channel_title: str | None = None
|
||
start_message_id: int | None = None
|
||
end_message_id: int | None = None
|
||
created_at: str | None = None
|
||
updated_at: str | None = None
|
||
confirmed_at: str | None = None
|
||
started_at: str | None = None
|
||
scanned_at: str | None = None
|
||
completed_at: str | None = None
|
||
current_message_id: int | None = None
|
||
examined_count: int = 0
|
||
eligible_count: int = 0
|
||
uploaded_count: int = 0
|
||
duplicate_count: int = 0
|
||
non_video_count: int = 0
|
||
skipped_count: int = 0
|
||
missing_count: int = 0
|
||
failed_count: int = 0
|
||
deleted_files_count: int = 0
|
||
last_error: str | None = None
|
||
|
||
@classmethod
|
||
def column_names(cls) -> list[str]:
|
||
return [f.name for f in fields(cls)]
|
||
|
||
@classmethod
|
||
def from_row(cls, row: Any) -> "UploadTask":
|
||
data = {name: row[name] for name in cls.column_names()}
|
||
data["status"] = TaskStatus(data["status"])
|
||
return cls(**data)
|
||
|
||
@property
|
||
def range_text(self) -> str:
|
||
if self.start_message_id is None:
|
||
return "-"
|
||
end = self.end_message_id if self.end_message_id is not None else "?"
|
||
return f"{self.start_message_id}–{end}"
|