| 1 | """Re-run a known list of failed aweme_ids against a previously-completed job. |
| 2 | |
| 3 | This module implements the "in-place retry" (方案 B) leg of the download |
| 4 | failure retry redesign. The entry point is :func:`retry_failed_awemes`, which: |
| 5 | |
| 6 | 1. Parses the original job URL to recover ``url_type`` and (when applicable) |
| 7 | ``sec_uid`` / ``mix_id`` / ``music_id``. |
| 8 | 2. Instantiates a :class:`core.downloader_base.BaseDownloader` via |
| 9 | :class:`DownloaderFactory`. The factory routes ``video``/``gallery`` to |
| 10 | :class:`VideoDownloader` and batch types (``user``/``collection``/ |
| 11 | ``music``) to their respective batch downloader; either way we only use |
| 12 | it as a host for the shared ``_download_aweme_assets`` method and the |
| 13 | ``mode`` semantics attached to the class, not for its batch-paging |
| 14 | behaviour. |
| 15 | 3. Derives the mode (``post`` / ``like`` / ``mix`` / ``music`` / ``None`` |
| 16 | for single-video) from the job's ``overrides``. |
| 17 | 4. Fetches each ``aweme_id``'s detail via ``api_client.get_video_detail``, |
| 18 | then runs ``_download_aweme_assets(aweme_data, author_name, mode=mode)``. |
| 19 | 5. Emits per-item progress through the reporter so SSE subscribers see the |
| 20 | retry unfold, and forwards per-item outcome to ``on_item_outcome`` so |
| 21 | :class:`server.jobs.JobManager` can update ``job.success`` / ``failed`` |
| 22 | in place. |
| 23 | |
| 24 | The sibling CLI project does not use this module — it ships in the shared |
| 25 | ``core/`` tree per ``AGENTS.md`` and will be synced, but only the desktop |
| 26 | sidecar currently wires a retry_executor. Keeping the helper here rather |
| 27 | than inside ``server/`` avoids leaking HTTP concerns into the download |
| 28 | strategies and keeps the sync story straightforward. |
| 29 | """ |
| 30 | |
| 31 | from __future__ import annotations |
| 32 | |
| 33 | from typing import Any, Callable, Dict, List, Optional |
| 34 | |
| 35 | from auth import CookieManager |
| 36 | from config import ConfigLoader |
| 37 | from control import QueueManager, RateLimiter, RetryHandler |
| 38 | from core.api_client import DouyinAPIClient |
| 39 | from core.downloader_factory import DownloaderFactory |
| 40 | from core.url_parser import URLParser |
| 41 | from storage import Database, FileManager |
| 42 | from utils.logger import setup_logger |
| 43 | from utils.validators import is_short_url, normalize_short_url |
| 44 | |
| 45 | logger = setup_logger("RetryExecutor") |
| 46 | |
| 47 | |
| 48 | def _derive_mode(overrides: Optional[Dict[str, Any]]) -> Optional[str]: |
| 49 | """Pull the first configured download mode from a job's overrides. |
| 50 | |
| 51 | Batch jobs persist their mode as ``overrides["mode"] = ["post"]`` etc. |
| 52 | Single-video jobs leave ``mode`` unset — which naturally means the |
| 53 | files land at ``<base>/<author>/<leaf>`` without a mode subdirectory. |
| 54 | Returning ``None`` for that case is intentional: ``get_save_path`` then |
| 55 | omits the middle segment, matching the original run. |
| 56 | """ |
| 57 | if not overrides: |
| 58 | return None |
| 59 | raw = overrides.get("mode") |
| 60 | if isinstance(raw, str): |
| 61 | candidate = raw.strip() |
| 62 | return candidate or None |
| 63 | if isinstance(raw, list) and raw: |
| 64 | first = raw[0] |
| 65 | if isinstance(first, str): |
| 66 | candidate = first.strip() |
| 67 | return candidate or None |
| 68 | return None |
| 69 | |
| 70 | |
| 71 | def _derive_url_type_for_factory(url_type: Optional[str]) -> str: |
| 72 | """Map the URL parser's type to a DownloaderFactory key. |
| 73 | |
| 74 | Per-aweme retry always uses the per-item code path |
| 75 | (``_download_aweme_assets``), so for batch-flavoured URLs we reuse |
| 76 | :class:`VideoDownloader` (single-item host). This keeps the factory |
| 77 | from kicking off a new paging run that would re-enumerate the whole |
| 78 | user profile / mix / music collection. |
| 79 | """ |
| 80 | # Every branch maps to ``video`` today; keeping the mapping explicit |
| 81 | # so future URL types (live/collect) slot in without silently falling |
| 82 | # back to batch enumeration. |
| 83 | if url_type in ("video", "gallery", "user", "collection", "music"): |
| 84 | return "video" |
| 85 | return "video" |
| 86 | |
| 87 | |
| 88 | async def retry_failed_awemes( |
| 89 | url: str, |
| 90 | *, |
| 91 | aweme_ids: List[str], |
| 92 | config: ConfigLoader, |
| 93 | file_manager: FileManager, |
| 94 | cookie_manager: CookieManager, |
| 95 | database: Optional[Database] = None, |
| 96 | rate_limiter: Optional[RateLimiter] = None, |
| 97 | retry_handler: Optional[RetryHandler] = None, |
| 98 | queue_manager: Optional[QueueManager] = None, |
| 99 | reporter: Any = None, |
| 100 | overrides: Optional[Dict[str, Any]] = None, |
| 101 | author_hint: Optional[Dict[str, Any]] = None, |
| 102 | on_item_outcome: Optional[Callable[[str], None]] = None, |
| 103 | ) -> Dict[str, int]: |
| 104 | """Retry the given aweme ids in place and return summary counts. |
| 105 | |
| 106 | Returns a dict with ``attempted`` / ``succeeded`` / ``failed`` / |
| 107 | ``skipped`` counters so callers that are not subscribed to the reporter |
| 108 | (e.g. unit tests) can assert outcomes. The actual SSE event stream is |
| 109 | driven by ``reporter`` when provided. |
| 110 | """ |
| 111 | counts = { |
| 112 | "attempted": 0, |
| 113 | "succeeded": 0, |
| 114 | "failed": 0, |
| 115 | "skipped": 0, |
| 116 | } |
| 117 | if not aweme_ids: |
| 118 | return counts |
| 119 | |
| 120 | # Apply overrides (mode/path/…) for the duration of the retry so the |
| 121 | # downloader's `get_save_path` picks up the same output_dir and folder |
| 122 | # template that the original run used. Snapshot + restore the values |
| 123 | # we overwrite, matching `_execute_download` in server/app.py. |
| 124 | snap: Dict[str, Any] = {} |
| 125 | if overrides: |
| 126 | for k in overrides.keys(): |
| 127 | snap[k] = config.get(k) |
| 128 | config.update(**overrides) |
| 129 | |
| 130 | try: |
| 131 | cookies = cookie_manager.get_cookies() |
| 132 | async with DouyinAPIClient(cookies) as api_client: |
| 133 | if is_short_url(url): |
| 134 | resolved = await api_client.resolve_short_url(normalize_short_url(url)) |
| 135 | if not resolved: |
| 136 | raise RuntimeError(f"Failed to resolve short URL during retry: {url}") |
| 137 | url = resolved |
| 138 | |
| 139 | parsed = URLParser.parse(url) |
| 140 | if not parsed: |
| 141 | raise RuntimeError(f"Unsupported URL during retry: {url}") |
| 142 | |
| 143 | mode = _derive_mode(overrides) |
| 144 | factory_type = _derive_url_type_for_factory(parsed.get("type")) |
| 145 | |
| 146 | downloader = DownloaderFactory.create( |
| 147 | factory_type, |
| 148 | config, |
| 149 | api_client, |
| 150 | file_manager, |
| 151 | cookie_manager, |
| 152 | database, |
| 153 | rate_limiter, |
| 154 | retry_handler, |
| 155 | queue_manager, |
| 156 | progress_reporter=reporter, |
| 157 | ) |
| 158 | if downloader is None: |
| 159 | raise RuntimeError(f"No downloader available for retry (url_type={factory_type})") |
| 160 | |
| 161 | if reporter is not None: |
| 162 | try: |
| 163 | reporter.on_job_start( |
| 164 | url=url, |
| 165 | url_type=parsed.get("type") or factory_type, |
| 166 | total=len(aweme_ids), |
| 167 | ) |
| 168 | except Exception: |
| 169 | pass |
| 170 | # set_item_total drives the percentage bar in the renderer; |
| 171 | # advance_item later hooks into the same running total. |
| 172 | try: |
| 173 | reporter.set_item_total( |
| 174 | len(aweme_ids), detail=f"重试 {len(aweme_ids)} 个失败作品" |
| 175 | ) |
| 176 | except Exception: |
| 177 | pass |
| 178 | |
| 179 | # Default author name falls back to whatever is stored in the |
| 180 | # per-aweme detail payload. Hint is only used when the detail |
| 181 | # lookup itself fails — in that case we still want a sensible |
| 182 | # folder for any partial artifacts the downloader writes. |
| 183 | hint_nickname = None |
| 184 | if author_hint and isinstance(author_hint, dict): |
| 185 | hint_nickname = author_hint.get("nickname") |
| 186 | |
| 187 | for aid in aweme_ids: |
| 188 | aid_str = str(aid or "").strip() |
| 189 | if not aid_str: |
| 190 | counts["skipped"] += 1 |
| 191 | counts["attempted"] += 1 |
| 192 | if on_item_outcome: |
| 193 | on_item_outcome("skipped") |
| 194 | if reporter is not None: |
| 195 | try: |
| 196 | reporter.on_log( |
| 197 | level="warning", |
| 198 | message="跳过 · aweme_id 为空", |
| 199 | ) |
| 200 | except Exception: |
| 201 | pass |
| 202 | continue |
| 203 | |
| 204 | aweme_data = await api_client.get_video_detail(aid_str) |
| 205 | if not aweme_data: |
| 206 | counts["failed"] += 1 |
| 207 | counts["attempted"] += 1 |
| 208 | if on_item_outcome: |
| 209 | on_item_outcome("failed") |
| 210 | if reporter is not None: |
| 211 | try: |
| 212 | reporter.on_log( |
| 213 | level="error", |
| 214 | message=f"获取作品详情失败 · {aid_str}", |
| 215 | ) |
| 216 | reporter.advance_item("failed", detail=aid_str) |
| 217 | except Exception: |
| 218 | pass |
| 219 | continue |
| 220 | |
| 221 | author = aweme_data.get("author") or {} |
| 222 | author_name = author.get("nickname") or hint_nickname or "unknown" |
| 223 | |
| 224 | ok = False |
| 225 | try: |
| 226 | ok = await downloader._download_aweme_assets(aweme_data, author_name, mode=mode) |
| 227 | except Exception as exc: # pragma: no cover — defensive |
| 228 | logger.warning( |
| 229 | "Retry of aweme %s raised %s: %s", |
| 230 | aid_str, |
| 231 | type(exc).__name__, |
| 232 | exc, |
| 233 | ) |
| 234 | ok = False |
| 235 | |
| 236 | counts["attempted"] += 1 |
| 237 | if ok: |
| 238 | counts["succeeded"] += 1 |
| 239 | if on_item_outcome: |
| 240 | on_item_outcome("ok") |
| 241 | if reporter is not None: |
| 242 | try: |
| 243 | reporter.advance_item("success", detail=aid_str) |
| 244 | except Exception: |
| 245 | pass |
| 246 | else: |
| 247 | counts["failed"] += 1 |
| 248 | if on_item_outcome: |
| 249 | on_item_outcome("failed") |
| 250 | if reporter is not None: |
| 251 | try: |
| 252 | reporter.advance_item("failed", detail=aid_str) |
| 253 | except Exception: |
| 254 | pass |
| 255 | finally: |
| 256 | if overrides: |
| 257 | config.update(**snap) |
| 258 | |
| 259 | return counts |
| 260 |