from __future__ import annotations import hashlib import io import json import re 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 ARCHIVE_ZIP_MAX_BYTES = 500 * 1024 * 1024 def sanitize_zip_segment(raw: str) -> str: value = re.sub(r'[\\/:*?"<>|]+', "_", str(raw or "").strip()) value = re.sub(r"\s+", " ", value).strip() return value or "未命名" def short_hash(text: str, *, n: int = 6) -> str: return hashlib.sha1(str(text or "").encode("utf-8")).hexdigest()[:n] def guess_ext_from_url(url: str) -> str: path = unquote(urlparse(str(url or "").split("?", 1)[0]).path) name = path.rsplit("/", 1)[-1] if "." not in name: return "bin" ext = name.rsplit(".", 1)[-1].lower() if not ext or len(ext) > 8 or "/" in ext: return "bin" return ext def build_attachment_filename(row_no: int, label: str, url: str) -> str: return ( f"{int(row_no):02d}_{sanitize_zip_segment(label)}_" f"{short_hash(url)}.{guess_ext_from_url(url)}" ) def first_row_no_by_plan_id(rows: Sequence[Mapping[str, Any]]) -> dict[str, int]: out: dict[str, int] = {} for row in rows: plan_id = str(row.get("planId") or "").strip() if not plan_id or plan_id in out: continue try: out[plan_id] = int(row.get("rowNo") or 0) except (TypeError, ValueError): out[plan_id] = 0 return out def build_hyperlink_formula(rel_path: str, display: str) -> str: path = str(rel_path).replace('"', '""') text = sanitize_zip_segment(display).replace('"', '""') return f'=HYPERLINK("{path}","{text}")' def parse_url_list(raw: Optional[str]) -> list[str]: if not raw: return [] text = str(raw).strip() if not text: return [] try: parsed = json.loads(text) if isinstance(parsed, list): return [str(x).strip() for x in parsed if str(x).strip()] except json.JSONDecodeError: pass return [p.strip() for p in text.split(",") if p.strip()] def collect_hazard_cert_items(products: Sequence[Any]) -> list[dict[str, str]]: out: list[dict[str, str]] = [] for product in products or []: pname = sanitize_zip_segment(getattr(product, "product_name", None) or "") specs = ( ("购买证", getattr(product, "purchase_cert_url", None)), ("运输证", getattr(product, "transport_cert_url", None)), ("合法使用证明", getattr(product, "legal_use_cert_url", None)), ) for base_label, raw in specs: for url in parse_url_list(raw): label = base_label if not pname or pname == "未命名" else f"{base_label}·{pname}" out.append({"url": url, "label": label}) return out def merge_plan_attachment_items( document_items: Sequence[Mapping[str, str]], hazard_items: Sequence[Mapping[str, str]], ) -> list[dict[str, str]]: seen: set[str] = set() out: list[dict[str, str]] = [] for item in list(document_items or []) + list(hazard_items or []): url = str(item.get("url") or "").strip() if not url or url in seen: continue seen.add(url) out.append({"url": url, "label": str(item.get("label") or "附件")}) return out def load_plan_products_by_plan_ids(db, *, tenant_no: str, plan_ids: list[str]): from app.models.pl_models import PlPlanProduct if not plan_ids: return {} rows = ( db.query(PlPlanProduct) .filter( PlPlanProduct.plan_id.in_(plan_ids), PlPlanProduct.tenant_no == tenant_no, PlPlanProduct.del_flag == 0, ) .order_by(PlPlanProduct.order_num, PlPlanProduct.create_time) .all() ) grouped: dict[str, list] = {pid: [] for pid in plan_ids} 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}, )