document-management/backend/dms/seed_documents.py

"""显式、幂等的B4开发文档、附件及挂载关系初始化。

执行：python -m dms.seed_documents
只创建缺失的约定演示数据，不覆盖有效同名文档，不恢复逻辑删除数据。
"""

from __future__ import annotations

import hashlib
import mimetypes
import uuid
from pathlib import Path

from docx import Document as WordDocument
from flask import Flask, current_app
from sqlalchemy import func, select

from dms import init_dms
from dms.common.enums import (
    AttachmentType,
    DocumentStatus,
    DocumentType,
    SecurityLevel,
    SubjectType,
    VisibilityType,
)
from dms.extensions import db
from dms.models import (
    AttachmentBinding,
    Category,
    Document,
    Organization,
    Permission,
    User,
)
from dms.storage.paths import ensure_storage_directories

DOCX_MIME = (
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
)


def _required(model, *criteria, description: str):
    value = db.session.scalar(select(model).where(*criteria))
    if value is None:
        raise RuntimeError(f"缺少{description}，请先执行B2和B3初始化工具")
    return value


def _create_docx(
    root: Path,
    *,
    title: str,
    paragraphs: list[str],
) -> tuple[str, int, str, Path]:
    relative_path = f"original/{uuid.uuid4()}.docx"
    target = root / Path(relative_path)
    temporary = root / "temporary" / f"{uuid.uuid4()}.docx"
    try:
        word = WordDocument()
        word.add_heading(title, level=1)
        for paragraph in paragraphs:
            word.add_paragraph(paragraph)
        word.save(temporary)
        # 必须能被python-docx重新打开，之后才进入最终目录。
        verified = WordDocument(temporary)
        if not verified.paragraphs:
            raise RuntimeError(f"生成的DOCX无法通过内容复验：{title}")
        target.parent.mkdir(parents=True, exist_ok=True)
        temporary.replace(target)
        content = target.read_bytes()
        return (
            relative_path,
            len(content),
            hashlib.sha256(content).hexdigest(),
            target,
        )
    except Exception:
        temporary.unlink(missing_ok=True)
        target.unlink(missing_ok=True)
        raise


def _document(
    *,
    root: Path,
    name: str,
    summary: str,
    document_type: DocumentType,
    status: DocumentStatus,
    security: SecurityLevel,
    visibility: VisibilityType,
    category: Category | None,
    actor: User,
    tags: list[str],
    new_files: list[Path],
    parent: Document | None = None,
    root_document: Document | None = None,
    attachment_type: AttachmentType | None = None,
) -> tuple[Document | None, bool, bool]:
    existing = db.session.scalar(
        select(Document).where(Document.document_name == name)
    )
    if existing is not None:
        if existing.is_deleted:
            return None, False, True
        if existing.document_type != document_type.value:
            raise RuntimeError(f"同名有效文档类型冲突：{name}")
        return existing, False, False
    relative_path, file_size, file_hash, target = _create_docx(
        root,
        title=name,
        paragraphs=[summary, "此文件由本地开发数据初始化工具生成。"],
    )
    new_files.append(target)
    document = Document(
        document_name=name,
        summary=summary,
        document_type=document_type.value,
        document_status=status.value,
        security_level=security.value,
        visibility_type=visibility.value,
        attachment_type=attachment_type.value if attachment_type else None,
        category_id=category.id if category else None,
        category_name=category.category_name if category else None,
        category_path=category.category_path if category else None,
        parent_document_id=parent.id if parent else None,
        root_document_id=root_document.id if root_document else None,
        tags=tags,
        original_file_name=f"{name}.docx",
        file_relative_path=relative_path,
        file_extension="docx",
        mime_type=mimetypes.guess_type(f"{name}.docx")[0] or DOCX_MIME,
        file_size=file_size,
        file_hash=file_hash,
        search_text=" ".join([name, summary, *tags]),
        created_by=actor.id,
        updated_by=actor.id,
        created_by_name=actor.real_name,
        updated_by_name=actor.real_name,
    )
    db.session.add(document)
    db.session.flush()
    return document, True, False


def _permission(
    document: Document,
    *,
    subject_type: SubjectType,
    subject_id: int,
    subject_name: str,
    actor_id: int,
) -> None:
    existing = db.session.scalar(
        select(Permission).where(
            Permission.document_id == document.id,
            Permission.subject_type == subject_type.value,
            Permission.subject_id == subject_id,
        )
    )
    if existing is not None:
        return
    db.session.add(
        Permission(
            document_id=document.id,
            subject_type=subject_type.value,
            subject_id=subject_id,
            subject_name=subject_name,
            can_view=True,
            can_download=True,
            can_edit=False,
            can_manage_permission=False,
            can_delete=False,
            created_by=actor_id,
            updated_by=actor_id,
        )
    )


def _binding(
    main: Document,
    attachment: Document,
    *,
    sort_no: int,
    actor_id: int,
) -> None:
    existing = db.session.scalar(
        select(AttachmentBinding).where(
            AttachmentBinding.main_document_id == main.id,
            AttachmentBinding.attachment_document_id == attachment.id,
        )
    )
    if existing is not None:
        return
    db.session.add(
        AttachmentBinding(
            main_document_id=main.id,
            attachment_document_id=attachment.id,
            sort_no=sort_no,
            created_by=actor_id,
            updated_by=actor_id,
        )
    )


def seed_documents(
    storage_root: str | Path | None = None,
) -> tuple[int, int, int]:
    root = ensure_storage_directories(
        storage_root or current_app.config["DMS_STORAGE_ROOT"]
    )
    actor = _required(
        User,
        User.username == "admin",
        User.is_deleted.is_(False),
        description="有效admin用户",
    )
    ordinary_user = _required(
        User,
        User.username == "user",
        User.is_deleted.is_(False),
        description="有效user用户",
    )
    ops = _required(
        Organization,
        Organization.org_code == "ORG_OPS",
        Organization.is_deleted.is_(False),
        description="ORG_OPS组织",
    )
    categories = {
        code: _required(
            Category,
            Category.category_code == code,
            Category.is_deleted.is_(False),
            Category.status == "ENABLED",
            description=f"{code}分类",
        )
        for code in ("STYLE_A1", "SCENE_B", "SCENE_C")
    }

    created_count = 0
    existing_count = 0
    failed_count = 0
    new_files: list[Path] = []
    try:
        mains: list[Document] = []
        for spec in (
            (
                "2024年度综合应急预案",
                "年度综合应急响应和资源协调方案。",
                "STYLE_A1",
                SecurityLevel.INTERNAL,
                VisibilityType.ALL_AUTHENTICATED,
                ["应急", "年度"],
            ),
            (
                "战略物资管理规程",
                "战略物资储备、调拨和保障管理规程。",
                "SCENE_B",
                SecurityLevel.SECRET,
                VisibilityType.ORGANIZATION,
                ["物资", "保障"],
            ),
            (
                "通信保障方案",
                "重要任务通信组织和应急保障方案。",
                "SCENE_C",
                SecurityLevel.SECRET,
                VisibilityType.CUSTOM,
                ["通信", "保障"],
            ),
        ):
            document, created, failed = _document(
                root=root,
                name=spec[0],
                summary=spec[1],
                document_type=DocumentType.MAIN,
                status=DocumentStatus.PUBLISHED,
                security=spec[3],
                visibility=spec[4],
                category=categories[spec[2]],
                actor=actor,
                tags=spec[5],
                new_files=new_files,
            )
            created_count += int(created)
            existing_count += int(not created and not failed)
            failed_count += int(failed)
            if document is None:
                raise RuntimeError(f"同名文档已逻辑删除，拒绝恢复：{spec[0]}")
            mains.append(document)

        for name, summary in (
            ("人员组织子案", "明确人员编组、岗位职责和协同关系。"),
            ("资源组织子案", "明确资源清单、配置方式和调用机制。"),
            ("重点任务子案", "明确重点任务分解、时序和保障要求。"),
        ):
            document, created, failed = _document(
                root=root,
                name=name,
                summary=summary,
                document_type=DocumentType.SUB_PLAN,
                status=DocumentStatus.PUBLISHED,
                security=SecurityLevel.INTERNAL,
                visibility=VisibilityType.ALL_AUTHENTICATED,
                category=categories["STYLE_A1"],
                actor=actor,
                tags=["子方案"],
                new_files=new_files,
                parent=mains[0],
                root_document=mains[0],
            )
            created_count += int(created)
            existing_count += int(not created and not failed)
            failed_count += int(failed)
            if document is None:
                raise RuntimeError(f"同名文档已逻辑删除，拒绝恢复：{name}")

        attachments: list[Document] = []
        for name, summary, attachment_type in (
            (
                "装备保障工作规范",
                "装备保障作业和协同工作规范。",
                AttachmentType.WORK_STANDARD,
            ),
            (
                "应急资源清单模板",
                "应急资源登记与汇总表模板。",
                AttachmentType.TABLE,
            ),
            (
                "通信保障流程图",
                "通信保障处置流程和协同关系图。",
                AttachmentType.DIAGRAM,
            ),
        ):
            document, created, failed = _document(
                root=root,
                name=name,
                summary=summary,
                document_type=DocumentType.ATTACHMENT,
                status=DocumentStatus.PUBLISHED,
                security=SecurityLevel.PUBLIC,
                visibility=VisibilityType.ALL_AUTHENTICATED,
                category=None,
                actor=actor,
                tags=["共享附件"],
                new_files=new_files,
                attachment_type=attachment_type,
            )
            created_count += int(created)
            existing_count += int(not created and not failed)
            failed_count += int(failed)
            if document is None:
                raise RuntimeError(f"同名附件已逻辑删除，拒绝恢复：{name}")
            attachments.append(document)

        _permission(
            mains[1],
            subject_type=SubjectType.ORG,
            subject_id=ops.id,
            subject_name=ops.org_name,
            actor_id=actor.id,
        )
        _permission(
            mains[2],
            subject_type=SubjectType.USER,
            subject_id=ordinary_user.id,
            subject_name=ordinary_user.real_name,
            actor_id=actor.id,
        )
        _binding(mains[0], attachments[0], sort_no=10, actor_id=actor.id)
        _binding(mains[0], attachments[1], sort_no=20, actor_id=actor.id)
        _binding(mains[1], attachments[0], sort_no=10, actor_id=actor.id)
        db.session.flush()

        for main in mains:
            main.child_count = (
                db.session.scalar(
                    select(func.count(Document.id)).where(
                        Document.parent_document_id == main.id,
                        Document.document_type == DocumentType.SUB_PLAN.value,
                        Document.is_deleted.is_(False),
                    )
                )
                or 0
            )
            main.attachment_count = (
                db.session.scalar(
                    select(func.count(AttachmentBinding.id)).where(
                        AttachmentBinding.main_document_id == main.id,
                        AttachmentBinding.is_deleted.is_(False),
                    )
                )
                or 0
            )
        for category in categories.values():
            category.document_count = (
                db.session.scalar(
                    select(func.count(Document.id)).where(
                        Document.category_id == category.id,
                        Document.document_type.in_(
                            [DocumentType.MAIN.value, DocumentType.SUB_PLAN.value]
                        ),
                        Document.is_deleted.is_(False),
                    )
                )
                or 0
            )
        db.session.commit()
    except Exception:
        db.session.rollback()
        for path in new_files:
            path.unlink(missing_ok=True)
        raise
    return created_count, existing_count, failed_count


def main() -> None:
    app = Flask("dms-seed-documents")
    init_dms(app)
    with app.app_context():
        try:
            created, existing, failed = seed_documents()
        except Exception:
            db.session.rollback()
            print("文档初始化失败：新增0个，已存在0个，失败1个。")
            raise
    print(
        f"文档初始化完成：新增{created}个，已存在{existing}个，失败{failed}个。"
    )


if __name__ == "__main__":
    main()