from __future__ import annotations import csv import json import shutil import zipfile from datetime import datetime, timezone from pathlib import Path from tempfile import TemporaryDirectory from typing import Any, Iterable from sqlalchemy import select from app.modules.billing.models import BillingInvoice, BillingInvoiceLine, BillingPayment, BillingInvoiceGenerationBatch from app.modules.clients.models import Client from app.modules.core.tenancy.models import FinancialYear, YearBackupExport from app.modules.documents.models import EngagementDocument, EngagementDocumentVersion from app.modules.notice_cases.models import NoticeCase, NoticeCaseDocument, NoticeCaseEvent, NoticeCaseHearing, NoticeCaseOrder from app.modules.services.models import ClientServiceSubscription, ClientServiceTaskInstance BACKUP_ROOT = Path("data/year_backups") def _safe_name(value: str) -> str: return "".join(ch if ch.isalnum() or ch in {"-", "_", "."} else "_" for ch in str(value or "")).strip("_") or "export" def _serialise(value: Any) -> str: if value is None: return "" if isinstance(value, (datetime,)): return value.isoformat() if hasattr(value, "isoformat"): try: return value.isoformat() except Exception: pass return str(value) def _columns(model: Any) -> list[str]: return [column.name for column in model.__table__.columns] def _write_csv(path: Path, model: Any, rows: Iterable[Any]) -> int: cols = _columns(model) count = 0 with path.open("w", newline="", encoding="utf-8") as fh: writer = csv.DictWriter(fh, fieldnames=cols) writer.writeheader() for row in rows: writer.writerow({col: _serialise(getattr(row, col, None)) for col in cols}) count += 1 return count def _relative_file_candidates(row: Any) -> list[Path]: candidates: list[Path] = [] rel = getattr(row, "local_relative_path", None) if rel: candidates.append(Path(str(rel))) candidates.append(Path("documents") / str(rel)) candidates.append(Path("data") / str(rel)) return candidates def _copy_known_files(staging_dir: Path, rows: Iterable[Any], subfolder: str) -> int: copied = 0 target_root = staging_dir / "files" / subfolder target_root.mkdir(parents=True, exist_ok=True) for row in rows: source = None for candidate in _relative_file_candidates(row): if candidate.exists() and candidate.is_file(): source = candidate break if not source: continue target_name = f"{getattr(row, 'id', 'file')}_{_safe_name(getattr(row, 'original_filename', None) or getattr(row, 'stored_filename', None) or source.name)}" shutil.copy2(source, target_root / target_name) copied += 1 return copied def _rows_by_fy(db, model: Any, *, tenant_id: int, financial_year: str): return db.execute( select(model).where(model.tenant_id == tenant_id, model.financial_year == financial_year) ).scalars().all() def build_year_backup_export(db, *, financial_year: FinancialYear, user_id: int | None) -> YearBackupExport: BACKUP_ROOT.mkdir(parents=True, exist_ok=True) now = datetime.now(timezone.utc) export_code = f"FY_{_safe_name(financial_year.year_code)}_{now.strftime('%Y%m%d_%H%M%S')}" staging_parent = BACKUP_ROOT / "_staging" staging_parent.mkdir(parents=True, exist_ok=True) with TemporaryDirectory(prefix=export_code + "_", dir=str(staging_parent)) as tmp: staging = Path(tmp) tenant_id = int(financial_year.tenant_id) fy = financial_year.year_code subscriptions = _rows_by_fy(db, ClientServiceSubscription, tenant_id=tenant_id, financial_year=fy) tasks = _rows_by_fy(db, ClientServiceTaskInstance, tenant_id=tenant_id, financial_year=fy) engagement_documents = _rows_by_fy(db, EngagementDocument, tenant_id=tenant_id, financial_year=fy) engagement_document_ids = [row.id for row in engagement_documents] engagement_versions = db.execute( select(EngagementDocumentVersion).where(EngagementDocumentVersion.document_id.in_(engagement_document_ids)) ).scalars().all() if engagement_document_ids else [] notice_cases = _rows_by_fy(db, NoticeCase, tenant_id=tenant_id, financial_year=fy) notice_case_ids = [row.id for row in notice_cases] notice_events = db.execute(select(NoticeCaseEvent).where(NoticeCaseEvent.case_id.in_(notice_case_ids))).scalars().all() if notice_case_ids else [] notice_hearings = db.execute(select(NoticeCaseHearing).where(NoticeCaseHearing.case_id.in_(notice_case_ids))).scalars().all() if notice_case_ids else [] notice_orders = db.execute(select(NoticeCaseOrder).where(NoticeCaseOrder.case_id.in_(notice_case_ids))).scalars().all() if notice_case_ids else [] notice_documents = db.execute(select(NoticeCaseDocument).where(NoticeCaseDocument.case_id.in_(notice_case_ids))).scalars().all() if notice_case_ids else [] invoices = _rows_by_fy(db, BillingInvoice, tenant_id=tenant_id, financial_year=fy) invoice_ids = [row.id for row in invoices] invoice_lines = db.execute(select(BillingInvoiceLine).where(BillingInvoiceLine.invoice_id.in_(invoice_ids))).scalars().all() if invoice_ids else [] payments = _rows_by_fy(db, BillingPayment, tenant_id=tenant_id, financial_year=fy) batches = _rows_by_fy(db, BillingInvoiceGenerationBatch, tenant_id=tenant_id, financial_year=fy) client_ids = sorted({row.client_id for row in subscriptions if getattr(row, "client_id", None)} | {row.client_id for row in notice_cases if getattr(row, "client_id", None)} | {row.client_id for row in invoices if getattr(row, "client_id", None)}) clients = db.execute(select(Client).where(Client.id.in_(client_ids))).scalars().all() if client_ids else [] manifest = { "export_code": export_code, "tenant_id": tenant_id, "financial_year": fy, "assessment_year": financial_year.assessment_year, "generated_at_utc": now.isoformat(), "generated_by_user_id": user_id, "record_counts": {}, "file_counts": {}, } datasets = [ ("financial_year.csv", FinancialYear, [financial_year]), ("clients.csv", Client, clients), ("engagements.csv", ClientServiceSubscription, subscriptions), ("tasks.csv", ClientServiceTaskInstance, tasks), ("engagement_documents.csv", EngagementDocument, engagement_documents), ("engagement_document_versions.csv", EngagementDocumentVersion, engagement_versions), ("notice_cases.csv", NoticeCase, notice_cases), ("notice_case_events.csv", NoticeCaseEvent, notice_events), ("notice_case_hearings.csv", NoticeCaseHearing, notice_hearings), ("notice_case_orders.csv", NoticeCaseOrder, notice_orders), ("notice_case_documents.csv", NoticeCaseDocument, notice_documents), ("billing_batches.csv", BillingInvoiceGenerationBatch, batches), ("billing_invoices.csv", BillingInvoice, invoices), ("billing_invoice_lines.csv", BillingInvoiceLine, invoice_lines), ("billing_payments.csv", BillingPayment, payments), ] for filename, model, rows in datasets: manifest["record_counts"][filename] = _write_csv(staging / filename, model, rows) manifest["file_counts"]["engagement_document_versions"] = _copy_known_files(staging, engagement_versions, "engagement_documents") manifest["file_counts"]["notice_case_documents"] = _copy_known_files(staging, notice_documents, "notice_cases") (staging / "manifest.json").write_text(json.dumps(manifest, indent=2, default=str), encoding="utf-8") zip_path = BACKUP_ROOT / f"{export_code}.zip" with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: for file_path in staging.rglob("*"): if file_path.is_file(): zf.write(file_path, file_path.relative_to(staging)) export = YearBackupExport( tenant_id=financial_year.tenant_id, financial_year_id=financial_year.id, year_code=financial_year.year_code, assessment_year=financial_year.assessment_year, export_status="completed", export_file_path=str(zip_path), file_size_bytes=zip_path.stat().st_size if zip_path.exists() else 0, manifest_json=json.dumps(manifest, default=str), generated_by_user_id=user_id, generated_at_utc=now, ) db.add(export) db.flush() return export