document-management/backend/dms/backfill_preview_cache.py

"""显式、幂等地为既有 DOC/DOCX/PDF 生成或校验预览缓存。

用法：
    python -m dms.backfill_preview_cache --dry-run
    python -m dms.backfill_preview_cache --limit 100

仅允许显式配置并连接既有 ``dms_test``：
``python -m dms.backfill_preview_cache``。
"""

from __future__ import annotations

import argparse
import logging
import os
from pathlib import Path
from datetime import datetime, timezone

from flask import Flask, current_app
from sqlalchemy import select, text
from sqlalchemy.engine import make_url

from dms import init_dms
from dms.common.enums import ContentExtractStatus
from dms.extensions import db
from dms.models import Document
from dms.services.document_converter import create_converter
from dms.services.preview_service import (
    PreviewService,
    extract_doc_content_via_pdf,
    preview_cache_key,
)
from dms.services.document_content_service import build_search_text
from dms.storage.paths import preview_cache_path, resolve_storage_path

logger = logging.getLogger(__name__)


def _database_gate() -> None:
    configured = os.environ.get("DMS_DATABASE_URL", "").strip()
    if not configured:
        raise RuntimeError("必须显式配置DMS_DATABASE_URL并指向dms_test")
    try:
        configured_database = make_url(configured).database
    except Exception as exc:
        raise RuntimeError("DMS_DATABASE_URL不是有效数据库连接配置") from exc
    if configured_database != "dms_test":
        raise RuntimeError("预览缓存回填命令只允许连接dms_test")
    database_name = db.session.scalar(text("SELECT DATABASE()"))
    if database_name != "dms_test":
        raise RuntimeError("预览缓存回填命令只允许连接dms_test")


def backfill(*, dry_run: bool = False, limit: int | None = None) -> dict[str, int]:
    _database_gate()
    result = {
        "processed": 0,
        "skipped": 0,
        "failed": 0,
        "doc_content_backfilled": 0,
    }

    converter = create_converter(
        {
            "DMS_OFFICE_PREVIEW_ENABLED": current_app.config.get(
                "DMS_OFFICE_PREVIEW_ENABLED", True
            ),
            "DMS_LIBREOFFICE_EXECUTABLE": current_app.config.get(
                "DMS_LIBREOFFICE_EXECUTABLE", ""
            ),
            "DMS_LIBREOFFICE_TIMEOUT_SECONDS": current_app.config.get(
                "DMS_LIBREOFFICE_TIMEOUT_SECONDS", 60
            ),
            "DMS_LIBREOFFICE_MAX_CONCURRENCY": current_app.config.get(
                "DMS_LIBREOFFICE_MAX_CONCURRENCY", 2
            ),
        }
    )
    storage_root = Path(
        current_app.config.get("DMS_STORAGE_ROOT", "")
    ).expanduser().resolve() or (
        Path(__file__).resolve().parents[2] / "dms-storage"
    )

    query = (
        select(Document)
        .where(
            Document.is_deleted.is_(False),
            Document.file_extension.in_(["doc", "docx", "pdf"]),
        )
        .order_by(Document.id.asc())
    )
    if limit:
        query = query.limit(limit)
    documents = db.session.scalars(query).all()

    for document in documents:
        extension = document.file_extension.lower().lstrip(".")
        try:
            original_path = resolve_storage_path(
                document.file_relative_path, storage_root
            )
            if not original_path.is_file():
                logger.warning("文件不存在：document_id=%s", document.id)
                result["failed"] += 1
                continue
            if extension == "pdf":
                result["skipped"] += 1
                continue

            preview_service = PreviewService(storage_root, converter)
            cache_path = preview_cache_path(
                preview_cache_key(document.file_hash, extension), storage_root
            )
            if preview_service.validate_cached_pdf(cache_path):
                result["skipped"] += 1
            elif dry_run:
                logger.info(
                    "dry-run：需要生成预览缓存 document_id=%s", document.id
                )
                result["processed"] += 1
            else:
                preview_service.preview_path_for(document)
                result["processed"] += 1

            if extension == "doc" and not dry_run:
                if document.content_extract_status in {
                    ContentExtractStatus.SUCCESS.value,
                    ContentExtractStatus.EMPTY.value,
                }:
                    pass
                else:
                    text_value, status = extract_doc_content_via_pdf(
                        original_path, converter, document.file_hash
                    )
                    if status != ContentExtractStatus.FAILED.value:
                        document.content_text = text_value
                        document.content_extract_status = status
                        document.content_extracted_at = datetime.now(
                            timezone.utc
                        ).replace(tzinfo=None)
                        document.search_text = build_search_text(
                            document.document_name,
                            document.summary,
                            document.tags,
                            text_value,
                        )
                        db.session.commit()
                        result["doc_content_backfilled"] += 1

        except Exception:
            logger.exception("回填预览缓存失败：document_id=%s", document.id)
            result["failed"] += 1
            db.session.rollback()

    return result


def main() -> None:
    parser = argparse.ArgumentParser(description="DMS 预览缓存回填命令")
    parser.add_argument(
        "--dry-run", action="store_true", help="仅输出需要处理的文档，不实际转换"
    )
    parser.add_argument(
        "--limit", type=int, default=None, help="限制处理的文档数量"
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    app = Flask("dms-preview-backfill")
    init_dms(app)
    with app.app_context():
        outcome = backfill(dry_run=args.dry_run, limit=args.limit)
    print(
        "预览缓存回填完成："
        f"处理{outcome['processed']}个，跳过{outcome['skipped']}个，"
        f"失败{outcome['failed']}个，"
        f"DOC正文回填{outcome['doc_content_backfilled']}个。"
    )


if __name__ == "__main__":
    main()