"""Enqueue + cache resolution for the download center (the per-user layer over MediaAsset). `enqueue` is the single entry point used by the API: it resolves (or creates) the shared MediaAsset for the requested source+format, links a per-user DownloadJob to it, and — on a cache hit (asset already ready) — marks the job done instantly with no re-download. The worker handles everything else. State transitions used by the routes (pause/resume/cancel/delete) and the worker live in `worker.py` / the routes; this module owns enqueue + asset bookkeeping. """ from datetime import datetime, timezone from sqlalchemy import func, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from app.config import settings from app.downloads import adapters from app.downloads import edit as editmod from app.downloads import formats, quota, storage from app.downloads.status import HOLDING, JobStatus from app.models import Channel, DownloadJob, MediaAsset, Video class EditError(Exception): """Bad edit request (empty spec / source not ready); `reason` is an i18n-mappable key.""" def __init__(self, reason: str): self.reason = reason super().__init__(reason) def source_url(source_kind: str, source_ref: str) -> str: """The URL yt-dlp should fetch. YouTube stores the 11-char id (the adapter reconstructs the watch URL); ad-hoc sources store the URL verbatim.""" return adapters.resolve_adapter(source_kind, source_ref).source_url(source_ref) def reference_url(job, asset) -> str | None: """The clean "downloaded from" URL for a job — shared by the private Downloads page and the public watch page so both surfaces always agree. None for edit derivatives (a clip's source is the user's own earlier download, not a web page); else the canonical page URL yt-dlp recorded; else one derived from source_kind+source_ref (covers queued/older rows before the worker fills it).""" if job is None or job.job_kind == "edit": return None if asset is not None and asset.source_webpage_url: return asset.source_webpage_url if job.source_kind in ("youtube", "url"): return source_url(job.source_kind, job.source_ref) return None def get_or_create_asset( db: Session, source_kind: str, source_ref: str, format_sig: str ) -> MediaAsset: """Fetch the cache row for this (source, format) or create a fresh `pending` one. Races on the unique cache key are resolved by re-selecting after an IntegrityError, so two simultaneous enqueues of the same video+format converge on one shared asset.""" stmt = select(MediaAsset).where( MediaAsset.source_kind == source_kind, MediaAsset.source_ref == source_ref, MediaAsset.format_sig == format_sig, ) asset = db.execute(stmt).scalar_one_or_none() if asset is not None: # A previously-FAILED shared asset must be re-attempted for the new job about to hold it — # reset it to pending (+clear the error) so the worker actually re-downloads instead of # short-circuiting the new job with the asset's stale error. (Mirrors resume_download; without # this a fresh enqueue of a once-failed (source,format) pair is permanently poisoned, since # errored assets carry no expires_at and GC never clears them.) if asset.status == "error": asset.status = "pending" asset.error = None return asset asset = MediaAsset( source_kind=source_kind, source_ref=source_ref, format_sig=format_sig, status="pending", ref_count=0, ) db.add(asset) try: db.flush() except IntegrityError: db.rollback() asset = db.execute(stmt).scalar_one() return asset def populate_from_catalog(db: Session, asset: MediaAsset) -> bool: """Fill an asset's display metadata (title/uploader/thumbnail/date/duration) from our own catalog when the source is a known YouTube video — so a queued/downloading row shows the real title + thumbnail immediately instead of the bare id. 0 network cost. Returns True if it found a catalog match.""" if asset.source_kind != "youtube" or asset.title: return False video = db.get(Video, asset.source_ref) if video is None: return False asset.title = video.title asset.thumbnail_url = video.thumbnail_url asset.duration_s = video.duration_seconds if video.published_at: asset.upload_date = video.published_at.date() channel = db.get(Channel, video.channel_id) if channel is not None: asset.uploader = channel.title return True def next_queue_pos(db: Session) -> int: return (db.execute(select(func.coalesce(func.max(DownloadJob.queue_pos), 0))).scalar() or 0) + 1 def incr_ref_count(db: Session, asset_id: int) -> None: """Atomically +1 an asset's hold count AT THE DB. A Python read-modify-write (`asset.ref_count = asset.ref_count + 1`) loses an update when two enqueues of the same (source,format) run in parallel — both read the same value and write the same +1 — leaving the count one short, so the first release can free a file another job still holds (B-sub-5).""" db.execute( update(MediaAsset) .where(MediaAsset.id == asset_id) .values(ref_count=MediaAsset.ref_count + 1) .execution_options(synchronize_session=False) ) def decr_ref_count(db: Session, asset_id: int) -> int: """Atomically -1 an asset's hold count (clamped at 0) and RETURN the new value. The caller decides whether to free the asset from the returned count while still holding the row's write lock, so two concurrent releases can't both observe the same pre-decrement value and double-free — Postgres serialises the two UPDATEs on the row lock (B-sub-5). synchronize_session=False: the ORM copy is intentionally left stale; callers read the returned value, not `asset.ref_count`.""" return db.execute( update(MediaAsset) .where(MediaAsset.id == asset_id) .values(ref_count=func.greatest(MediaAsset.ref_count - 1, 0)) .returning(MediaAsset.ref_count) .execution_options(synchronize_session=False) ).scalar_one() def enqueue( db: Session, user_id: int, source_kind: str, source_ref: str, spec: dict, profile_id: int | None = None, display_name: str | None = None, ) -> DownloadJob: """Queue a download for `user_id`. Returns the DownloadJob (already `done` on a cache hit). Raises quota.QuotaExceeded if the user is at their job-count or footprint cap.""" quota.check_enqueue(db, user_id) spec = formats.normalize(spec) sig = formats.format_sig(spec) asset = get_or_create_asset(db, source_kind, source_ref, sig) # Show a real title + thumbnail the moment it's queued (feed downloads are catalog videos). populate_from_catalog(db, asset) job = DownloadJob( user_id=user_id, asset_id=asset.id, source_kind=source_kind, source_ref=source_ref, profile_id=profile_id, profile_snapshot=spec, display_name=display_name or (asset.title[:255] if asset.status == "ready" and asset.title else None), status="queued", queue_pos=next_queue_pos(db), ) db.add(job) incr_ref_count(db, asset.id) # Cache hit: the file already exists — no worker round-trip needed. if asset.status == "ready": job.status = "done" job.progress = 100 asset.last_access_at = datetime.now(timezone.utc) db.commit() db.refresh(job) return job def _clip_title(source_title: str | None, display_name: str | None) -> str: if display_name: return display_name[:255] base = source_title or "Clip" return f"{base} (clip)"[:255] def enqueue_edit( db: Session, user_id: int, source_job: DownloadJob, edit_spec: dict, display_name: str | None = None, ) -> DownloadJob: """Queue a phase-2 editor derivative (trim/crop) of `source_job` for `user_id`. The derivative is per-user (never shared-cached): its asset uses `source_kind='edit'` with a cache key that embeds the user id, so an identical re-edit by the same user is a cache hit but it never dedups across users. Reuses the same worker/quota/GC path as a download. Raises quota.QuotaExceeded (holdings cap) or EditError (empty spec / source not ready).""" src_asset = db.get(MediaAsset, source_job.asset_id) if source_job.asset_id else None if src_asset is None or src_asset.status != "ready" or not src_asset.rel_path: raise EditError("source_not_ready") spec = editmod.normalize_edit_spec(edit_spec) if not spec: raise EditError("empty_edit") quota.check_enqueue(db, user_id) sig = editmod.edit_sig(spec) ref = f"{user_id}:{src_asset.id}"[:512] asset = get_or_create_asset(db, "edit", ref, sig) if not asset.title: asset.title = _clip_title(src_asset.title, display_name) asset.uploader = src_asset.uploader asset.thumbnail_url = src_asset.thumbnail_url asset.upload_date = src_asset.upload_date asset.duration_s = editmod.clip_duration(spec, src_asset.duration_s) job = DownloadJob( user_id=user_id, asset_id=asset.id, job_kind="edit", source_kind="edit", source_ref=ref, source_job_id=source_job.id, source_asset_id=src_asset.id, edit_spec=spec, profile_snapshot={}, display_name=display_name or asset.title, status="queued", queue_pos=next_queue_pos(db), ) db.add(job) incr_ref_count(db, asset.id) if asset.status == "ready": # identical clip already produced for this user → instant job.status = "done" job.progress = 100 asset.last_access_at = datetime.now(timezone.utc) db.commit() db.refresh(job) return job def finalize_if_running(db: Session, job_id: int, **fields) -> bool: """Write terminal `fields` to a job in ONE atomic ``UPDATE ... WHERE status='running'`` and report whether it fired. This is the single owner of the invariant "a finishing worker child must never resurrect a job a route just stopped": the atomic WHERE closes the read-then-write race an ``if job.status == 'running'`` guard leaves open, so a concurrent cancel/pause can no longer be overwritten with done/error (which would re-open the double-release → file-delete-under-live-holder class, since `done`/`error` hold the asset but `canceled` does not — see release_asset). Runs in the CALLER's session and does NOT commit — the caller commits it together with its asset writes in one transaction. It flushes first so any pending asset write takes its row lock BEFORE this takes the job's, matching the cancel route's asset-then-job order (no deadlock). `fields` are mapped-column kwargs (never user input). Returns True when the row was still running and got updated, False when a route had already moved it out of 'running'.""" if not fields: return False db.flush() # asset-then-job lock order: push pending asset writes before we lock the job row # A Core UPDATE with the guard as a WHERE — same atomic "only while still running" write as the # rest of the codebase's 2.0 statements, no hand-built SQL string. synchronize_session=False: this # is a terminal write, so we don't pay to sync the (about-to-be-discarded) ORM identity map. Bound # StrEnum values bind as their plain str literal (StrEnum is a str subclass). result = db.execute( update(DownloadJob) .where(DownloadJob.id == job_id, DownloadJob.status == JobStatus.RUNNING.value) .values(**fields) .execution_options(synchronize_session=False) ) return result.rowcount > 0 def release_asset(db: Session, job: DownloadJob) -> None: """Drop this job's hold on its asset when it leaves a holding state. Once NO job holds the asset anymore, delete the file + row immediately — a deleted download should free its disk, and the shared cache only needs to span *overlapping* holders (a later re-add just downloads again). `error` counts as holding: enqueue always +1'd ref_count and the worker never decrements on failure (ref_count is the route's job), so an errored job releases here on delete — else its +1 leaks and a later resume→ready keeps the file pinned past its last holder. Idempotent for the cancel path: cancel already released while the job was holding, then set it `canceled` (∉ the set below), so delete-after-cancel is a no-op and never double-releases. We DON'T delete an errored asset row here even at ref==0: another request may be concurrently re-enqueuing the same (source,format) — get_or_create_asset would reset that row to `pending` and +1 it, and deleting it under that would FK-null the new job's asset (a lost update on ref_count). A ready asset is still freed at ref==0 (its file is the point); an orphaned errored row is cheap (no file) and gets reused+reset by the next enqueue, or reclaimed by a later GC pass.""" if not job.asset_id or job.status not in HOLDING: return asset = db.get(MediaAsset, job.asset_id) if asset is None: return new_count = decr_ref_count(db, asset.id) if new_count == 0 and asset.status == "ready": if asset.rel_path: storage.delete_asset_files(settings.download_root, asset.rel_path) db.delete(asset) # --- Job state transitions --------------------------------------------------------------------- # The single home for the pause/resume/cancel/delete/redownload transitions + their asset # bookkeeping (release, failed-asset reset, fresh-asset rebuild), so the routes are thin callers and # the worker/GC could share the same definitions (B-sub-8). Each commits its own transition; the # route keeps only auth + HTTP mapping + serialization. def pause(db: Session, job: DownloadJob) -> None: """Pause a queued/running job (a running worker aborts cooperatively via the cancel poll); a no-op in any other state.""" if job.status in (JobStatus.QUEUED, JobStatus.RUNNING): job.status = JobStatus.PAUSED db.commit() def resume(db: Session, job: DownloadJob) -> None: """Requeue a paused/errored job for a fresh attempt, clearing its error. If the SHARED asset itself failed, reset it to pending so the worker actually re-downloads — otherwise the requeued job would instantly re-inherit the asset's stale error (the worker short-circuits an errored asset).""" if job.status in (JobStatus.PAUSED, JobStatus.ERROR): job.progress = 0 job.error = None job.error_code = None if job.asset_id: asset = db.get(MediaAsset, job.asset_id) if asset and asset.status == "error": # MediaAsset.status (no enum), not a JobStatus asset.status = "pending" asset.error = None job.status = JobStatus.QUEUED db.commit() def cancel(db: Session, job: DownloadJob) -> None: """Cancel a job unless it's already done/canceled, releasing its asset hold WHILE it still counts as holding — `canceled` ∉ HOLDING, so a later delete won't double-release.""" if job.status not in (JobStatus.DONE, JobStatus.CANCELED): release_asset(db, job) # release while the job still counts as holding job.status = JobStatus.CANCELED db.commit() def delete(db: Session, job: DownloadJob) -> None: """Delete a job, releasing its asset hold first (which frees the file iff this was the last holder).""" release_asset(db, job) db.delete(job) db.commit() def redownload(db: Session, job: DownloadJob) -> None: """Rebuild a finished/failed job's file into a FRESH, private asset (the cache sig is salted with the job id + current asset id, so it never reuses the dedup pool) and requeue the SAME job row — so any public share link (DownloadLink -> job_id) keeps resolving to the newly-downloaded file. The old hold is dropped first (freeing the old file iff this job was its last holder). The caller has already rejected edit clips + still-in-progress jobs (HTTP 409).""" spec = formats.normalize(job.profile_snapshot) sig = formats.format_sig(spec, salt=f"{job.id}:{job.asset_id}") asset = get_or_create_asset(db, job.source_kind, job.source_ref, sig) populate_from_catalog(db, asset) # Drop the old hold BEFORE re-pointing the job — release_asset keys off the job's current asset. release_asset(db, job) incr_ref_count(db, asset.id) job.asset_id = asset.id job.status = JobStatus.QUEUED job.progress = 0 job.phase = None job.error = None job.speed_bps = None job.eta_s = None # Re-derive the display name from the freshly cleaned title on completion (else an older, un-cleaned # name would survive the rebuild). job.display_name = None job.queue_pos = next_queue_pos(db) # tail of the queue, not its stale original slot db.commit()