wuzilongandCursor
81e52477de
Pack archive xlsx plus attachment files under 附件/, with HYPERLINK columns on each plan's first row. Co-authored-by: Cursor <cursoragent@cursor.com>
283 lines
9.4 KiB
Python
283 lines
9.4 KiB
Python
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},
|
||
)
|