| 1 | """FastAPI REST 服务入口。 |
| 2 | |
| 3 | HTTP 层薄封装： |
| 4 | - 接收 URL，创建 job，返回 job_id |
| 5 | - 实际下载委托给 cli.main.download_url 的简化复用 |
| 6 | |
| 7 | fastapi/uvicorn 是**可选**依赖。若未安装，导入本模块会 ImportError。 |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | from contextlib import asynccontextmanager |
| 13 | from typing import Any, Dict, List |
| 14 | |
| 15 | from fastapi import FastAPI, HTTPException |
| 16 | from pydantic import BaseModel |
| 17 | |
| 18 | from auth import CookieManager |
| 19 | from config import ConfigLoader |
| 20 | from control import QueueManager, RateLimiter, RetryHandler |
| 21 | from core import DouyinAPIClient, DownloaderFactory, URLParser |
| 22 | from server.jobs import JobManager |
| 23 | from storage import FileManager |
| 24 | from utils.logger import setup_logger |
| 25 | from utils.validators import is_short_url, normalize_short_url |
| 26 | |
| 27 | logger = setup_logger("REST") |
| 28 | |
| 29 | |
| 30 | class DownloadRequest(BaseModel): |
| 31 | url: str |
| 32 | |
| 33 | |
| 34 | class JobResponse(BaseModel): |
| 35 | job_id: str |
| 36 | status: str |
| 37 | url: str |
| 38 | |
| 39 | |
| 40 | class _ServerDeps: |
| 41 | """跨请求复用的重量级依赖。 |
| 42 | |
| 43 | REST 服务在进程生命周期内只需要一份 FileManager / RateLimiter / RetryHandler / |
| 44 | QueueManager / CookieManager；每个请求重新构造既浪费又会触发文件系统 mkdir。 |
| 45 | DouyinAPIClient 由于持有 aiohttp.ClientSession，依旧按请求创建，避免跨请求泄漏 |
| 46 | 连接状态或触发 "Session is closed" 错误。 |
| 47 | """ |
| 48 | |
| 49 | def __init__(self, config: ConfigLoader): |
| 50 | self.config = config |
| 51 | # Resolve the cookie file path relative to the config file's directory |
| 52 | # so the sidecar can find it regardless of its working directory (which |
| 53 | # on macOS is often '/' when launched by Electron). |
| 54 | if config.config_path: |
| 55 | from pathlib import Path |
| 56 | |
| 57 | cookie_file = str(Path(config.config_path).resolve().parent / ".cookies.json") |
| 58 | else: |
| 59 | cookie_file = ".cookies.json" |
| 60 | self.cookie_manager = CookieManager(cookie_file=cookie_file) |
| 61 | # Load cookies from the config (env var / YAML cookie key) first, then |
| 62 | # fall back to whatever is already on disk in the cookie file. This |
| 63 | # ensures that cookies saved by a previous session are picked up on |
| 64 | # restart even when the config doesn't embed them inline. |
| 65 | initial_cookies = config.get_cookies() |
| 66 | if initial_cookies: |
| 67 | self.cookie_manager.set_cookies(initial_cookies) |
| 68 | else: |
| 69 | # Trigger a load from disk so get_cookies() returns the persisted |
| 70 | # session without requiring a fresh login on every app restart. |
| 71 | self.cookie_manager.get_cookies() |
| 72 | self.file_manager = FileManager(config.get("path")) |
| 73 | self.rate_limiter = RateLimiter(max_per_second=float(config.get("rate_limit", 2) or 2)) |
| 74 | self.retry_handler = RetryHandler(max_retries=int(config.get("retry_times", 3) or 3)) |
| 75 | self.queue_manager = QueueManager(max_workers=int(config.get("thread", 5) or 5)) |
| 76 | |
| 77 | |
| 78 | async def _execute_download(url: str, deps: "_ServerDeps") -> Dict[str, int]: |
| 79 | """简化版 download_url：只负责执行并返回成功/失败计数。 |
| 80 | |
| 81 | 有意不复用 cli.main.download_url —— 后者绑定了 progress_display 的 rich 状态。 |
| 82 | API client 仍按请求创建（aiohttp session 不跨请求复用）；其余重量级依赖从 |
| 83 | _ServerDeps 共享。 |
| 84 | """ |
| 85 | async with DouyinAPIClient(deps.cookie_manager.get_cookies()) as api_client: |
| 86 | if is_short_url(url): |
| 87 | resolved = await api_client.resolve_short_url(normalize_short_url(url)) |
| 88 | if not resolved: |
| 89 | raise RuntimeError(f"Failed to resolve short URL: {url}") |
| 90 | url = resolved |
| 91 | |
| 92 | parsed = URLParser.parse(url) |
| 93 | if not parsed: |
| 94 | raise RuntimeError(f"Unsupported URL: {url}") |
| 95 | |
| 96 | downloader = DownloaderFactory.create( |
| 97 | parsed["type"], |
| 98 | deps.config, |
| 99 | api_client, |
| 100 | deps.file_manager, |
| 101 | deps.cookie_manager, |
| 102 | None, # database 不在 server 场景里启用，避免单例冲突 |
| 103 | deps.rate_limiter, |
| 104 | deps.retry_handler, |
| 105 | deps.queue_manager, |
| 106 | progress_reporter=None, |
| 107 | ) |
| 108 | if downloader is None: |
| 109 | raise RuntimeError(f"No downloader for url_type={parsed['type']}") |
| 110 | |
| 111 | result = await downloader.download(parsed) |
| 112 | return { |
| 113 | "total": result.total, |
| 114 | "success": result.success, |
| 115 | "failed": result.failed, |
| 116 | "skipped": result.skipped, |
| 117 | } |
| 118 | |
| 119 | |
| 120 | def build_app(config: ConfigLoader) -> FastAPI: |
| 121 | deps = _ServerDeps(config) |
| 122 | |
| 123 | async def executor(url: str) -> Dict[str, int]: |
| 124 | return await _execute_download(url, deps) |
| 125 | |
| 126 | server_cfg = config.get("server") or {} |
| 127 | if not isinstance(server_cfg, dict): |
| 128 | server_cfg = {} |
| 129 | manager = JobManager( |
| 130 | executor=executor, |
| 131 | max_concurrency=int(config.get("thread", 2) or 2), |
| 132 | max_jobs=int(server_cfg.get("max_jobs") or JobManager.DEFAULT_MAX_JOBS), |
| 133 | job_ttl_seconds=float( |
| 134 | server_cfg.get("job_ttl_seconds") or JobManager.DEFAULT_JOB_TTL_SECONDS |
| 135 | ), |
| 136 | ) |
| 137 | |
| 138 | @asynccontextmanager |
| 139 | async def lifespan(app: FastAPI): |
| 140 | yield |
| 141 | await manager.shutdown() |
| 142 | |
| 143 | app = FastAPI( |
| 144 | title="Douyin Downloader API", |
| 145 | version="1.0", |
| 146 | description="REST API for dispatching Douyin download jobs.", |
| 147 | lifespan=lifespan, |
| 148 | ) |
| 149 | app.state.job_manager = manager |
| 150 | app.state.deps = deps |
| 151 | |
| 152 | @app.get("/api/v1/health") |
| 153 | async def health() -> Dict[str, str]: |
| 154 | return {"status": "ok"} |
| 155 | |
| 156 | @app.post("/api/v1/download", response_model=JobResponse) |
| 157 | async def create_job(req: DownloadRequest) -> JobResponse: |
| 158 | if not req.url: |
| 159 | raise HTTPException(status_code=400, detail="url is required") |
| 160 | job = await manager.submit(req.url) |
| 161 | return JobResponse(job_id=job.job_id, status=job.status, url=job.url) |
| 162 | |
| 163 | @app.get("/api/v1/jobs/{job_id}") |
| 164 | async def get_job(job_id: str) -> Dict[str, Any]: |
| 165 | job = await manager.get(job_id) |
| 166 | if job is None: |
| 167 | raise HTTPException(status_code=404, detail="job not found") |
| 168 | return job.to_dict() |
| 169 | |
| 170 | @app.get("/api/v1/jobs") |
| 171 | async def list_jobs() -> Dict[str, List[Dict[str, Any]]]: |
| 172 | jobs = await manager.list_jobs() |
| 173 | return {"jobs": [j.to_dict() for j in jobs]} |
| 174 | |
| 175 | return app |
| 176 | |
| 177 | |
| 178 | async def run_server(config: ConfigLoader, *, host: str, port: int) -> None: |
| 179 | import uvicorn |
| 180 | |
| 181 | app = build_app(config) |
| 182 | uv_config = uvicorn.Config(app, host=host, port=port, log_level="info") |
| 183 | server = uvicorn.Server(uv_config) |
| 184 | await server.serve() |
| 185 |