"""Discovery and the strictly sequential processing pipeline.
Exactly one worker task exists per process, and InstanceLock guarantees one
process, so at most one video is downloading or uploading at any time.
Per item: activate -> reconcile with Drive -> resolve folder -> download ->
choose name -> upload (resumable) -> verify -> COMPLETED -> clean up.
Every step is idempotent and driven by the persisted record, so the same code
path handles first runs, retries and recovery after a crash or dyno restart.
"""
from __future__ import annotations
import asyncio
import logging
import random
import threading
import time
from dataclasses import dataclass
from datetime import datetime, timedelta
from caption_parser import parse_caption
from database import iso, retry_time, utcnow
from drive_client import SOURCE_PROPERTY
from errors import (
BotError,
Interrupted,
PermanentError,
QueueHaltError,
SkipItem,
TransientError,
UploadSessionExpired,
)
from models import (
FINAL_TASK_STATUSES,
SETUP_TASK_STATUSES,
DriveFile,
LocalFile,
SourceMessage,
State,
TaskStatus,
UploadTask,
VideoRecord,
)
from naming import add_suffix, build_filename, clean_component, detect_video_extension, mime_for_name, normalize_key, truncate_name
from progress import HtmlReply, ProgressReporter, esc, format_duration
log = logging.getLogger(__name__)
CURSOR_KEY = "last_discovered_message_id"
COMPLETED_CURSOR_KEY = "last_completed_message_id"
GLOBAL_RUN_KEY = "queue_global_run" # "1" after /run, "0" after /stop
PAUSED_KEY = "queue_paused" # "1" after /pause, "0" after /resume
CANCELLED_PREFIX = "cancelled:"
SESSION_MAX_AGE = timedelta(days=6) # Drive resumable sessions expire after about a week
MAX_DB_FAILURES = 5
UNFINISHED = ("PENDING", "RETRY_WAIT", "DOWNLOADING", "DOWNLOADED", "UPLOADING")
@dataclass
class RunStats:
uploaded: int = 0
reused: int = 0
skipped: int = 0
failed: int = 0
retries: int = 0
@dataclass
class SyncResult:
new: int = 0
known: int = 0
ignored: int = 0
skipped: int = 0
note: str | None = None
def add(self, outcome: str) -> None:
setattr(self, outcome, getattr(self, outcome) + 1)
def summary(self) -> str:
text = (
f"New videos queued: {self.new}\nAlready known: {self.known}\n"
f"Skipped by filters/duplicates: {self.skipped}\nNon-video messages ignored: {self.ignored}"
)
return f"{text}\nNote: {self.note}" if self.note else text
def public_error(err: BaseException) -> str:
text = getattr(err, "public", None) or type(err).__name__
return text if len(text) <= 300 else text[:297] + "..."
class QueueManager:
def __init__(
self,
config,
db,
drive,
downloader,
folders,
*,
source=None,
reporter: ProgressReporter | None = None,
db_error_delay: float = 30.0,
):
self.config = config
self.db = db
self.drive = drive
self.downloader = downloader
self.folders = folders
self.source = source
self.reporter = reporter
self.db_error_delay = db_error_delay
self.listening = False # True while live channel updates are being received
self.halt_reason: str | None = None
self.stats = RunStats()
self.current: VideoRecord | None = None
self.active_uploads = 0
self.max_concurrent_uploads = 0
self._task: asyncio.Task | None = None
self._paused = False
self._stop_requested = False
self._resume_event = asyncio.Event()
self._wake = asyncio.Event()
self._stop_event = asyncio.Event()
self._interrupt = threading.Event() # read from the upload thread
self._sync_lock = asyncio.Lock()
self._item_lock = asyncio.Lock() # held while one item is processed
self._db_failures = 0
self._last_cleanup = time.monotonic()
downloader.reclaim = lambda: downloader.cleanup_orphans(db, keep_failed=config.keep_failed_downloads)
# -- discovery ------------------------------------------------------------------
async def sync(self, *, rescan: bool = False) -> SyncResult:
"""Discover channel messages and queue unseen videos. Safe to repeat."""
if self.source is None:
raise RuntimeError("no Telegram source configured")
if self._sync_lock.locked():
return SyncResult(note="a sync is already running")
async with self._sync_lock:
cursor = self.db.get_int(CURSOR_KEY, 0)
if self.config.process_mode == "new" and cursor == 0 and not rescan:
return SyncResult(note="PROCESS_MODE=new and no cursor yet; only new posts will be queued")
start = self.config.start_message_id if rescan else max(self.config.start_message_id, cursor + 1)
result = SyncResult()
async for msg in self.source.iter_history(min_id=start, start_date=self.config.start_date):
result.add(self.register(msg))
self._wake.set()
log.info("Sync finished: %s", result.summary().replace("\n", "; "))
return result
def register(self, msg: SourceMessage, *, task_id: int | None = None) -> str:
"""Record one discovered message. Returns new / known / skipped / ignored."""
if msg.media is None:
self.db.raise_int(CURSOR_KEY, msg.message_id)
return "ignored"
if self.db.exists(msg.channel_id, msg.message_id):
return "known"
parsed = parse_caption(msg.caption)
reason = self._filter_reason(parsed)
if reason is None and self.config.duplicate_media_policy == "skip" and msg.media.media_id is not None:
for other in self.db.find_by_media_id(msg.media.media_id):
if other.message_id != msg.message_id and other.state != State.SKIPPED:
reason = f"duplicate of message {other.message_id} (same Telegram media)"
break
state = State.SKIPPED if reason else State.PENDING
sequence = None if reason else self.db.next_sequence()
inserted = self.db.insert_video(
msg, parsed, state=state, skip_reason=reason, sequence_no=sequence, task_id=task_id
)
self.db.raise_int(CURSOR_KEY, msg.message_id)
if not inserted:
return "known"
if parsed.missing or parsed.warnings:
log.warning("Message %s caption issues: %s", msg.message_id, "; ".join(parsed.warnings))
self._wake.set()
return "skipped" if reason else "new"
def _filter_reason(self, parsed) -> str | None:
if self.config.filter_topics and normalize_key(parsed.topic or "") not in self.config.filter_topics:
return "filtered: topic not in FILTER_TOPICS"
if self.config.filter_batches and normalize_key(parsed.batch or "") not in self.config.filter_batches:
return "filtered: batch not in FILTER_BATCHES"
return None
# -- control ----------------------------------------------------------------------
#
# Two persisted switches (kv table) describe what the operator asked for:
# queue_global_run = 1 -> /run: process every discovered video
# queue_paused = 1 -> /pause: nothing runs until /resume
# /upload tasks run whenever the queue is not paused, even without /run.
@property
def running(self) -> bool:
return self._task is not None and not self._task.done()
@property
def status(self) -> str:
if self.running:
if self._stop_requested:
return "stopping"
return "paused" if self._paused else "running"
return "halted" if self.halt_reason else "idle"
@property
def global_enabled(self) -> bool:
return self.db.get_kv(GLOBAL_RUN_KEY, "0") == "1"
@property
def paused_flag(self) -> bool:
return self.db.get_kv(PAUSED_KEY, "0") == "1"
def desired_state(self) -> str:
if self.paused_flag:
return "paused"
return "running" if self.global_enabled else "stopped"
def _persist_desired(self, value: str) -> None:
"""Back-compat helper: 'running' | 'paused' | 'stopped'."""
if value == "paused":
self.db.set_kv(PAUSED_KEY, "1")
else:
self.db.set_kv(PAUSED_KEY, "0")
self.db.set_kv(GLOBAL_RUN_KEY, "1" if value == "running" else "0")
def start(self) -> bool:
"""/run: process the whole queue. Returns False if the worker was already running."""
self._persist_desired("running")
return self.ensure_worker()
def ensure_worker(self) -> bool:
"""Start the worker if needed (used by /run, /upload tasks and startup)."""
if self.running:
if self._paused:
return self.resume()
self._wake.set()
return False
self.halt_reason = None
self._paused = False
self._stop_requested = False
self._interrupt.clear()
self._stop_event.clear()
self._resume_event.set()
self._task = asyncio.create_task(self._worker(), name="queue-worker")
return True
def has_runnable_work(self) -> bool:
return self.global_enabled or self.db.next_runnable_task() is not None
def boot(self) -> str:
"""Decide what to do after a (re)start, from durable state only."""
if self.paused_flag:
return "paused"
if self.has_runnable_work():
self.ensure_worker()
return "running"
return "idle"
def pause(self) -> bool:
if not self.running or self._paused:
return False
self.db.set_kv(PAUSED_KEY, "1")
self.db.set_task_status(TaskStatus.RUNNING, TaskStatus.PAUSED)
self._paused = True
self._resume_event.clear()
self._interrupt.set()
self._stop_event.set()
return True
def resume(self) -> bool:
self.db.set_kv(PAUSED_KEY, "0")
resumed_tasks = self.db.set_task_status(TaskStatus.PAUSED, TaskStatus.RUNNING)
if not self.running:
if self.has_runnable_work():
return self.ensure_worker()
return False
if not self._paused:
self._wake.set()
return resumed_tasks > 0
self._paused = False
self._interrupt.clear()
self._stop_event.clear()
self._resume_event.set()
return True
async def stop(self, *, persist: bool = True, timeout: float | None = None) -> None:
"""Stop at the next resumable checkpoint and end the worker.
With ``persist`` (the /stop command) the whole-queue run is switched off and a
running /upload task is parked as PAUSED. Without it (SIGTERM / dyno restart)
nothing durable changes, so processing resumes automatically after restart.
"""
if persist:
self._persist_desired("stopped")
self.db.set_task_status(TaskStatus.RUNNING, TaskStatus.PAUSED)
if not self.running:
return
self._stop_requested = True
self._interrupt.set()
self._stop_event.set()
self._resume_event.set()
self._wake.set()
try:
await asyncio.wait_for(asyncio.shield(self._task), timeout)
except asyncio.TimeoutError:
log.warning("Worker did not reach a checkpoint within %ss; its state remains resumable", timeout)
async def shutdown(self, timeout: float = 20.0) -> None:
"""Process termination (SIGTERM): checkpoint without changing the operator's choice."""
await self.stop(persist=False, timeout=timeout)
async def wait(self) -> None:
if self._task:
await self._task
def wake(self) -> None:
self._wake.set()
def retry_failed(self) -> int:
count = self.db.reset_failed(has_local=self.downloader.has_complete_local)
self._wake.set()
return count
def remaining(self) -> int:
counts = self.db.counts()
return sum(counts[s] for s in UNFINISHED)
# -- /upload tasks --------------------------------------------------------------------
def submit_task(self, task_id: int) -> UploadTask:
task = self.db.update_task(task_id, status=TaskStatus.QUEUED, confirmed_at=iso(utcnow()))
log.info("Upload task #%s queued (range %s, channel %s)", task.id, task.range_text, task.channel_id)
if not self.paused_flag:
self.ensure_worker()
self._wake.set()
return task
async def cancel_task(self, task_id: int, *, reason: str = "cancelled by administrator") -> str:
"""Cancel a setup or an active task. Completed uploads are always kept.
Returns one of: unknown, finished, setup_cancelled, cancelled.
"""
task = self.db.get_task(task_id)
if task is None:
return "unknown"
if task.status in SETUP_TASK_STATUSES:
self.db.update_task(task_id, status=TaskStatus.CANCELLED, completed_at=iso(utcnow()), last_error=reason)
return "setup_cancelled"
if task.status in FINAL_TASK_STATUSES:
return "finished"
if self.current is not None and self.current.task_id == task_id:
self._interrupt.set() # current transfer stops at its next checkpoint
async with self._item_lock:
await self._finalize_task(self.db.get_task(task_id), TaskStatus.CANCELLED, reason=reason)
if not self._paused and not self._stop_requested:
self._interrupt.clear()
self._wake.set()
return "cancelled"
async def _scan_task(self, task: UploadTask) -> None:
"""Examine every message ID of the task's inclusive range (idempotent)."""
self.db.update_task(task.id, status=TaskStatus.RUNNING, started_at=task.started_at or iso(utcnow()), last_error=None)
async for message_id, msg in self.source.iter_range(task.channel_id, task.start_message_id, task.end_message_id):
self._attach(task, message_id, msg)
task = self.db.update_task(task.id, scanned_at=iso(utcnow()))
counts = self._sync_task_counts(task.id)
log.info("Upload task #%s scanned: %s", task.id, counts)
if self.reporter:
self.reporter.new_message()
await self.reporter.notify(
HtmlReply(
f"🚀 Upload Task #{task.id} started\nRange: {task.range_text}\n"
f"Eligible videos: {counts['eligible_count']} "
f"(already uploaded: {counts['duplicate_count']}, non-video messages: {counts['non_video_count']})"
),
extra_chat_ids=(task.admin_user_id,),
)
def _attach(self, task: UploadTask, message_id: int, msg: SourceMessage | None) -> str:
key = (task.channel_id, message_id)
if msg is None:
result = "missing"
elif msg.media is None:
result = "non_video"
else:
existing = self.db.get(key)
if existing is None:
outcome = self.register(msg, task_id=task.id)
if outcome in ("new", "known"):
result = "pending"
else:
rec = self.db.get(key)
result = "duplicate" if (rec.skip_reason or "").startswith("duplicate") else "skipped"
else:
result = self._claim(existing, task)
self.db.add_task_item(task.id, task.channel_id, message_id, result)
return result
def _claim(self, rec: VideoRecord, task: UploadTask) -> str:
"""Attach an already-known video to ``task``; never re-uploads completed work."""
note = f"claimed by upload task #{task.id}"
if rec.state == State.COMPLETED:
return "duplicate"
if rec.state == State.SKIPPED:
reason = rec.skip_reason or ""
if reason.startswith(CANCELLED_PREFIX):
self.db.transition(rec.key, State.PENDING, note=note, task_id=task.id, skip_reason=None)
return "pending"
return "duplicate" if reason.startswith("duplicate") else "skipped"
if rec.state == State.FAILED:
target = State.DOWNLOADED if self.downloader.has_complete_local(rec) else State.PENDING
self.db.transition(rec.key, target, note=note, task_id=task.id, attempts=0, next_retry_at=None)
return "pending"
self.db.update_fields(rec.key, task_id=task.id)
return "pending"
def _sync_task_counts(self, task_id: int) -> dict[str, int]:
r = self.db.task_item_results(task_id)
total = sum(r.values())
counts = {
"examined_count": total - r.get("missing", 0),
"eligible_count": sum(r.get(k, 0) for k in ("pending", "uploaded", "duplicate", "skipped", "failed", "cancelled")),
"uploaded_count": r.get("uploaded", 0),
"duplicate_count": r.get("duplicate", 0),
"non_video_count": r.get("non_video", 0),
"skipped_count": r.get("skipped", 0) + r.get("cancelled", 0),
"missing_count": r.get("missing", 0),
"failed_count": r.get("failed", 0),
}
self.db.update_task(task_id, **counts)
if self.reporter and self.reporter.board.task and self.reporter.board.task.get("id") == task_id:
self.reporter.board.task.update(
completed=counts["uploaded_count"],
skipped=counts["duplicate_count"] + counts["non_video_count"] + counts["skipped_count"] + counts["missing_count"],
failed=counts["failed_count"],
)
return counts
async def _finalize_task(self, task: UploadTask, status: TaskStatus, *, reason: str | None = None) -> None:
if status == TaskStatus.CANCELLED:
for rec in self.db.videos_for_task(task.id):
if rec.state in (State.COMPLETED, State.SKIPPED, State.FAILED):
continue
if rec.state == State.DOWNLOADING:
rec = self.db.transition(rec.key, State.PENDING, note="checkpoint before cancel")
skip = f"{CANCELLED_PREFIX} task #{task.id}"
self.db.transition(rec.key, State.SKIPPED, note=skip, skip_reason=skip, upload_session_uri=None)
self.db.set_task_item_result(task.id, rec.channel_id, rec.message_id, "cancelled")
deleted = self.downloader.cleanup(rec)
if deleted:
self.db.add_task_counter(task.id, "deleted_files_count", deleted)
counts = self._sync_task_counts(task.id)
task = self.db.update_task(
task.id, status=status, completed_at=iso(utcnow()), current_message_id=None,
last_error=reason if status != TaskStatus.COMPLETED else task.last_error,
)
log.info("Upload task #%s %s: %s", task.id, status.value, counts)
if self.reporter:
if self.reporter.board.task and self.reporter.board.task.get("id") == task.id:
self.reporter.board.task = None
await self.reporter.notify(self.task_summary(task), extra_chat_ids=(task.admin_user_id,))
self.reporter.new_message()
def task_summary(self, task: UploadTask) -> HtmlReply:
title = {
TaskStatus.COMPLETED: "✅ Upload Task Completed",
TaskStatus.CANCELLED: "🛑 Upload Task Cancelled",
TaskStatus.FAILED: "❌ Upload Task Failed",
}.get(task.status, f"Upload Task {esc(task.status.value)}")
folders = self.db.task_folders(task.id)
lines = [
title,
"",
f"Task: #{task.id}",
f"Requested range: {task.range_text}",
f"Messages examined: {task.examined_count}" + (f" (+{task.missing_count} missing/inaccessible)" if task.missing_count else ""),
f"Eligible videos found: {task.eligible_count}",
f"Uploaded: {task.uploaded_count}",
f"Duplicates skipped: {task.duplicate_count}",
f"Non-video messages skipped: {task.non_video_count}",
f"Other skipped: {task.skipped_count}",
f"Failed: {task.failed_count}",
f"Local temporary files deleted: {task.deleted_files_count}",
"Drive folders used: " + (", ".join(esc(f) for f in folders) if folders else "-"),
]
if task.last_error and task.status != TaskStatus.COMPLETED:
lines.append(f"Reason: {esc(task.last_error)}")
return HtmlReply("\n".join(lines))
async def _run_task_step(self, task: UploadTask) -> bool:
"""Advance ``task`` by one step. Returns False when it is waiting for a retry."""
if task.status == TaskStatus.QUEUED or task.scanned_at is None:
retry_in = None
async with self._item_lock: # a concurrent /cancel waits for the scan to finish
task = self.db.get_task(task.id)
if task.status not in (TaskStatus.QUEUED, TaskStatus.RUNNING):
return True
try:
await self._scan_task(task)
except PermanentError as exc:
await self._finalize_task(task, TaskStatus.FAILED, reason=public_error(exc))
except TransientError as exc:
self.db.update_task(task.id, last_error=public_error(exc))
log.warning("Upload task #%s scan failed (%s); will retry", task.id, exc)
retry_in = exc.retry_after or max(self.config.retry_base_delay, 5.0)
if retry_in:
await self._sleep(retry_in)
return True
rec = self.db.next_item(task_id=task.id)
if rec is None:
if self.db.task_unfinished_count(task.id) == 0:
await self._finalize_task(task, TaskStatus.COMPLETED)
return True
return False
if self.reporter:
board = self.reporter.board
if not board.task or board.task.get("id") != task.id:
board.task = {"id": task.id, "start": task.start_message_id, "end": task.end_message_id}
self._sync_task_counts(task.id)
board.task["current"] = rec.message_id
self.db.update_task(task.id, current_message_id=rec.message_id)
await self._process_locked(rec)
self._sync_task_counts(task.id)
return True
async def _process_locked(self, rec: VideoRecord) -> None:
async with self._item_lock:
fresh = self.db.get(rec.key)
# A cancellation may have finalized this item while we waited for the lock.
if fresh is None or fresh.state in (State.COMPLETED, State.SKIPPED, State.FAILED):
return
await self.process_item(fresh)
if self.config.processing_delay and not self._stop_requested and not self._paused:
await self._sleep(self.config.processing_delay)
def _maybe_cleanup_orphans(self) -> None:
interval = self.config.temp_cleanup_interval
now = time.monotonic()
if interval and now - self._last_cleanup >= interval:
self._last_cleanup = now
deleted = self.downloader.cleanup_orphans(self.db, keep_failed=self.config.keep_failed_downloads)
if deleted:
log.info("Periodic cleanup removed %d abandoned temporary file(s)", deleted)
# -- worker -------------------------------------------------------------------------
async def _worker(self) -> None:
self.stats = RunStats()
if self.reporter:
await self.reporter.start()
log.info("Queue worker started")
try:
while not self._stop_requested:
if self._paused:
if self.reporter:
self.reporter.board.status = "Paused"
await self._resume_event.wait()
continue
self._wake.clear()
try:
self._maybe_cleanup_orphans()
task = self.db.next_runnable_task()
if task is not None:
if not await self._run_task_step(task):
await self._wait_for_work(force=True)
continue
if not self.global_enabled:
break
rec = self.db.next_item()
except self.db.errors as exc:
await self._on_db_error(exc)
continue
if rec is None:
if not await self._wait_for_work():
break
continue
await self._process_locked(rec)
finally:
log.info("Queue worker stopped (%s)", self.halt_reason or ("stop requested" if self._stop_requested else "queue empty"))
if self.reporter:
self.reporter.board.status = "Idle"
self.reporter.board.task = None
await self.reporter.stop()
await self.reporter.notify(self._summary())
async def _wait_for_work(self, *, force: bool = False) -> bool:
next_retry = self.db.earliest_retry_at()
if next_retry is None and not self.listening and not force:
return False
timeout = 60.0 if next_retry is None else max(1.0, (next_retry - utcnow()).total_seconds())
if self.reporter:
self.reporter.board.status = "Waiting"
try:
await asyncio.wait_for(self._wake.wait(), min(timeout, 300.0))
except asyncio.TimeoutError:
pass
return True
async def _sleep(self, seconds: float) -> None:
try:
await asyncio.wait_for(self._stop_event.wait(), seconds)
except asyncio.TimeoutError:
pass
def _summary(self) -> str:
s = self.stats
lines = [
"Run finished" if not self.halt_reason else f"Processing halted: {self.halt_reason}",
f"Uploaded: {s.uploaded}",
f"Already in Drive (duplicates reused): {s.reused}",
f"Skipped: {s.skipped}",
f"Failed: {s.failed}",
f"Remaining: {self.remaining()}",
]
return "\n".join(lines)
async def _notify(self, text: str) -> None:
if self.reporter:
await self.reporter.notify(text)
async def _halt(self, reason: str) -> None:
self.halt_reason = reason
self._stop_requested = True
log.error("Queue halted: %s", reason)
async def _on_db_error(self, exc: BaseException) -> None:
self._db_failures += 1
log.error("Database error (%d/%d): %s", self._db_failures, MAX_DB_FAILURES, exc)
if self._db_failures >= MAX_DB_FAILURES:
await self._halt("database unavailable; check DATABASE_URL / database health, then /resume")
else:
await self._sleep(self.db_error_delay)
# -- per item -------------------------------------------------------------------------
def folder_path(self, rec: VideoRecord) -> list[str]:
strict = self.config.strict_filenames
topic = clean_component(rec.topic, strict=strict) or self.config.fallback_topic
if self.config.folder_structure == "batch_topic":
return [clean_component(rec.batch, strict=strict) or self.config.fallback_batch, topic]
return [topic]
async def process_item(self, rec: VideoRecord) -> None:
key = rec.key
self.current = rec
if self.reporter:
counts = self.db.counts()
total = sum(counts.values()) - counts["SKIPPED"]
self.reporter.board.begin_item(
counts["COMPLETED"] + 1, total, " / ".join(self.folder_path(rec)), rec.title or f"(message {rec.message_id})"
)
try:
try:
await self._process(rec)
self._db_failures = 0
except Interrupted as exc:
log.info("Message %s: %s", rec.message_id, exc)
except SkipItem as exc:
self._skip(key, exc.reason)
except QueueHaltError as exc:
self._park(key, exc)
await self._notify(f"Processing halted on message {rec.message_id}: {public_error(exc)}")
await self._halt(public_error(exc))
except PermanentError as exc:
await self._fail(key, exc)
except TransientError as exc:
await self._retry_or_fail(key, exc)
except self.db.errors:
raise
except Exception as exc: # unexpected bug: record it and keep the queue moving
log.exception("Unexpected error processing message %s", rec.message_id)
wrapped = TransientError(f"{type(exc).__name__}: {exc}", public=f"Internal error ({type(exc).__name__}); see logs")
await self._retry_or_fail(key, wrapped)
except self.db.errors as exc:
await self._on_db_error(exc)
finally:
self.current = None
async def _process(self, rec: VideoRecord) -> None:
key = rec.key
rec = self._activate(rec)
board = self.reporter.board if self.reporter else None
# A previous attempt may have uploaded the file and crashed before saving.
if rec.upload_started_at or rec.drive_file_id:
if board:
board.status = "Checking Google Drive"
existing = await asyncio.to_thread(self._find_existing_upload, rec, None)
if existing:
await self._complete(rec, existing, None, reused=True)
return
if board:
board.status = "Resolving folder"
folder_id = await asyncio.to_thread(self.folders.resolve, self.folder_path(rec))
if rec.state == State.PENDING:
local = await self._download(rec)
else:
local = await self.downloader.local_file(rec)
rec = self.db.get(key)
drive_file, reused = await self._upload(rec, local, folder_id)
await self._complete(self.db.get(key), drive_file, local, reused=reused)
def _activate(self, rec: VideoRecord) -> VideoRecord:
has_local = self.downloader.has_complete_local(rec)
if rec.state == State.RETRY_WAIT:
return self.db.transition(rec.key, State.DOWNLOADED if has_local else State.PENDING, note="retry")
if rec.state == State.DOWNLOADING:
return self.db.transition(rec.key, State.PENDING, note="resume interrupted download")
if rec.state in (State.DOWNLOADED, State.UPLOADING) and not has_local:
# e.g. a Heroku dyno restart wiped the ephemeral disk; the upload
# session URI is kept and reused if the re-download is byte-identical.
return self.db.transition(rec.key, State.PENDING, note="local file missing; downloading again")
return rec
async def _download(self, rec: VideoRecord) -> LocalFile:
key = rec.key
rec = self.db.transition(key, State.DOWNLOADING, download_started_at=iso(utcnow()))
board = self.reporter.board if self.reporter else None
if board:
board.start_download(rec.expected_size)
def on_progress(done: int, total: int | None) -> None:
if board and board.download:
board.download.update(done, total)
try:
local = await self.downloader.download(rec, on_progress=on_progress, should_stop=self._interrupt.is_set)
except Interrupted:
self.db.transition(key, State.PENDING, note="checkpoint during download")
raise
fields = dict(
local_path=str(local.path), local_size=local.size, md5=local.md5, sha256=local.sha256, downloaded_at=iso(utcnow())
)
if rec.upload_session_uri and rec.md5 and rec.md5 != local.md5:
log.warning("Message %s: re-downloaded file differs; discarding old upload session", rec.message_id)
fields["upload_session_uri"] = None
self.db.transition(key, State.DOWNLOADED, **fields)
return local
def _find_existing_upload(self, rec: VideoRecord, local: LocalFile | None) -> DriveFile | None:
"""Find a verified Drive file previously uploaded for this exact message."""
candidates: list[DriveFile] = []
if rec.drive_file_id:
known = self.drive.get_file(rec.drive_file_id)
if known:
candidates.append(known)
candidates += self.drive.find_by_source_key(rec.source_key)
size = local.size if local else (rec.local_size or rec.expected_size)
md5 = local.md5 if local else rec.md5
for f in candidates:
if f.trashed or (size is not None and f.size != size):
continue
if md5 and f.md5 and f.md5 != md5:
continue
log.info("Message %s already uploaded as Drive file %s; reusing it", rec.message_id, f.id)
return f
if candidates:
log.warning("Message %s has Drive files that failed verification; uploading again", rec.message_id)
return None
def _choose_name(self, rec: VideoRecord, local: LocalFile, folder_id: str) -> tuple[str, DriveFile | None]:
cfg = self.config
name, warnings = build_filename(
rec.title,
message_id=rec.message_id,
original_filename=rec.original_filename,
mime_type=rec.mime_type,
detected_ext=detect_video_extension(local.path),
missing_title_policy=cfg.missing_title_policy,
mismatch_policy=cfg.extension_mismatch_policy,
strict=cfg.strict_filenames,
)
if cfg.sequence_prefix_enabled and rec.sequence_no:
name = truncate_name(cfg.sequence_prefix_format.format(seq=rec.sequence_no) + name)
new_warnings = [w for w in warnings if w not in rec.warnings]
if new_warnings:
self.db.update_fields(rec.key, warnings=rec.warnings + new_warnings)
candidate, n = name, 1
while True:
clashes = self.drive.find_by_name(folder_id, candidate)
if not clashes:
return candidate, None
for f in clashes:
if f.app_properties.get(SOURCE_PROPERTY) == rec.source_key and f.size == local.size:
return candidate, f
if cfg.reuse_identical_files and f.size == local.size and f.md5 and f.md5 == local.md5:
log.info("Identical file %r (same size and MD5) already in folder; reusing it", candidate)
return candidate, f
policy = cfg.duplicate_name_policy
if policy == "skip":
raise SkipItem(f"a different file named {candidate!r} already exists in the folder")
if policy == "fail":
raise PermanentError(
f"name collision: {candidate!r}", public=f"Name collision: {candidate!r} already exists (DUPLICATE_POLICY=fail)"
)
n += 1
if n > 1000:
raise PermanentError(f"could not find a free name for {name!r}")
if policy == "suffix_message_id":
suffix = f"[{rec.message_id}]" if n == 2 else f"[{rec.message_id}] ({n - 1})"
else:
suffix = f"({n})"
candidate = add_suffix(name, suffix)
async def _upload(self, rec: VideoRecord, local: LocalFile, folder_id: str) -> tuple[DriveFile, bool]:
key = rec.key
session = rec.upload_session_uri
if session and rec.upload_session_created_at:
if utcnow() - datetime.fromisoformat(rec.upload_session_created_at) > SESSION_MAX_AGE:
session = None
if session and rec.drive_folder_id and rec.drive_folder_id != folder_id:
session = None
if session:
name = rec.final_name
else:
existing = await asyncio.to_thread(self._find_existing_upload, rec, local)
if existing:
return existing, True
name, reuse = await asyncio.to_thread(self._choose_name, rec, local, folder_id)
if reuse:
return reuse, True
fields = dict(final_name=name, drive_folder_id=folder_id, upload_session_uri=session)
if rec.state != State.UPLOADING:
# upload_started_at is written *before* the upload begins, so a crash at
# any later point triggers the Drive reconciliation on restart.
self.db.transition(key, State.UPLOADING, upload_started_at=rec.upload_started_at or iso(utcnow()), **fields)
else:
self.db.update_fields(key, **fields)
board = self.reporter.board if self.reporter else None
if board:
board.title = name
board.start_upload(local.size)
def on_session(uri: str) -> None:
self.db.update_fields(key, upload_session_uri=uri, upload_session_created_at=iso(utcnow()))
def on_progress(done: int, total: int) -> None:
if board and board.upload:
board.upload.update(done, total)
self.active_uploads += 1
self.max_concurrent_uploads = max(self.max_concurrent_uploads, self.active_uploads)
try:
uploaded = await asyncio.to_thread(
self.drive.upload_file,
local.path,
name=name,
parent_id=folder_id,
mime_type=mime_for_name(name, rec.mime_type),
size=local.size,
app_properties={SOURCE_PROPERTY: rec.source_key},
description=self._description(rec),
session_uri=session,
on_session=on_session,
on_progress=on_progress,
should_stop=self._interrupt.is_set,
)
except UploadSessionExpired:
self.db.update_fields(key, upload_session_uri=None)
raise
finally:
self.active_uploads -= 1
if board:
board.status = "Verifying"
try:
verified = await asyncio.to_thread(self._verify, uploaded, local)
except TransientError:
self.db.update_fields(key, upload_session_uri=None)
raise
return verified, False
def _verify(self, uploaded: DriveFile, local: LocalFile) -> DriveFile:
found = self.drive.get_file(uploaded.id)
if found is None or found.trashed:
raise TransientError(f"uploaded file {uploaded.id} not found during verification", public="Upload verification failed")
if found.size != local.size:
raise TransientError(
f"size mismatch after upload: Drive {found.size} vs local {local.size}", public="Upload verification failed (size)"
)
if found.md5 and found.md5 != local.md5:
raise TransientError("MD5 mismatch after upload", public="Upload verification failed (checksum)")
return found
def _description(self, rec: VideoRecord) -> str:
parts = [f"Telegram message {rec.message_id}"]
if rec.batch:
parts.append(f"Batch: {rec.batch}")
if rec.indexes:
parts.append(f"Index: {', '.join(rec.indexes)}")
return " | ".join(parts)
async def _complete(self, rec: VideoRecord, drive_file: DriveFile, local: LocalFile | None, *, reused: bool) -> None:
if self.config.share_mode == "anyone_with_link" and not reused:
await asyncio.to_thread(self.drive.share_with_link, drive_file.id)
fields = dict(
drive_file_id=drive_file.id,
final_name=drive_file.name,
completed_at=iso(utcnow()),
upload_session_uri=None,
last_error=None,
error_kind=None,
next_retry_at=None,
)
if drive_file.parents:
fields["drive_folder_id"] = drive_file.parents[0]
if local:
fields.update(local_size=local.size, md5=local.md5, sha256=local.sha256)
self.db.transition(rec.key, State.COMPLETED, note="reused existing Drive file" if reused else "uploaded and verified", **fields)
self.db.raise_int(COMPLETED_CURSOR_KEY, rec.message_id)
if rec.task_id is not None:
self.db.set_task_item_result(
rec.task_id, rec.channel_id, rec.message_id, "duplicate" if reused else "uploaded",
folder=" / ".join(self.folder_path(rec)),
)
# Local files are removed only after the COMPLETED state is committed.
if not self.config.keep_local_files:
self._cleanup_files(rec)
if reused:
self.stats.reused += 1
else:
self.stats.uploaded += 1
log.info("Message %s completed: %r -> Drive %s", rec.message_id, drive_file.name, drive_file.id)
def _cleanup_files(self, rec: VideoRecord) -> int:
deleted = self.downloader.cleanup(rec)
if deleted and rec.task_id is not None:
self.db.add_task_counter(rec.task_id, "deleted_files_count", deleted)
return deleted
def _skip(self, key, reason: str) -> None:
rec = self.db.transition(key, State.SKIPPED, skip_reason=reason, note=reason)
if rec.task_id is not None:
result = "duplicate" if reason.startswith("duplicate") else "skipped"
self.db.set_task_item_result(rec.task_id, rec.channel_id, rec.message_id, result)
if not self.config.keep_local_files:
self._cleanup_files(rec)
self.stats.skipped += 1
log.warning("Message %s skipped: %s", key[1], reason)
def _park(self, key, err: BotError) -> None:
"""Keep the item for a later retry without consuming an attempt."""
rec = self.db.get(key)
fields = dict(last_error=str(err)[:2000], error_kind=type(err).__name__, next_retry_at=iso(utcnow()))
if rec.state == State.RETRY_WAIT:
self.db.update_fields(key, **fields)
else:
self.db.transition(key, State.RETRY_WAIT, note="queue halted", **fields)
async def _retry_or_fail(self, key, err: TransientError) -> None:
rec = self.db.get(key)
attempts = rec.attempts + (0 if err.rate_limited else 1)
if attempts >= self.config.max_attempts:
await self._fail(key, err, attempts=attempts)
return
if err.retry_after is not None:
delay = float(err.retry_after)
else:
base = self.config.retry_base_delay * (2 ** max(attempts - 1, 0))
delay = min(base, self.config.retry_max_delay) * random.uniform(0.9, 1.1)
fields = dict(
attempts=attempts, last_error=str(err)[:2000], error_kind=type(err).__name__, next_retry_at=retry_time(delay)
)
if rec.state == State.RETRY_WAIT:
self.db.update_fields(key, **fields)
else:
self.db.transition(key, State.RETRY_WAIT, note="transient error", **fields)
self.stats.retries += 1
log.warning("Message %s: %s; retry in %.0fs (attempt %d/%d)", key[1], err, delay, attempts, self.config.max_attempts)
await self._notify(
f"Message {key[1]}: {public_error(err)}\nWill retry in {format_duration(delay)} "
f"(attempt {attempts} of {self.config.max_attempts})."
)
async def _fail(self, key, err: BotError, *, attempts: int | None = None) -> None:
rec = self.db.get(key)
fields = dict(last_error=str(err)[:2000], error_kind=type(err).__name__, next_retry_at=None)
if attempts is not None:
fields["attempts"] = attempts
if rec.state == State.FAILED:
rec = self.db.get(key)
self.db.update_fields(key, **fields)
else:
rec = self.db.transition(key, State.FAILED, note="failed", **fields)
if rec.task_id is not None:
self.db.set_task_item_result(rec.task_id, rec.channel_id, rec.message_id, "failed")
# The failure is recorded first; the partial/complete local file can always be
# downloaded again, so it is removed unless KEEP_FAILED_DOWNLOADS=true.
if not self.config.keep_failed_downloads and not self.config.keep_local_files:
self._cleanup_files(rec)
self.stats.failed += 1
log.error("Message %s failed: %s", key[1], err)
await self._notify(f"Message {key[1]} FAILED: {public_error(err)}\nNot retried automatically; use /retry.")
if self.config.strict_stop:
await self._halt(f"STRICT_STOP is enabled and message {key[1]} failed")