384 lines
12 KiB
Python
384 lines
12 KiB
Python
"""列表导出：Excel 总表 + 附件目录 ZIP（计划归档 / 充值申请等共用）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import io
|
||
import json
|
||
import re
|
||
import zipfile
|
||
from typing import Any, Callable, Iterator, Mapping, Optional, Sequence, Tuple
|
||
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
|
||
|
||
EXPORT_ZIP_ATTACH_COL_MAX = 20
|
||
EXPORT_ZIP_MAX_FILES = 2000
|
||
EXPORT_ZIP_MAX_BYTES = 500 * 1024 * 1024
|
||
|
||
# 兼容旧常量名
|
||
ARCHIVE_ZIP_ATTACH_COL_MAX = EXPORT_ZIP_ATTACH_COL_MAX
|
||
ARCHIVE_ZIP_MAX_FILES = EXPORT_ZIP_MAX_FILES
|
||
ARCHIVE_ZIP_MAX_BYTES = EXPORT_ZIP_MAX_BYTES
|
||
|
||
|
||
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_entity_id(
|
||
rows: Sequence[Mapping[str, Any]], *, id_field: str
|
||
) -> dict[str, int]:
|
||
out: dict[str, int] = {}
|
||
for row in rows:
|
||
entity_id = str(row.get(id_field) or "").strip()
|
||
if not entity_id or entity_id in out:
|
||
continue
|
||
try:
|
||
out[entity_id] = int(row.get("rowNo") or 0)
|
||
except (TypeError, ValueError):
|
||
out[entity_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 urls_to_attachment_items(
|
||
raw: Optional[str], *, label_prefix: str = "附件"
|
||
) -> list[dict[str, str]]:
|
||
"""将逗号/JSON URL 列表转为 {url, label} 列表。"""
|
||
out: list[dict[str, str]] = []
|
||
for i, url in enumerate(parse_url_list(raw), start=1):
|
||
out.append({"url": url, "label": f"{label_prefix}{i}"})
|
||
return out
|
||
|
||
|
||
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 _plan_attachment_jobs(
|
||
*,
|
||
rows: Sequence[Mapping[str, Any]],
|
||
attachments_by_id: Mapping[str, Sequence[Mapping[str, str]]],
|
||
id_field: str,
|
||
attach_dir: str,
|
||
) -> Tuple[dict[str, list[tuple[str, str, str]]], list[tuple[str, str, str]]]:
|
||
"""预先确定附件 zip 路径与超链接（尚未下载）。
|
||
|
||
返回:
|
||
entity_file_meta: entity_id -> [(zip_path, label, formula_rel), ...]
|
||
download_jobs: [(zip_path, label, url), ...] 按写入顺序
|
||
"""
|
||
row_no_by_id = first_row_no_by_entity_id(rows, id_field=id_field)
|
||
entity_file_meta: dict[str, list[tuple[str, str, str]]] = {}
|
||
download_jobs: list[tuple[str, str, str]] = []
|
||
planned = 0
|
||
|
||
for entity_id, items in (attachments_by_id or {}).items():
|
||
seq = row_no_by_id.get(entity_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
|
||
planned += 1
|
||
if planned > EXPORT_ZIP_MAX_FILES:
|
||
raise ExportValidationError(
|
||
f"附件文件数超过 {EXPORT_ZIP_MAX_FILES}，请缩小筛选范围"
|
||
)
|
||
item_label = str(item.get("label") or "附件")
|
||
filename = build_attachment_filename(seq, item_label, url)
|
||
zip_path = f"{attach_dir}/{filename}"
|
||
formula_rel = f"附件/{filename}"
|
||
meta.append((zip_path, item_label, formula_rel))
|
||
download_jobs.append((zip_path, item_label, url))
|
||
if meta:
|
||
entity_file_meta[entity_id] = meta
|
||
return entity_file_meta, download_jobs
|
||
|
||
|
||
def _build_xlsx_bytes(
|
||
*,
|
||
db: Optional[Session],
|
||
tenant_no: str,
|
||
columns: Sequence[ExportColumn],
|
||
rows: Sequence[Mapping[str, Any]],
|
||
entity_file_meta: Mapping[str, Sequence[tuple[str, str, str]]],
|
||
id_field: str,
|
||
sheet: str,
|
||
) -> bytes:
|
||
max_att = 0
|
||
for meta in entity_file_meta.values():
|
||
max_att = max(max_att, len(meta))
|
||
k = min(EXPORT_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)
|
||
entity_id = str(data.get(id_field) or "").strip()
|
||
for i in range(1, k + 1):
|
||
data[_attach_field(i)] = ""
|
||
if entity_id and entity_id not in seen_first:
|
||
seen_first.add(entity_id)
|
||
meta = list(entity_file_meta.get(entity_id) or [])
|
||
for i, (_zip_path, item_label, formula_rel) in enumerate(meta[:k], start=1):
|
||
data[_attach_field(i)] = build_hyperlink_formula(formula_rel, item_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 = sheet
|
||
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)
|
||
return xlsx_buf.getvalue()
|
||
|
||
|
||
def build_attachment_zip_bytes(
|
||
*,
|
||
db: Optional[Session],
|
||
tenant_no: str,
|
||
columns: Sequence[ExportColumn],
|
||
rows: Sequence[Mapping[str, Any]],
|
||
attachments_by_id: Mapping[str, Sequence[Mapping[str, str]]],
|
||
stamp: str,
|
||
package_label: str,
|
||
id_field: str = "id",
|
||
sheet_title: Optional[str] = None,
|
||
xlsx_basename: Optional[str] = None,
|
||
download_fn: Optional[Callable[[str], bytes]] = None,
|
||
) -> tuple[bytes, list[str]]:
|
||
"""打包 ZIP：xlsx + 附件/；返回 (zip_bytes, failure_messages)。"""
|
||
failures: list[str] = []
|
||
chunks = list(
|
||
iter_attachment_zip_chunks(
|
||
db=db,
|
||
tenant_no=tenant_no,
|
||
columns=columns,
|
||
rows=rows,
|
||
attachments_by_id=attachments_by_id,
|
||
stamp=stamp,
|
||
package_label=package_label,
|
||
id_field=id_field,
|
||
sheet_title=sheet_title,
|
||
xlsx_basename=xlsx_basename,
|
||
download_fn=download_fn,
|
||
failures_out=failures,
|
||
)
|
||
)
|
||
return b"".join(chunks), failures
|
||
|
||
|
||
def iter_attachment_zip_chunks(
|
||
*,
|
||
db: Optional[Session],
|
||
tenant_no: str,
|
||
columns: Sequence[ExportColumn],
|
||
rows: Sequence[Mapping[str, Any]],
|
||
attachments_by_id: Mapping[str, Sequence[Mapping[str, str]]],
|
||
stamp: str,
|
||
package_label: str,
|
||
id_field: str = "id",
|
||
sheet_title: Optional[str] = None,
|
||
xlsx_basename: Optional[str] = None,
|
||
download_fn: Optional[Callable[[str], bytes]] = None,
|
||
failures_out: Optional[list[str]] = None,
|
||
) -> Iterator[bytes]:
|
||
"""边下载边产出 ZIP 字节，便于代理层在长时间打包时保持连接。"""
|
||
download = download_fn or _default_download
|
||
label = sanitize_zip_segment(package_label) or "导出"
|
||
sheet = (sheet_title or label)[:31] or "导出"
|
||
xlsx_name = xlsx_basename or f"(总表){label}.xlsx"
|
||
root = f"{label}({stamp})"
|
||
attach_dir = f"{root}/附件"
|
||
|
||
entity_file_meta, download_jobs = _plan_attachment_jobs(
|
||
rows=rows,
|
||
attachments_by_id=attachments_by_id,
|
||
id_field=id_field,
|
||
attach_dir=attach_dir,
|
||
)
|
||
xlsx_bytes = _build_xlsx_bytes(
|
||
db=db,
|
||
tenant_no=tenant_no,
|
||
columns=columns,
|
||
rows=rows,
|
||
entity_file_meta=entity_file_meta,
|
||
id_field=id_field,
|
||
sheet=sheet,
|
||
)
|
||
|
||
failures: list[str] = []
|
||
total_bytes = 0
|
||
file_count = 0
|
||
buf = io.BytesIO()
|
||
sent = 0
|
||
|
||
def flush() -> bytes:
|
||
nonlocal sent
|
||
data = buf.getvalue()
|
||
if len(data) <= sent:
|
||
return b""
|
||
chunk = data[sent:]
|
||
sent = len(data)
|
||
return chunk
|
||
|
||
with zipfile.ZipFile(buf, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
|
||
zf.writestr(f"{root}/{xlsx_name}", xlsx_bytes)
|
||
chunk = flush()
|
||
if chunk:
|
||
yield chunk
|
||
|
||
for zip_path, item_label, url in download_jobs:
|
||
try:
|
||
content = download(url)
|
||
except Exception as exc: # noqa: BLE001 — 单文件失败不阻断整包
|
||
failures.append(f"{item_label}: {url} ({exc})")
|
||
continue
|
||
file_count += 1
|
||
total_bytes += len(content)
|
||
if file_count > EXPORT_ZIP_MAX_FILES:
|
||
raise ExportValidationError(
|
||
f"附件文件数超过 {EXPORT_ZIP_MAX_FILES}，请缩小筛选范围"
|
||
)
|
||
if total_bytes > EXPORT_ZIP_MAX_BYTES:
|
||
raise ExportValidationError("附件总体积过大，请缩小筛选范围后重试")
|
||
zf.writestr(zip_path, content)
|
||
chunk = flush()
|
||
if chunk:
|
||
yield chunk
|
||
|
||
if failures:
|
||
note = "以下附件下载失败，已跳过：\n" + "\n".join(failures)
|
||
zf.writestr(f"{root}/导出说明.txt", note.encode("utf-8"))
|
||
chunk = flush()
|
||
if chunk:
|
||
yield chunk
|
||
|
||
# ZipFile 关闭后写入中央目录
|
||
chunk = flush()
|
||
if chunk:
|
||
yield chunk
|
||
|
||
if failures_out is not None:
|
||
failures_out.extend(failures)
|
||
|
||
|
||
def build_attachment_zip_streaming_response(
|
||
*,
|
||
db: Session,
|
||
tenant_no: str,
|
||
columns: Sequence[ExportColumn],
|
||
rows: Sequence[Mapping[str, Any]],
|
||
attachments_by_id: Mapping[str, Sequence[Mapping[str, str]]],
|
||
stamp: str,
|
||
package_label: str,
|
||
id_field: str = "id",
|
||
sheet_title: Optional[str] = None,
|
||
xlsx_basename: Optional[str] = None,
|
||
download_fn: Optional[Callable[[str], bytes]] = None,
|
||
) -> StreamingResponse:
|
||
"""立即返回 StreamingResponse：先发响应头，再边下附件边推送 ZIP 字节。"""
|
||
label = sanitize_zip_segment(package_label) or "导出"
|
||
filename = f"{label}({stamp}).zip"
|
||
encoded = quote(filename)
|
||
disposition = f"attachment; filename={encoded}; filename*=UTF-8''{encoded}"
|
||
|
||
def generate() -> Iterator[bytes]:
|
||
yield from iter_attachment_zip_chunks(
|
||
db=db,
|
||
tenant_no=tenant_no,
|
||
columns=columns,
|
||
rows=rows,
|
||
attachments_by_id=attachments_by_id,
|
||
stamp=stamp,
|
||
package_label=package_label,
|
||
id_field=id_field,
|
||
sheet_title=sheet_title,
|
||
xlsx_basename=xlsx_basename,
|
||
download_fn=download_fn,
|
||
)
|
||
|
||
return StreamingResponse(
|
||
generate(),
|
||
media_type="application/zip",
|
||
headers={
|
||
"Content-Disposition": disposition,
|
||
# 避免中间代理整包缓冲导致长时间无字节
|
||
"X-Accel-Buffering": "no",
|
||
},
|
||
)
|