document-management/backend/dms/seed_categories.py

"""显式、幂等的方案分类开发数据初始化。

执行：python -m dms.seed_categories
仅创建缺失编码，不修改或删除任何已有分类、组织、用户或文档。
"""

from __future__ import annotations

from flask import Flask
from sqlalchemy import select

from dms import init_dms
from dms.common.enums import CategoryType, EnabledStatus
from dms.extensions import db
from dms.models import Category


def _find(category_code: str) -> Category | None:
    return db.session.scalar(
        select(Category).where(Category.category_code == category_code)
    )


def _create(
    *,
    category_code: str,
    category_name: str,
    category_type: CategoryType,
    sort_no: int,
    parent: Category | None,
) -> tuple[Category, bool]:
    existing = _find(category_code)
    if existing is not None:
        return existing, False
    if parent is not None and (parent.is_deleted or parent.status != "ENABLED"):
        raise RuntimeError(f"父分类{parent.category_code}无效，拒绝创建{category_code}")
    category_path = (
        f"{parent.category_path.rstrip('/')}/{category_name}"
        if parent is not None
        else f"/{category_name}"
    )
    category = Category(
        category_code=category_code,
        category_name=category_name,
        category_type=category_type.value,
        parent_id=parent.id if parent is not None else None,
        category_path=category_path,
        sort_no=sort_no,
        document_count=0,
        status=EnabledStatus.ENABLED.value,
    )
    db.session.add(category)
    db.session.flush()
    return category, True


def seed_categories() -> tuple[int, int]:
    created_count = 0
    existing_count = 0

    scene_a, created = _create(
        category_code="SCENE_A",
        category_name="XXXX场景A",
        category_type=CategoryType.SCENE,
        sort_no=10,
        parent=None,
    )
    created_count += int(created)
    existing_count += int(not created)

    for code, name, sort_no in (
        ("STYLE_A1", "XXXX样式A1", 10),
        ("STYLE_A2", "XXXX样式A2", 20),
    ):
        _, created = _create(
            category_code=code,
            category_name=name,
            category_type=CategoryType.STYLE,
            sort_no=sort_no,
            parent=scene_a,
        )
        created_count += int(created)
        existing_count += int(not created)

    for code, name, sort_no in (
        ("SCENE_B", "XXXX场景B", 20),
        ("SCENE_C", "XXXX场景C", 30),
    ):
        _, created = _create(
            category_code=code,
            category_name=name,
            category_type=CategoryType.SCENE,
            sort_no=sort_no,
            parent=None,
        )
        created_count += int(created)
        existing_count += int(not created)

    db.session.commit()
    return created_count, existing_count


def main() -> None:
    app = Flask("dms-seed-categories")
    init_dms(app)
    with app.app_context():
        try:
            created_count, existing_count = seed_categories()
        except Exception:
            db.session.rollback()
            raise
    print(
        f"分类初始化完成：新增{created_count}个，已存在{existing_count}个。"
    )


if __name__ == "__main__":
    main()