diff --git a/backend/app/services/pl_plan_archive_zip_export.py b/backend/app/services/pl_plan_archive_zip_export.py index 83ac9e687..b910333be 100644 --- a/backend/app/services/pl_plan_archive_zip_export.py +++ b/backend/app/services/pl_plan_archive_zip_export.py @@ -1,10 +1,17 @@ from __future__ import annotations import hashlib +import io import json import re -from typing import Any, Mapping, Optional, Sequence -from urllib.parse import unquote, urlparse +import zipfile +from typing import Any, Callable, Mapping, Optional, Sequence +from urllib.parse import quote, unquote, urlparse + +from fastapi.responses import StreamingResponse +from sqlalchemy.orm import Session + +from app.utils.excel_export import DictLabelResolver, ExportColumn, ExportValidationError ARCHIVE_ZIP_ATTACH_COL_MAX = 20 ARCHIVE_ZIP_MAX_FILES = 2000 @@ -123,3 +130,153 @@ def load_plan_products_by_plan_ids(db, *, tenant_no: str, plan_ids: list[str]): for row in rows: grouped.setdefault(row.plan_id, []).append(row) return grouped + + +def _default_download(url: str) -> bytes: + from app.services.bx.storage import download_file_bytes + + return download_file_bytes(url) + + +def _attach_field(index: int) -> str: + return f"attach{index}" + + +def build_archive_zip_bytes( + *, + db: Optional[Session], + tenant_no: str, + columns: Sequence[ExportColumn], + rows: Sequence[Mapping[str, Any]], + attachments_by_plan: Mapping[str, Sequence[Mapping[str, str]]], + stamp: str, + download_fn: Optional[Callable[[str], bytes]] = None, +) -> tuple[bytes, list[str]]: + """打包归档 ZIP：xlsx + 附件/；返回 (zip_bytes, failure_messages)。""" + download = download_fn or _default_download + row_no_by_plan = first_row_no_by_plan_id(rows) + root = f"计划归档({stamp})" + attach_dir = f"{root}/附件" + + # plan_id -> list of (rel_path_in_zip_under_附件, label, formula_rel) + plan_file_meta: dict[str, list[tuple[str, str, str]]] = {} + zip_file_entries: list[tuple[str, bytes]] = [] + failures: list[str] = [] + total_bytes = 0 + file_count = 0 + + for plan_id, items in (attachments_by_plan or {}).items(): + seq = row_no_by_plan.get(plan_id) or 0 + if seq <= 0: + continue + meta: list[tuple[str, str, str]] = [] + for item in items or []: + url = str(item.get("url") or "").strip() + if not url: + continue + label = str(item.get("label") or "附件") + filename = build_attachment_filename(seq, label, url) + zip_path = f"{attach_dir}/{filename}" + try: + content = download(url) + except Exception as exc: # noqa: BLE001 — 单文件失败不阻断整包 + failures.append(f"{label}: {url} ({exc})") + continue + file_count += 1 + total_bytes += len(content) + if file_count > ARCHIVE_ZIP_MAX_FILES: + raise ExportValidationError( + f"附件文件数超过 {ARCHIVE_ZIP_MAX_FILES}，请缩小筛选范围" + ) + if total_bytes > ARCHIVE_ZIP_MAX_BYTES: + raise ExportValidationError( + "附件总体积过大，请缩小筛选范围后重试" + ) + zip_file_entries.append((zip_path, content)) + formula_rel = f"附件/{filename}" + meta.append((zip_path, label, formula_rel)) + if meta: + plan_file_meta[plan_id] = meta + + max_att = 0 + for meta in plan_file_meta.values(): + max_att = max(max_att, len(meta)) + k = min(ARCHIVE_ZIP_ATTACH_COL_MAX, max_att) + + attach_columns = [ + ExportColumn(title=f"附件_{i}", field=_attach_field(i)) for i in range(1, k + 1) + ] + all_columns = list(columns) + attach_columns + + enriched_rows: list[dict[str, Any]] = [] + seen_first: set[str] = set() + for row in rows: + data = dict(row) + plan_id = str(data.get("planId") or "").strip() + for i in range(1, k + 1): + data[_attach_field(i)] = "" + if plan_id and plan_id not in seen_first: + seen_first.add(plan_id) + meta = plan_file_meta.get(plan_id) or [] + for i, (_zip_path, label, formula_rel) in enumerate(meta[:k], start=1): + data[_attach_field(i)] = build_hyperlink_formula(formula_rel, label) + enriched_rows.append(data) + + if not enriched_rows: + raise ExportValidationError("没有可导出的数据") + if not all_columns: + raise ExportValidationError("没有可导出的列") + + resolver = DictLabelResolver(db, all_columns, tenant_no=tenant_no) + from openpyxl import Workbook + + wb = Workbook() + ws = wb.active + ws.title = "计划归档" + ws.append([c.title for c in all_columns]) + for row in enriched_rows: + cells = resolver.row_to_cells(row, all_columns) + ws.append(cells) + + xlsx_buf = io.BytesIO() + wb.save(xlsx_buf) + xlsx_bytes = xlsx_buf.getvalue() + + out = io.BytesIO() + with zipfile.ZipFile(out, mode="w", compression=zipfile.ZIP_DEFLATED) as zf: + zf.writestr(f"{root}/(总表)计划归档.xlsx", xlsx_bytes) + for path, content in zip_file_entries: + zf.writestr(path, content) + if failures: + note = "以下附件下载失败，已跳过：\n" + "\n".join(failures) + zf.writestr(f"{root}/导出说明.txt", note.encode("utf-8")) + return out.getvalue(), failures + + +def build_archive_zip_streaming_response( + *, + db: Session, + tenant_no: str, + columns: Sequence[ExportColumn], + rows: Sequence[Mapping[str, Any]], + attachments_by_plan: Mapping[str, Sequence[Mapping[str, str]]], + stamp: str, + download_fn: Optional[Callable[[str], bytes]] = None, +) -> StreamingResponse: + data, _failures = build_archive_zip_bytes( + db=db, + tenant_no=tenant_no, + columns=columns, + rows=rows, + attachments_by_plan=attachments_by_plan, + stamp=stamp, + download_fn=download_fn, + ) + filename = f"计划归档({stamp}).zip" + encoded = quote(filename) + disposition = f"attachment; filename={encoded}; filename*=UTF-8''{encoded}" + return StreamingResponse( + io.BytesIO(data), + media_type="application/zip", + headers={"Content-Disposition": disposition}, + ) diff --git a/backend/tests/test_pl_plan_archive_zip_export.py b/backend/tests/test_pl_plan_archive_zip_export.py index f87ac5359..45cdcea07 100644 --- a/backend/tests/test_pl_plan_archive_zip_export.py +++ b/backend/tests/test_pl_plan_archive_zip_export.py @@ -1,3 +1,8 @@ +import io +import zipfile + +from openpyxl import load_workbook + from app.services.pl_plan_archive_zip_export import ( sanitize_zip_segment, build_attachment_filename, @@ -6,7 +11,9 @@ from app.services.pl_plan_archive_zip_export import ( parse_url_list, collect_hazard_cert_items, merge_plan_attachment_items, + build_archive_zip_bytes, ) +from app.utils.excel_export import ExportColumn def test_sanitize_zip_segment(): @@ -60,3 +67,46 @@ def test_merge_dedupes_url(): merged = merge_plan_attachment_items(doc, haz) assert len(merged) == 1 assert merged[0]["label"] == "司机" + + +def test_build_archive_zip_bytes_structure_and_hyperlink(): + columns = [ + ExportColumn(title="序号", field="rowNo"), + ExportColumn(title="计划编号", field="planNo"), + ] + rows = [ + {"planId": "p1", "rowNo": 1, "planNo": "PL001"}, + {"planId": "p1", "rowNo": 2, "planNo": "PL001"}, + ] + attachments_by_plan = { + "p1": [{"url": "https://cdn.example/a.jpg", "label": "司机"}], + } + + def fake_download(url: str) -> bytes: + return b"fake-image-bytes" + + stamp = "202609231230" + data, failures = build_archive_zip_bytes( + db=None, + tenant_no="t1", + columns=columns, + rows=rows, + attachments_by_plan=attachments_by_plan, + stamp=stamp, + download_fn=fake_download, + ) + assert failures == [] + with zipfile.ZipFile(io.BytesIO(data)) as zf: + names = zf.namelist() + assert f"计划归档({stamp})/(总表)计划归档.xlsx" in names + att = [n for n in names if n.startswith(f"计划归档({stamp})/附件/")] + assert len(att) == 1 + assert "01_司机_" in att[0] + xlsx_bytes = zf.read(f"计划归档({stamp})/(总表)计划归档.xlsx") + wb = load_workbook(io.BytesIO(xlsx_bytes)) + ws = wb.active + headers = [c.value for c in ws[1]] + assert "附件_1" in headers + attach_col = headers.index("附件_1") + 1 + assert str(ws.cell(2, attach_col).value).startswith("=HYPERLINK(") + assert ws.cell(3, attach_col).value in (None, "")