"""Celery tasks for the Import Engine. ``process_import_job`` Main worker task. Reads the uploaded CSV or Excel file, validates every row against the registered ``ImportSpec``, calls the entity handler, and updates job progress in real-time. The task is **idempotent**: re-running a job that is already DONE or FAILED is a no-op (returns immediately). A job is claimed by transitioning from PENDING/VALIDATING → PROCESSING before any row is processed; if the transition fails the task aborts to avoid double-runs. """ from __future__ import annotations import csv import io import logging from typing import Any import structlog from celery import shared_task from django.utils import timezone _log = structlog.get_logger("simorgh.platform_core.import_engine") # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _read_rows(file_content: bytes, filename: str) -> list[dict[str, Any]]: """Parse *file_content* into a list of row dicts. Supports: - ``.csv`` — standard Python csv.DictReader - ``.xlsx`` / ``.xls`` — openpyxl (if installed); falls back to csv on ImportError so tests can run without the extra dependency. """ name_lower = filename.lower() if name_lower.endswith(".xlsx") or name_lower.endswith(".xls"): try: import openpyxl # type: ignore[import] wb = openpyxl.load_workbook(io.BytesIO(file_content), read_only=True, data_only=True) ws = wb.active rows = list(ws.iter_rows(values_only=True)) if not rows: return [] headers = [str(h) if h is not None else "" for h in rows[0]] return [dict(zip(headers, row)) for row in rows[1:]] except ImportError: _log.warning("openpyxl not installed — falling back to CSV parser for xlsx") # Fall through to CSV parsing (will likely fail gracefully) text = file_content.decode("utf-8-sig", errors="replace") reader = csv.DictReader(io.StringIO(text)) return list(reader) def _coerce_row(row: dict[str, Any]) -> dict[str, str]: """Normalise all values to stripped strings; drop keys with None value.""" return {k: str(v).strip() for k, v in row.items() if v is not None} # --------------------------------------------------------------------------- # Celery task # --------------------------------------------------------------------------- @shared_task(bind=True, max_retries=3, default_retry_delay=60, name="platform_core.process_import_job") def process_import_job(self, job_id: int) -> dict[str, Any]: """Process a single :class:`~simorgh.apps.platform_core.models.ImportJob`. Returns a summary dict with keys ``status``, ``processed``, ``errors``. """ from simorgh.apps.platform_core.models import ImportJob, ImportJobStatus from simorgh.apps.platform_core.import_registry import import_registry # ------------------------------------------------------------------ # Load job # ------------------------------------------------------------------ try: job = ImportJob.objects.select_related("file", "tenant", "created_by").get(pk=job_id) except ImportJob.DoesNotExist: _log.warning("process_import_job: job not found", job_id=job_id) return {"status": "not_found"} # Idempotency: don't re-process finished jobs if job.status in (ImportJobStatus.DONE, ImportJobStatus.FAILED): _log.info("process_import_job: already finished", job_id=job_id, status=job.status) return {"status": job.status} # ------------------------------------------------------------------ # Transition to PROCESSING (claim) # ------------------------------------------------------------------ now = timezone.now() updated = ImportJob.objects.filter( pk=job_id, status__in=(ImportJobStatus.PENDING, ImportJobStatus.VALIDATING), ).update(status=ImportJobStatus.PROCESSING, started_at=now) if not updated: _log.warning("process_import_job: could not claim job", job_id=job_id) return {"status": "already_running"} job.refresh_from_db() # ------------------------------------------------------------------ # Resolve ImportSpec # ------------------------------------------------------------------ spec = import_registry.get(job.entity_type) if spec is None: _finalize_job(job, status=ImportJobStatus.FAILED, error_report=[ {"row": 0, "errors": [f"No importer registered for entity type '{job.entity_type}'."]} ]) return {"status": "failed", "reason": "no_spec"} # ------------------------------------------------------------------ # Read file # ------------------------------------------------------------------ try: from simorgh.apps.storage.providers import get_provider provider = get_provider(job.file.storage_backend) file_content: bytes = provider.read(job.file.path) except Exception as exc: _log.exception("process_import_job: failed to read file", job_id=job_id, exc_info=exc) _finalize_job(job, status=ImportJobStatus.FAILED, error_report=[ {"row": 0, "errors": [f"Could not read source file: {exc}"]} ]) return {"status": "failed", "reason": "read_error"} rows = _read_rows(file_content, job.file.filename) # Update total_rows ImportJob.objects.filter(pk=job_id).update(total_rows=len(rows)) # ------------------------------------------------------------------ # Process rows # ------------------------------------------------------------------ error_report: list[dict[str, Any]] = [] processed = 0 error_count = 0 for i, raw_row in enumerate(rows, start=1): row = _coerce_row(raw_row) # Validate row_errors = import_registry.validate_row(spec, row, row_number=i) if row_errors: error_report.append({"row": i, "errors": row_errors}) error_count += 1 continue # Call handler try: spec.handler( row, tenant=job.tenant, actor=job.created_by, organization_node_id=job.organization_node_id, ) processed += 1 except Exception as exc: _log.warning( "process_import_job: handler error", job_id=job_id, row=i, exc=str(exc), ) error_report.append({"row": i, "errors": [str(exc)]}) error_count += 1 # Persist progress every 50 rows if i % 50 == 0: ImportJob.objects.filter(pk=job_id).update( processed_rows=processed, error_rows=error_count, ) # ------------------------------------------------------------------ # Finalize # ------------------------------------------------------------------ final_status = ImportJobStatus.DONE if not error_report else ImportJobStatus.DONE # Jobs complete as DONE even with partial errors; only infrastructure # errors (unreadable file / missing spec) mark it as FAILED. _finalize_job( job, status=final_status, processed_rows=processed, error_rows=error_count, error_report=error_report, ) # Dispatch completion event try: from simorgh.apps.events.bus import dispatch dispatch("core.import_job_completed", { "tenant_id": str(job.tenant_id), "job_id": str(job_id), "entity_type": job.entity_type, "processed_rows": processed, "error_rows": error_count, }) except Exception: pass # event dispatch is best-effort return { "status": final_status, "processed": processed, "errors": error_count, } def _finalize_job( job: Any, *, status: str, processed_rows: int = 0, error_rows: int = 0, error_report: list[dict[str, Any]] | None = None, ) -> None: from simorgh.apps.platform_core.models import ImportJob ImportJob.objects.filter(pk=job.pk).update( status=status, processed_rows=processed_rows, error_rows=error_rows, error_report=error_report or [], finished_at=timezone.now(), ) # --------------------------------------------------------------------------- # Export Engine — Celery task # --------------------------------------------------------------------------- _export_log = structlog.get_logger("simorgh.platform_core.export_engine") @shared_task( bind=True, max_retries=3, default_retry_delay=60, name="platform_core.process_export_job", ) def process_export_job(self, job_id: int) -> dict[str, Any]: """Process a single :class:`~simorgh.apps.platform_core.models.ExportJob`. 1. Loads the job and transitions it to PROCESSING. 2. Gets the :class:`~simorgh.apps.platform_core.export_registry.ExportSpec`. 3. Calls ``queryset_fn`` with the job's filters to fetch rows. 4. Serialises rows to CSV / XLSX / JSON bytes using the spec's columns. 5. Saves the result via the storage layer and sets ``job.file``. 6. Marks the job DONE (or FAILED on infrastructure error). """ from simorgh.apps.platform_core.models import ExportJob, ExportJobStatus, ExportFormat from simorgh.apps.platform_core.export_registry import export_registry # ------------------------------------------------------------------ # Load job # ------------------------------------------------------------------ try: job = ExportJob.objects.select_related("tenant", "created_by").get(pk=job_id) except ExportJob.DoesNotExist: _export_log.warning("process_export_job: job not found", job_id=job_id) return {"status": "not_found"} # Idempotency if job.status in (ExportJobStatus.DONE, ExportJobStatus.FAILED): return {"status": job.status} # ------------------------------------------------------------------ # Claim job # ------------------------------------------------------------------ now = timezone.now() updated = ExportJob.objects.filter( pk=job_id, status=ExportJobStatus.PENDING, ).update(status=ExportJobStatus.PROCESSING, started_at=now) if not updated: return {"status": "already_running"} job.refresh_from_db() # ------------------------------------------------------------------ # Resolve ExportSpec # ------------------------------------------------------------------ spec = export_registry.get(job.entity_type) if spec is None: ExportJob.objects.filter(pk=job_id).update( status=ExportJobStatus.FAILED, error_message=f"No exporter registered for entity type '{job.entity_type}'.", finished_at=timezone.now(), ) return {"status": "failed", "reason": "no_spec"} # ------------------------------------------------------------------ # Fetch rows # ------------------------------------------------------------------ try: objects = list( spec.queryset_fn( tenant=job.tenant, actor=job.created_by, filters=job.filters, ) ) except Exception as exc: _export_log.exception("process_export_job: queryset_fn failed", job_id=job_id) ExportJob.objects.filter(pk=job_id).update( status=ExportJobStatus.FAILED, error_message=f"Failed to fetch data: {exc}", finished_at=timezone.now(), ) return {"status": "failed", "reason": "queryset_error"} rows = [export_registry.build_row(spec, obj) for obj in objects] # ------------------------------------------------------------------ # Serialise # ------------------------------------------------------------------ try: file_bytes, filename, content_type = _serialise_rows(rows, spec, job.format) except Exception as exc: _export_log.exception("process_export_job: serialisation failed", job_id=job_id) ExportJob.objects.filter(pk=job_id).update( status=ExportJobStatus.FAILED, error_message=f"Serialisation error: {exc}", finished_at=timezone.now(), ) return {"status": "failed", "reason": "serialise_error"} # ------------------------------------------------------------------ # Save to storage # ------------------------------------------------------------------ try: file_meta = _save_export_file( job=job, file_bytes=file_bytes, filename=filename, content_type=content_type, ) except Exception as exc: _export_log.exception("process_export_job: storage save failed", job_id=job_id) ExportJob.objects.filter(pk=job_id).update( status=ExportJobStatus.FAILED, error_message=f"Storage error: {exc}", finished_at=timezone.now(), ) return {"status": "failed", "reason": "storage_error"} # ------------------------------------------------------------------ # Finalize # ------------------------------------------------------------------ ExportJob.objects.filter(pk=job_id).update( status=ExportJobStatus.DONE, file=file_meta, row_count=len(rows), finished_at=timezone.now(), error_message="", ) try: from simorgh.apps.events.bus import dispatch dispatch("core.export_job_completed", { "tenant_id": str(job.tenant_id), "job_id": str(job_id), "entity_type": job.entity_type, "row_count": len(rows), "format": job.format, }) except Exception: pass return {"status": "done", "row_count": len(rows)} # --------------------------------------------------------------------------- # Export serialisation helpers # --------------------------------------------------------------------------- def _serialise_rows( rows: list[dict[str, Any]], spec: Any, fmt: str, ) -> tuple[bytes, str, str]: """Return (file_bytes, filename, content_type) for *rows* in *fmt*.""" safe_name = spec.entity_type.replace(".", "_").replace("/", "_") if fmt == "csv": return _to_csv(rows, spec), f"{safe_name}_export.csv", "text/csv" elif fmt == "xlsx": return _to_xlsx(rows, spec), f"{safe_name}_export.xlsx", \ "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" elif fmt == "json": return _to_json(rows), f"{safe_name}_export.json", "application/json" else: raise ValueError(f"Unknown export format: {fmt!r}") def _to_csv(rows: list[dict[str, Any]], spec: Any) -> bytes: buf = io.StringIO() writer = csv.DictWriter(buf, fieldnames=spec.column_keys(), extrasaction="ignore") writer.writeheader() writer.writerows(rows) return buf.getvalue().encode("utf-8-sig") # BOM for Excel compatibility def _to_xlsx(rows: list[dict[str, Any]], spec: Any) -> bytes: try: import openpyxl # type: ignore[import] except ImportError: # Fall back to CSV if openpyxl is not installed return _to_csv(rows, spec) wb = openpyxl.Workbook() ws = wb.active ws.title = spec.entity_type[:31] # Excel sheet name limit # Header row headers = spec.column_labels() keys = spec.column_keys() ws.append(headers) for row in rows: ws.append([row.get(k) for k in keys]) buf = io.BytesIO() wb.save(buf) return buf.getvalue() def _to_json(rows: list[dict[str, Any]]) -> bytes: import json as _json return _json.dumps(rows, ensure_ascii=False, default=str).encode("utf-8") def _save_export_file( job: Any, file_bytes: bytes, filename: str, content_type: str, ) -> Any: """Store *file_bytes* and return a :class:`~storage.FileMetadata` instance.""" from simorgh.apps.storage.services import store_file return store_file( filename=filename, content=io.BytesIO(file_bytes), content_type=content_type, tenant_id=job.tenant_id, organization_node_id=job.organization_node_id, uploaded_by_id=job.created_by_id, app_context="export", )