document-management/backend/dms/services/organization_service.py

"""组织树查询与祖先保留。"""

from __future__ import annotations

from sqlalchemy import select

from dms.common.enums import EnabledStatus
from dms.common.errors import InvalidArgumentError
from dms.common.response import serialize_id
from dms.extensions import db
from dms.models import Organization


def organization_tree(
    *, keyword: str | None, status: str | None
) -> list[dict[str, object]]:
    organizations = db.session.scalars(
        select(Organization)
        .where(Organization.is_deleted.is_(False))
        .order_by(Organization.sort_no.asc(), Organization.id.asc())
    ).all()
    by_id = {item.id: item for item in organizations}

    if status:
        try:
            normalized_status = EnabledStatus(status).value
        except ValueError as exc:
            raise InvalidArgumentError("status必须是ENABLED或DISABLED") from exc
    else:
        normalized_status = EnabledStatus.ENABLED.value

    base_ids = {item.id for item in organizations if item.status == normalized_status}
    if keyword and keyword.strip():
        needle = keyword.strip().casefold()
        base_ids = {
            item.id
            for item in organizations
            if item.id in base_ids
            and (needle in item.org_name.casefold() or needle in item.org_code.casefold())
        }

    included_ids = set(base_ids)
    for item_id in tuple(base_ids):
        parent_id = by_id[item_id].parent_id
        while parent_id is not None and parent_id in by_id:
            included_ids.add(parent_id)
            parent_id = by_id[parent_id].parent_id

    nodes = {
        item.id: {
            "id": serialize_id(item.id),
            "parentId": serialize_id(item.parent_id),
            "orgCode": item.org_code,
            "orgName": item.org_name,
            "orgPath": item.org_path,
            "sortNo": item.sort_no,
            "status": item.status,
            "children": [],
        }
        for item in organizations
        if item.id in included_ids
    }
    roots: list[dict[str, object]] = []
    for item in organizations:
        if item.id not in nodes:
            continue
        node = nodes[item.id]
        if item.parent_id in nodes:
            nodes[item.parent_id]["children"].append(node)
        else:
            roots.append(node)
    return roots