shanxia2.0/backend/app/utils/excel_export.py

"""Excel 导出工具：集中处理字典 code → label（对齐 Java ExcelDictConvert / ExcelDictPreloadConvert）。"""
from __future__ import annotations

import io
from dataclasses import dataclass
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence
from urllib.parse import quote

from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.models.models import SysDictData

CellFormatter = Callable[[Any], Any]

EXPORT_MAX_ROWS = 50000


class ExportValidationError(ValueError):
    """导出业务校验失败（空数据 / 超限等），路由层转为 400。"""


@dataclass(frozen=True)
class ExportColumn:
    """导出列定义，dict_type 非空时自动将字典 value 转为 label。"""

    title: str
    field: str
    dict_type: Optional[str] = None
    formatter: Optional[CellFormatter] = None


def load_dict_label_maps(
    db: Session,
    dict_types: Iterable[str],
    tenant_no: Optional[str] = None,
) -> Dict[str, Dict[str, str]]:
    """预加载多个字典类型的 value→label 映射（单次查询，避免 N+1）。"""
    types = sorted({t.strip() for t in dict_types if t and str(t).strip()})
    if not types:
        return {}

    q = db.query(SysDictData).filter(
        SysDictData.del_flag == 0,
        SysDictData.dict_type.in_(types),
    )
    if tenant_no:
        q = q.filter(SysDictData.tenant_no == tenant_no)

    maps: Dict[str, Dict[str, str]] = {t: {} for t in types}
    for row in q.all():
        dtype = row.dict_type
        if dtype not in maps:
            continue
        value = row.dict_value
        if value is None:
            continue
        maps[dtype][str(value)] = row.dict_label or str(value)
    return maps


class DictLabelResolver:
    """导出单元格解析：按列配置做字典翻译或自定义格式化。"""

    def __init__(
        self,
        db: Session,
        columns: Sequence[ExportColumn],
        tenant_no: Optional[str] = None,
        label_maps: Optional[Mapping[str, Mapping[str, str]]] = None,
    ):
        dict_types = [c.dict_type for c in columns if c.dict_type]
        self._maps = dict(label_maps) if label_maps is not None else load_dict_label_maps(
            db, dict_types, tenant_no
        )

    def resolve(self, row: Mapping[str, Any], column: ExportColumn) -> Any:
        raw = row.get(column.field, "")
        if column.formatter is not None:
            return column.formatter(raw)
        if column.dict_type and raw not in (None, ""):
            label = self._maps.get(column.dict_type, {}).get(str(raw))
            if label is not None:
                return label
        if raw is None:
            return ""
        return raw

    def row_to_cells(self, row: Mapping[str, Any], columns: Sequence[ExportColumn]) -> List[Any]:
        return [self.resolve(row, col) for col in columns]


def _build_workbook(
    sheet_title: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    resolver: DictLabelResolver,
):
    from openpyxl import Workbook

    wb = Workbook()
    ws = wb.active
    ws.title = sheet_title
    ws.append([c.title for c in columns])
    for row in rows:
        ws.append(resolver.row_to_cells(row, columns))
    return wb


def workbook_to_streaming_response(
    wb,
    *,
    filename_stem: str,
    include_legacy_filename: bool = True,
) -> StreamingResponse:
    output = io.BytesIO()
    wb.save(output)
    output.seek(0)
    encoded = quote(f"{filename_stem}.xlsx")
    disposition = f"attachment; filename*=UTF-8''{encoded}"
    if include_legacy_filename:
        disposition = f"attachment; filename={encoded}; {disposition}"
    return StreamingResponse(
        output,
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        headers={"Content-Disposition": disposition},
    )


def filter_export_columns(
    all_columns: Sequence[ExportColumn],
    requested: Optional[Sequence[str]],
) -> List[ExportColumn]:
    """按前端可见列 prop 裁剪；未知 field 忽略；requested 为空则返回全部默认列。"""
    if not requested:
        return list(all_columns)
    by_field = {c.field: c for c in all_columns}
    out: List[ExportColumn] = []
    for key in requested:
        col = by_field.get(str(key))
        if col is not None:
            out.append(col)
    return out if out else list(all_columns)


def assert_export_row_count(count: int) -> None:
    if count <= 0:
        raise ExportValidationError("没有可导出的数据")
    if count > EXPORT_MAX_ROWS:
        raise ExportValidationError(f"导出超过 {EXPORT_MAX_ROWS} 条，请缩小筛选范围")


def export_rows_to_xlsx_response(
    db: Session,
    *,
    tenant_no: Optional[str],
    sheet_title: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    filename_stem: str,
    include_legacy_filename: bool = True,
) -> StreamingResponse:
    """CRM 列表导出统一入口：表头 + 行数据 + 字典 label 转换 + 文件流响应。"""
    assert_export_row_count(len(rows))
    if not columns:
        raise ExportValidationError("没有可导出的列")
    resolver = DictLabelResolver(db, columns, tenant_no=tenant_no)
    wb = _build_workbook(sheet_title, columns, rows, resolver)
    return workbook_to_streaming_response(
        wb,
        filename_stem=filename_stem,
        include_legacy_filename=include_legacy_filename,
    )