import io import logging from collections.abc import Callable from typing import Any from celery import shared_task from django.utils import timezone logger = logging.getLogger(__name__) class PipelineStepRegistry: """Registry of processing steps so apps add steps without editing the kit. Register in ``apps.py:ready()``:: from infrasynth.files.processing import PipelineStepRegistry def strip_exif(data: bytes, mime_type: str, params: dict): ... return data, mime_type PipelineStepRegistry.register("strip_exif", strip_exif) Or decorate: ``@pipeline_step("strip_exif")``. A registered step takes ``(data, mime_type, params)`` and returns ``(data, mime_type)``. """ _steps: dict[str, Callable] = {} @classmethod def register(cls, name: str, func: Callable | None = None): if func is None: def decorator(inner: Callable) -> Callable: cls._steps[name] = inner return inner return decorator cls._steps[name] = func return func @classmethod def get(cls, name: str) -> Callable | None: return cls._steps.get(name) @classmethod def all(cls) -> dict[str, Callable]: return dict(cls._steps) def pipeline_step(name: str): """Decorator form of :meth:`PipelineStepRegistry.register`.""" return PipelineStepRegistry.register(name) def get_pipeline_executor() -> Any: """Returns an executor instance (``INFRASYNTH_FILES["PIPELINE_EXECUTOR"]``).""" from django.utils.module_loading import import_string from infrasynth.shared.settings_utils import get_setting path = get_setting("INFRASYNTH_FILES", "PIPELINE_EXECUTOR", "infrasynth.files.processing.PipelineExecutor") return import_string(path)() class PipelineExecutor: """Executes a processing pipeline over a stored file, step by step.""" def execute(self, execution): from .models import PipelineExecution execution.status = PipelineExecution.Status.RUNNING execution.started_at = timezone.now() execution.save(update_fields=["status", "started_at"]) try: data, mime_type = self._read_source(execution.file) pipeline = execution.pipeline steps = pipeline.steps if pipeline else [] for step in steps: step_type = step.get("type") params = step.get("params", {}) registered = PipelineStepRegistry.get(step_type) handler = registered or getattr(self, f"_step_{step_type}", None) if handler is None: raise ValueError(f"Unknown pipeline step type: '{step_type}'") data, mime_type = handler(data, mime_type, params) output = self._store_output(execution, data, mime_type) execution.status = PipelineExecution.Status.COMPLETED execution.output_file = output execution.completed_at = timezone.now() execution.error = "" execution.save(update_fields=["status", "output_file", "completed_at", "error"]) except Exception as exc: # noqa: BLE001 logger.exception("Pipeline execution %s failed", execution.pk) execution.status = PipelineExecution.Status.FAILED execution.completed_at = timezone.now() execution.error = str(exc) execution.save(update_fields=["status", "completed_at", "error"]) from .signals import file_processed file_processed.send( sender=PipelineExecution, file_id=execution.file_id, pipeline_name=execution.pipeline.slug if execution.pipeline else "", output_file_id=execution.output_file_id, status=execution.status, ) return execution def _read_source(self, stored_file): from .storage import get_storage_backend backend = get_storage_backend(stored_file.storage_backend) fh = backend.open(stored_file.storage_key, "rb") return fh.read(), stored_file.mime_type def _step_resize(self, data: bytes, mime_type: str, params: dict): from PIL import Image img = Image.open(io.BytesIO(data)) width = int(params.get("width", 800)) height = params.get("height") if height: img.thumbnail((width, int(height))) else: img.thumbnail((width, width)) output = io.BytesIO() img.save(output, format=img.format or "PNG") return output.getvalue(), mime_type def _step_optimize(self, data: bytes, mime_type: str, params: dict): from PIL import Image quality = int(params.get("quality", 80)) img = Image.open(io.BytesIO(data)) fmt = img.format or "PNG" if fmt.upper() == "PNG": img = img.convert("P", palette=Image.Palette.ADAPTIVE, colors=256) # type: ignore[assignment] output = io.BytesIO() img.save(output, format="PNG", optimize=True) else: output = io.BytesIO() img.save(output, format=fmt, quality=quality, optimize=True) return output.getvalue(), mime_type def _step_watermark(self, data: bytes, mime_type: str, params: dict): from PIL import Image, ImageDraw, ImageFont text = params.get("text", "Confidential") img = Image.open(io.BytesIO(data)).convert("RGBA") layer = Image.new("RGBA", img.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(layer) try: font = ImageFont.load_default(size=48) except TypeError: font = ImageFont.load_default() width, height = img.size draw.text((width // 4, height // 2), text, fill=(255, 255, 255, 120), font=font) out = Image.alpha_composite(img, layer) output = io.BytesIO() out.save(output, format="PNG") return output.getvalue(), "image/png" def _step_scan(self, data: bytes, mime_type: str, params: dict): from .scanner import get_scanner result = get_scanner().scan(data) if not result.clean: raise ValueError(f"Virus scan rejected the file: {result.threat}") logger.info("Virus scan passed for %d bytes (scanner=%s)", len(data), result.scanner) return data, mime_type def _store_output(self, execution, data: bytes, mime_type: str): from django.core.files.base import ContentFile from .services import FileService original = execution.file name = f"processed_{execution.pipeline.slug}_{original.original_filename}" content_file = ContentFile(data, name=name) content_file.content_type = mime_type # type: ignore[attr-defined] service = FileService() output = service.upload( content_file, filename=name, user=original.uploaded_by, metadata={"source_file_id": original.id, "pipeline": execution.pipeline.slug}, ) return output @shared_task(name="infrasynth.files.run_pipeline_execution", bind=True, max_retries=3) def run_pipeline_execution(self, execution_id, tenant_id=None): """Celery task wrapper around the pipeline executor.""" from infrasynth.tenancy.context import tenant_context from infrasynth.tenancy.models import Tenant from .models import PipelineExecution try: execution = PipelineExecution.all_objects.select_related("file", "pipeline").get(pk=execution_id) except PipelineExecution.DoesNotExist: logger.warning("Pipeline execution %s not found", execution_id) return None tenant = Tenant.objects.filter(pk=tenant_id).first() if tenant_id else execution.tenant with tenant_context(tenant): return get_pipeline_executor().execute(execution)