新增PDF批量预生成

backend/app/services/pdf_pregeneration.py

@@ -0,0 +1,182 @@
+import queue
+import threading
+import uuid
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from pathlib import Path
+
+from app.db.session import SessionLocal
+from app.models import Certificate, CertificateAccessToken, Learner, ProjectCourse
+from app.services.pdf import cached_pdf_is_fresh, certificate_template_path, pdf_cache_path, render_certificate_pdf
+
+MAX_PDF_PREGENERATION_JOBS = 30
+
+
+def _utc_now() -> datetime:
+    return datetime.now(UTC).replace(tzinfo=None)
+
+
+@dataclass
+class PdfPregenerationJob:
+    id: str
+    certificate_ids: list[int]
+    concurrency_limit: int
+    status: str = "running"
+    total: int = 0
+    completed: int = 0
+    generated: int = 0
+    cached: int = 0
+    skipped: int = 0
+    failed: int = 0
+    errors: list[str] = field(default_factory=list)
+    created_at: datetime = field(default_factory=_utc_now)
+    finished_at: datetime | None = None
+
+    def to_dict(self) -> dict[str, object]:
+        percent = round((self.completed / self.total) * 100) if self.total else 0
+        return {
+            "id": self.id,
+            "status": self.status,
+            "total": self.total,
+            "completed": self.completed,
+            "generated": self.generated,
+            "cached": self.cached,
+            "skipped": self.skipped,
+            "failed": self.failed,
+            "percent": percent,
+            "concurrency_limit": self.concurrency_limit,
+            "errors": self.errors[-10:],
+            "created_at": self.created_at,
+            "finished_at": self.finished_at,
+        }
+
+
+class PdfPregenerationManager:
+    def __init__(self) -> None:
+        self._jobs: dict[str, PdfPregenerationJob] = {}
+        self._lock = threading.Lock()
+
+    def start(self, certificate_ids: list[int], concurrency_limit: int) -> PdfPregenerationJob:
+        unique_ids = _unique_ids(certificate_ids)
+        if not unique_ids:
+            raise ValueError("请选择需要预生成PDF的证书")
+        job = PdfPregenerationJob(
+            id=uuid.uuid4().hex,
+            certificate_ids=unique_ids,
+            concurrency_limit=max(1, concurrency_limit),
+            total=len(unique_ids),
+        )
+        with self._lock:
+            self._jobs[job.id] = job
+            self._trim_jobs()
+        thread = threading.Thread(target=self._run_job, args=(job.id,), daemon=True)
+        thread.start()
+        return job
+
+    def get(self, job_id: str) -> PdfPregenerationJob | None:
+        with self._lock:
+            return self._jobs.get(job_id)
+
+    def list(self) -> list[PdfPregenerationJob]:
+        with self._lock:
+            return sorted(self._jobs.values(), key=lambda item: item.created_at, reverse=True)
+
+    def _run_job(self, job_id: str) -> None:
+        job = self.get(job_id)
+        if not job:
+            return
+
+        work_queue: queue.Queue[int] = queue.Queue()
+        for certificate_id in job.certificate_ids:
+            work_queue.put(certificate_id)
+
+        workers = [
+            threading.Thread(target=self._worker, args=(job_id, work_queue), daemon=True)
+            for _ in range(min(job.concurrency_limit, job.total))
+        ]
+        for worker in workers:
+            worker.start()
+        for worker in workers:
+            worker.join()
+
+        with self._lock:
+            job.status = "completed_with_errors" if job.failed else "completed"
+            job.finished_at = _utc_now()
+
+    def _worker(self, job_id: str, work_queue: queue.Queue[int]) -> None:
+        while True:
+            try:
+                certificate_id = work_queue.get_nowait()
+            except queue.Empty:
+                return
+            try:
+                job = self.get(job_id)
+                if not job:
+                    return
+                result = pre_generate_certificate_pdf(certificate_id, job.concurrency_limit)
+                self._record(job_id, result)
+            except Exception as exc:
+                self._record(job_id, "failed", f"证书ID {certificate_id}：{exc}")
+            finally:
+                work_queue.task_done()
+
+    def _record(self, job_id: str, result: str, error: str | None = None) -> None:
+        job = self.get(job_id)
+        if not job:
+            return
+        with self._lock:
+            job.completed += 1
+            if result == "generated":
+                job.generated += 1
+            elif result == "cached":
+                job.cached += 1
+            elif result == "skipped":
+                job.skipped += 1
+            else:
+                job.failed += 1
+                if error:
+                    job.errors.append(error)
+
+    def _trim_jobs(self) -> None:
+        if len(self._jobs) <= MAX_PDF_PREGENERATION_JOBS:
+            return
+        old_jobs = sorted(self._jobs.values(), key=lambda item: item.created_at)
+        for job in old_jobs[: len(self._jobs) - MAX_PDF_PREGENERATION_JOBS]:
+            self._jobs.pop(job.id, None)
+
+
+def pre_generate_certificate_pdf(certificate_id: int, concurrency_limit: int) -> str:
+    with SessionLocal() as db:
+        certificate = db.get(Certificate, certificate_id)
+        if not certificate or certificate.status != "valid":
+            return "skipped"
+        learner = db.get(Learner, certificate.learner_id)
+        token = db.get(CertificateAccessToken, certificate.qr_token_id or certificate.public_token_id)
+        project = db.query(ProjectCourse).filter(ProjectCourse.code == certificate.project_code).first()
+        if not learner or not token:
+            raise ValueError("证书数据不完整")
+
+        was_cached = certificate_pdf_is_ready(certificate)
+        render_certificate_pdf(certificate, learner, project, token, concurrency_limit=concurrency_limit)
+        db.commit()
+        return "cached" if was_cached else "generated"
+
+
+def certificate_pdf_is_ready(certificate: Certificate) -> bool:
+    output_path = pdf_cache_path(certificate.certificate_no)
+    template_path = certificate_template_path()
+    latest_renderer_mtime = max(template_path.stat().st_mtime, Path(__file__).with_name("pdf.py").stat().st_mtime)
+    return cached_pdf_is_fresh(output_path) and output_path.stat().st_mtime >= latest_renderer_mtime
+
+
+def _unique_ids(values: list[int]) -> list[int]:
+    seen: set[int] = set()
+    result: list[int] = []
+    for value in values:
+        if value not in seen:
+            seen.add(value)
+            result.append(value)
+    return result
+
+
+pdf_pregeneration_manager = PdfPregenerationManager()