145 lines
5.5 KiB
Python
145 lines
5.5 KiB
Python
import io
|
|
import logging
|
|
|
|
from celery import shared_task
|
|
from django.utils import timezone
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
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", {})
|
|
handler = 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):
|
|
logger.info("Virus scan step executed on %d bytes (no scanner configured)", len(data))
|
|
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):
|
|
"""Celery task wrapper around the pipeline executor."""
|
|
from .models import PipelineExecution
|
|
|
|
try:
|
|
execution = PipelineExecution.objects.select_related("file", "pipeline").get(pk=execution_id)
|
|
except PipelineExecution.DoesNotExist:
|
|
logger.warning("Pipeline execution %s not found", execution_id)
|
|
return None
|
|
return PipelineExecutor().execute(execution)
|