document-management/backend/dms/backfill_document_content.py

"""显式、幂等地回填既有PDF/DOCX正文及search_text。

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

from __future__ import annotations

import os

from flask import Flask
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_content_service import (
    build_search_text,
    extract_document_content,
)
from dms.services.recycle_bin_service import _verify_file


def _database_gate() -> None:
    """安全闸门：必须显式配置 DMS_DATABASE_URL，且只允许连接 dms / dms_test。"""
    configured = os.environ.get("DMS_DATABASE_URL", "").strip()
    if not configured:
        raise RuntimeError("必须显式配置DMS_DATABASE_URL")
    try:
        configured_database = make_url(configured).database
    except Exception as exc:
        raise RuntimeError("DMS_DATABASE_URL不是有效数据库连接配置") from exc
    allowed = {"dms", "dms_test"}
    if configured_database not in allowed:
        raise RuntimeError(
            f"正文回填命令只允许连接 {' 或 '.join(sorted(allowed))}，"
            f"当前指向 {configured_database!r}"
        )
    database_name = db.session.scalar(text("SELECT DATABASE()"))
    if database_name not in allowed:
        raise RuntimeError(
            f"正文回填命令实际连接的库 {database_name!r} 不在允许列表内"
        )


def backfill() -> dict[str, int]:
    _database_gate()
    result = {"processed": 0, "skipped": 0, "failed": 0}
    document_ids = db.session.scalars(
        select(Document.id).order_by(Document.id.asc())
    ).all()
    for document_id in document_ids:
        document = db.session.get(Document, document_id)
        assert document is not None
        extension = document.file_extension.lower().lstrip(".")
        if extension not in {"pdf", "docx"}:
            if document.content_extract_status != ContentExtractStatus.UNSUPPORTED:
                document.content_text = None
                document.content_extract_status = (
                    ContentExtractStatus.UNSUPPORTED.value
                )
                document.content_extracted_at = None
                document.search_text = build_search_text(
                    document.document_name,
                    document.summary,
                    document.tags,
                    None,
                )
                db.session.commit()
                result["processed"] += 1
            else:
                result["skipped"] += 1
            continue
        if document.content_extract_status in {
            ContentExtractStatus.SUCCESS.value,
            ContentExtractStatus.EMPTY.value,
        } and document.content_extracted_at is not None:
            result["skipped"] += 1
            continue
        verified = None
        try:
            verified = _verify_file(document)
            extraction = extract_document_content(verified.path, extension)
            document.content_text = extraction.text
            document.content_extract_status = extraction.status
            document.content_extracted_at = extraction.extracted_at
            document.search_text = build_search_text(
                document.document_name,
                document.summary,
                document.tags,
                extraction.text,
            )
            db.session.commit()
            result["processed"] += 1
            if extraction.status == ContentExtractStatus.FAILED.value:
                result["failed"] += 1
        except Exception:
            db.session.rollback()
            document = db.session.get(Document, document_id)
            if document is not None:
                document.content_text = None
                document.content_extract_status = ContentExtractStatus.FAILED.value
                document.content_extracted_at = None
                document.search_text = build_search_text(
                    document.document_name,
                    document.summary,
                    document.tags,
                    None,
                )
                db.session.commit()
            result["processed"] += 1
            result["failed"] += 1
        finally:
            if verified is not None:
                verified.close()
    return result


def main() -> None:
    app = Flask("dms-content-backfill")
    init_dms(app)
    with app.app_context():
        outcome = backfill()
    print(
        "正文回填完成："
        f"处理{outcome['processed']}个，跳过{outcome['skipped']}个，"
        f"失败{outcome['failed']}个。"
    )


if __name__ == "__main__":
    main()