401 lines
15 KiB
Python
401 lines
15 KiB
Python
from __future__ import annotations
|
||
|
||
import io
|
||
import json
|
||
import zipfile
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, timezone
|
||
from typing import Any
|
||
from urllib.parse import quote
|
||
|
||
from openpyxl import Workbook
|
||
from openpyxl.worksheet.worksheet import Worksheet
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from src.db.models.app_user import AppUser
|
||
from src.db.models.form_type import FormTypeEnum
|
||
from src.services.budget_form_service import BudgetFormService
|
||
from src.services.export_normalizers import (
|
||
ExcelValueSerializer,
|
||
Form1RowNormalizer,
|
||
Form3ReportRowNormalizer,
|
||
SmetaRowNormalizer,
|
||
)
|
||
from src.services.export_writers import (
|
||
FORM1_FLAT_COLUMNS,
|
||
FORM3_REPORT_COLUMNS,
|
||
SMETA_COLUMNS,
|
||
TabularSheetWriter,
|
||
)
|
||
from src.services.project_service import ProjectService
|
||
from src.services.sheet_service import SheetService
|
||
|
||
XLSX_MEDIA_TYPE = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
||
ZIP_MEDIA_TYPE = "application/zip"
|
||
FORM1_DIRECTION_REQUIRED_SHEETS = {"AHR", "CAP"}
|
||
SHEETS_WITH_SECTIONS = {"AHR", "CAP", "OPER"}
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ExportSheetPlan:
|
||
sheet_name: str
|
||
effective_direction: str | None
|
||
effective_sections: list[str] | None
|
||
validation_direction: str | None
|
||
validation_sections: list[str] | None
|
||
|
||
|
||
class ExportService:
|
||
def __init__(self, db: AsyncSession):
|
||
self.budget_form_service = BudgetFormService(db)
|
||
self.sheet_service = SheetService(db)
|
||
self.project_service = ProjectService(db)
|
||
self.value_serializer = ExcelValueSerializer()
|
||
self.form1_row_normalizer = Form1RowNormalizer(self.value_serializer)
|
||
self.form3_report_row_normalizer = Form3ReportRowNormalizer()
|
||
self.smeta_row_normalizer = SmetaRowNormalizer()
|
||
self.sheet_writers: dict[str, TabularSheetWriter] = {
|
||
"__default__": TabularSheetWriter(
|
||
columns=FORM1_FLAT_COLUMNS,
|
||
row_normalizer=self.form1_row_normalizer.normalize,
|
||
value_serializer=self.value_serializer.to_excel_value,
|
||
),
|
||
"SMETA": TabularSheetWriter(
|
||
columns=SMETA_COLUMNS,
|
||
row_normalizer=self.smeta_row_normalizer.normalize,
|
||
value_serializer=self.value_serializer.to_excel_value,
|
||
),
|
||
}
|
||
self.form3_report_writer = TabularSheetWriter(
|
||
columns=FORM3_REPORT_COLUMNS,
|
||
row_normalizer=self.form3_report_row_normalizer.normalize,
|
||
value_serializer=self.value_serializer.to_excel_value,
|
||
)
|
||
|
||
async def export_form_payload(
|
||
self,
|
||
form_id: int,
|
||
current_user: AppUser,
|
||
sheet: str | None = None,
|
||
direction: str | None = None,
|
||
sections: list[str] | None = None,
|
||
ignore_non_applicable_query_params: bool = False,
|
||
) -> tuple[io.BytesIO, str, str]:
|
||
stream, filename = await self.export_form_to_xlsx(
|
||
form_id=form_id,
|
||
current_user=current_user,
|
||
sheet=sheet,
|
||
direction=direction,
|
||
sections=sections,
|
||
ignore_non_applicable_query_params=ignore_non_applicable_query_params,
|
||
)
|
||
return stream, filename, XLSX_MEDIA_TYPE
|
||
|
||
async def export_project_report_payload(
|
||
self,
|
||
project_id: int,
|
||
year: int,
|
||
report_type: str,
|
||
current_user: AppUser,
|
||
sections: list[str] | None = None,
|
||
) -> tuple[io.BytesIO, str, str]:
|
||
stream, filename = await self.export_project_report_to_xlsx(
|
||
project_id=project_id,
|
||
year=year,
|
||
report_type=report_type,
|
||
current_user=current_user,
|
||
sections=sections,
|
||
)
|
||
return stream, filename, XLSX_MEDIA_TYPE
|
||
|
||
async def export_bulk_payload(
|
||
self,
|
||
form_ids: list[int],
|
||
current_user: AppUser,
|
||
skip_failed: bool = True,
|
||
direction: str | None = None,
|
||
sections: list[str] | None = None,
|
||
ignore_non_applicable_query_params: bool = False,
|
||
) -> tuple[io.BytesIO, str, str]:
|
||
archive_stream = io.BytesIO()
|
||
exported_count = 0
|
||
errors: list[dict[str, Any]] = []
|
||
used_names: set[str] = set()
|
||
|
||
with zipfile.ZipFile(
|
||
archive_stream,
|
||
mode="w",
|
||
compression=zipfile.ZIP_DEFLATED,
|
||
) as archive:
|
||
for form_id in self._deduplicate_ids(form_ids):
|
||
try:
|
||
form_stream, filename = await self.export_form_to_xlsx(
|
||
form_id=form_id,
|
||
current_user=current_user,
|
||
direction=direction,
|
||
sections=sections,
|
||
ignore_non_applicable_query_params=ignore_non_applicable_query_params,
|
||
)
|
||
archive_name = self._make_unique_filename(filename, used_names)
|
||
archive.writestr(archive_name, form_stream.getvalue())
|
||
exported_count += 1
|
||
except ValueError as exc:
|
||
if not skip_failed:
|
||
raise
|
||
errors.append({"form_id": form_id, "error": str(exc)})
|
||
|
||
if errors:
|
||
archive.writestr(
|
||
"errors.json",
|
||
json.dumps({"errors": errors}, ensure_ascii=False, indent=2).encode("utf-8"),
|
||
)
|
||
|
||
if exported_count == 0:
|
||
raise ValueError("Не удалось экспортировать ни одной формы.")
|
||
|
||
archive_stream.seek(0)
|
||
filename = f"export_forms_{datetime.now(timezone.utc).strftime('%Y-%m-%d')}.zip"
|
||
return archive_stream, filename, ZIP_MEDIA_TYPE
|
||
|
||
async def export_form_to_xlsx(
|
||
self,
|
||
form_id: int,
|
||
current_user: AppUser,
|
||
sheet: str | None = None,
|
||
direction: str | None = None,
|
||
sections: list[str] | None = None,
|
||
ignore_non_applicable_query_params: bool = False,
|
||
) -> tuple[io.BytesIO, str]:
|
||
form = await self._get_form_for_export(form_id=form_id, current_user=current_user)
|
||
sheets = self._resolve_requested_sheets(form=form, requested_sheet=sheet)
|
||
self._validate_sections_for_form(form=form, sections=sections)
|
||
plans = self._build_sheet_plans(
|
||
form=form,
|
||
sheets=sheets,
|
||
direction=direction,
|
||
sections=sections,
|
||
ignore_non_applicable_query_params=ignore_non_applicable_query_params,
|
||
)
|
||
|
||
workbook = self._create_workbook_with_metadata(form=form)
|
||
exported_sheets = 0
|
||
for plan in plans:
|
||
self._validate_sheet_query_params(
|
||
form_type_code=form.form_type.code,
|
||
sheet=plan.sheet_name,
|
||
direction=plan.validation_direction,
|
||
sections=plan.validation_sections,
|
||
)
|
||
rows = await self.sheet_service.get(
|
||
form_id=form_id,
|
||
sheet=plan.sheet_name,
|
||
direction=plan.effective_direction,
|
||
sections=plan.effective_sections,
|
||
)
|
||
worksheet = workbook.create_sheet(title=self._safe_sheet_name(plan.sheet_name))
|
||
self._write_sheet_rows(
|
||
worksheet=worksheet,
|
||
rows=rows,
|
||
sheet_name=plan.sheet_name,
|
||
)
|
||
exported_sheets += 1
|
||
|
||
if exported_sheets == 0:
|
||
raise ValueError("Не удалось экспортировать листы формы.")
|
||
|
||
stream = io.BytesIO()
|
||
workbook.save(stream)
|
||
workbook.close()
|
||
stream.seek(0)
|
||
filename = self._build_filename(form_id=form.id, form_type_code=form.form_type.code, year=form.year, sheet=sheet)
|
||
return stream, filename
|
||
|
||
async def export_project_report_to_xlsx(
|
||
self,
|
||
project_id: int,
|
||
year: int,
|
||
report_type: str,
|
||
current_user: AppUser,
|
||
sections: list[str] | None = None,
|
||
) -> tuple[io.BytesIO, str]:
|
||
rows = await self.project_service.get_report_rows(
|
||
project_id=project_id,
|
||
year=year,
|
||
report_type=report_type,
|
||
sections=sections,
|
||
user=current_user,
|
||
)
|
||
report_type_upper = report_type.upper()
|
||
workbook = Workbook()
|
||
meta_sheet = workbook.active
|
||
meta_sheet.title = "metadata"
|
||
metadata_rows = [
|
||
("Тип формы", "FORM_3"),
|
||
("Project ID", project_id),
|
||
("Year", year),
|
||
("Report Type", report_type_upper),
|
||
("Sections", ", ".join(sections) if sections else None),
|
||
("Экспортировано", datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")),
|
||
]
|
||
for row_idx, (key, value) in enumerate(metadata_rows, start=1):
|
||
meta_sheet.cell(row=row_idx, column=1, value=key)
|
||
meta_sheet.cell(row=row_idx, column=2, value=value)
|
||
worksheet = workbook.create_sheet(title=self._safe_sheet_name(report_type_upper))
|
||
self.form3_report_writer.write(worksheet, rows)
|
||
stream = io.BytesIO()
|
||
workbook.save(stream)
|
||
workbook.close()
|
||
stream.seek(0)
|
||
date_str = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||
filename = f"{year}_FORM_3_{project_id}_{report_type_upper}_{date_str}.xlsx"
|
||
return stream, filename
|
||
|
||
@staticmethod
|
||
def build_content_disposition(filename: str) -> str:
|
||
encoded = quote(filename)
|
||
return f"attachment; filename*=UTF-8''{encoded}"
|
||
|
||
def _fill_metadata(self, sheet: Worksheet, form: Any) -> None:
|
||
metadata_rows = [
|
||
("Форма ID", form.id),
|
||
("Тип формы", form.form_type_code),
|
||
("Год", form.year),
|
||
("Org Unit ID", form.org_unit_id),
|
||
("Org Unit", form.org_unit.title if form.org_unit else None),
|
||
("Экспортировано", datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")),
|
||
]
|
||
for row_idx, (key, value) in enumerate(metadata_rows, start=1):
|
||
sheet.cell(row=row_idx, column=1, value=key)
|
||
sheet.cell(row=row_idx, column=2, value=value)
|
||
|
||
def _write_sheet_rows(
|
||
self,
|
||
worksheet: Worksheet,
|
||
rows: list[tuple],
|
||
sheet_name: str,
|
||
) -> None:
|
||
writer = self.sheet_writers.get(sheet_name, self.sheet_writers["__default__"])
|
||
writer.write(worksheet, rows)
|
||
|
||
async def _get_form_for_export(self, form_id: int, current_user: AppUser) -> Any:
|
||
form = await self.budget_form_service.get(
|
||
budget_form_id=form_id,
|
||
user=current_user,
|
||
load_form_type=True,
|
||
load_org=True,
|
||
)
|
||
if not form:
|
||
raise ValueError(f"Форма {form_id} не найдена")
|
||
return form
|
||
|
||
def _resolve_requested_sheets(self, form: Any, requested_sheet: str | None) -> list[str]:
|
||
if requested_sheet and requested_sheet not in form.form_type.sheet_list:
|
||
raise ValueError(f"Лист {requested_sheet} формы {form.id} не найден")
|
||
if requested_sheet:
|
||
return [requested_sheet]
|
||
return list(form.form_type.sheet_list)
|
||
|
||
@staticmethod
|
||
def _validate_sections_for_form(form: Any, sections: list[str] | None) -> None:
|
||
if not sections:
|
||
return
|
||
invalid = sorted(set(sections) - set(form.form_type.section_list))
|
||
if invalid:
|
||
raise ValueError(f"Недопустимые sections: {', '.join(invalid)}")
|
||
|
||
def _create_workbook_with_metadata(self, form: Any) -> Workbook:
|
||
workbook = Workbook()
|
||
meta_sheet = workbook.active
|
||
meta_sheet.title = "metadata"
|
||
self._fill_metadata(meta_sheet, form)
|
||
return workbook
|
||
|
||
def _build_sheet_plans(
|
||
self,
|
||
form: Any,
|
||
sheets: list[str],
|
||
direction: str | None,
|
||
sections: list[str] | None,
|
||
ignore_non_applicable_query_params: bool,
|
||
) -> list[ExportSheetPlan]:
|
||
plans: list[ExportSheetPlan] = []
|
||
for sheet_name in sheets:
|
||
effective_direction = direction if self._requires_direction(form.form_type.code, sheet_name) else None
|
||
effective_sections = sections if sheet_name in SHEETS_WITH_SECTIONS else None
|
||
validation_direction = effective_direction if ignore_non_applicable_query_params else direction
|
||
validation_sections = effective_sections if ignore_non_applicable_query_params else sections
|
||
plans.append(
|
||
ExportSheetPlan(
|
||
sheet_name=sheet_name,
|
||
effective_direction=effective_direction,
|
||
effective_sections=effective_sections,
|
||
validation_direction=validation_direction,
|
||
validation_sections=validation_sections,
|
||
)
|
||
)
|
||
return plans
|
||
|
||
|
||
def _validate_sheet_query_params(
|
||
self,
|
||
form_type_code: FormTypeEnum,
|
||
sheet: str,
|
||
direction: str | None,
|
||
sections: list[str] | None,
|
||
) -> None:
|
||
if form_type_code == FormTypeEnum.FORM_1 and sheet in FORM1_DIRECTION_REQUIRED_SHEETS and not direction:
|
||
raise ValueError("direction обязателен для FORM_1 листов AHR/CAP")
|
||
if direction and not self._requires_direction(form_type_code, sheet):
|
||
raise ValueError("direction допустим только для FORM_1 листов AHR/CAP")
|
||
if sections and sheet not in SHEETS_WITH_SECTIONS:
|
||
raise ValueError(f"sections не поддерживается для листа {sheet}")
|
||
|
||
@staticmethod
|
||
def _requires_direction(form_type_code: FormTypeEnum, sheet: str) -> bool:
|
||
return form_type_code == FormTypeEnum.FORM_1 and sheet in FORM1_DIRECTION_REQUIRED_SHEETS
|
||
|
||
@staticmethod
|
||
def _safe_sheet_name(value: str) -> str:
|
||
safe = value
|
||
for char in ("[", "]", ":", "*", "?", "/", "\\"):
|
||
safe = safe.replace(char, "_")
|
||
safe = safe.strip()
|
||
return (safe or "Sheet")[:31]
|
||
|
||
@staticmethod
|
||
def _build_filename(form_id: int, form_type_code: str, year: int, sheet: str | None = None) -> str:
|
||
date_str = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||
if sheet:
|
||
return f"{year}_{form_type_code}_{form_id}_{sheet}_{date_str}.xlsx"
|
||
return f"{year}_{form_type_code}_{form_id}_{date_str}.xlsx"
|
||
|
||
@staticmethod
|
||
def _deduplicate_ids(form_ids: list[int]) -> list[int]:
|
||
seen: set[int] = set()
|
||
result: list[int] = []
|
||
for form_id in form_ids:
|
||
if form_id not in seen:
|
||
seen.add(form_id)
|
||
result.append(form_id)
|
||
return result
|
||
|
||
@staticmethod
|
||
def _make_unique_filename(filename: str, used_names: set[str]) -> str:
|
||
if filename not in used_names:
|
||
used_names.add(filename)
|
||
return filename
|
||
|
||
if "." in filename:
|
||
base, ext = filename.rsplit(".", 1)
|
||
ext = f".{ext}"
|
||
else:
|
||
base, ext = filename, ""
|
||
|
||
suffix = 2
|
||
while True:
|
||
candidate = f"{base}_{suffix}{ext}"
|
||
if candidate not in used_names:
|
||
used_names.add(candidate)
|
||
return candidate
|
||
suffix += 1
|