xiaoxia-saas/apps/worker/worker_app/tasks/ingest.py

import subprocess
from datetime import datetime, timezone
from typing import Optional

from celery import Celery
from celery.app.task import Task
from celery.utils.log import get_task_logger
from worker_app.celery_app import celery_app
from worker_app.core.asset_types import infer_mime_type_from_storage_key
from worker_app.db import SessionLocal

from packages.adapters.sqlalchemy_impl import (
    SQLAlchemyAssetRepository,
    SQLAlchemyIngestJobRepository,
)
from packages.domain import Asset, AssetStatus, IngestJobStatus

logger = get_task_logger(__name__)


def _safe_parse_fps(fps_str: str) -> float:
    """Safely parse fps from a fraction string like \"30/1\" or \"30000/1001\"."""
    try:
        if "/" in fps_str:
            num, den = fps_str.split("/", 1)
            den_val = float(den)
            if den_val == 0:
                return 0.0
            return float(num) / den_val
        return float(fps_str)
    except (ValueError, ZeroDivisionError):
        return 0.0


def extract_media_metadata(file_url: str, media_type: str) -> dict:
    """
    提取媒体文件的元数据。

    Args:
        file_url: 媒体文件 URL 或本地路径
        media_type: 媒体类型 (video, audio, image)

    Returns:
        提取的元数据字典，失败时返回空字典
    """
    metadata = {}

    try:
        if media_type == "video":
            # 使用 ffprobe 提取视频元数据
            cmd = [
                "ffprobe",
                "-v",
                "quiet",
                "-print_format",
                "json",
                "-show_format",
                "-show_streams",
                file_url,
            ]
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=30,
            )
            if result.returncode == 0:
                import json as json_lib

                probe_data = json_lib.loads(result.stdout)

                # 提取视频流信息
                for stream in probe_data.get("streams", []):
                    if stream.get("codec_type") == "video":
                        metadata["width"] = int(stream.get("width", 0))
                        metadata["height"] = int(stream.get("height", 0))
                        metadata["codec"] = stream.get("codec_name", "")
                        metadata["fps"] = (
                            _safe_parse_fps(stream.get("r_frame_rate", "0/1")) if stream.get("r_frame_rate") else 0
                        )
                        break

                # 提取格式信息
                format_info = probe_data.get("format", {})
                metadata["duration"] = float(format_info.get("duration", 0))
                metadata["size_bytes"] = int(format_info.get("size", 0))
                metadata["bitrate"] = int(format_info.get("bit_rate", 0))

        elif media_type == "image":
            # 使用 Pillow 提取图片元数据
            try:
                from PIL import Image

                with Image.open(file_url) as img:
                    metadata["width"] = img.width
                    metadata["height"] = img.height
                    metadata["format"] = img.format
                    metadata["mode"] = img.mode
                    if hasattr(img, "_getexif") and img._getexif():
                        exif = img._getexif()
                        if exif:
                            metadata["exif"] = {k: str(v) for k, v in exif.items() if isinstance(v, (str, int, float))}
            except ImportError:
                logger.warning("Pillow not available for image metadata extraction")
            except Exception as e:
                logger.warning(f"Failed to extract image metadata: {e}")

    except subprocess.TimeoutExpired:
        logger.warning(f"Timeout extracting metadata from {file_url}")
    except FileNotFoundError:
        logger.warning("ffprobe not found, cannot extract video metadata")
    except Exception as e:
        logger.warning(f"Failed to extract metadata: {e}")

    return metadata


@celery_app.task(name="worker.ingest_asset")
def ingest_asset(job_id: str) -> dict:
    """
    Ingest asset task.

    Steps:
    1. Fetch IngestJob from repository
    2. Extract metadata from storage_key
    3. Create Asset entity
    4. Update IngestJob status to COMPLETED
    5. Return result
    """
    db = SessionLocal()
    try:
        job_repo = SQLAlchemyIngestJobRepository(db)
        asset_repo = SQLAlchemyAssetRepository(db)

        job = job_repo.get(job_id)
        if job is None:
            return {"status": "failed", "error": "job not found"}

        # Update job status to PROCESSING
        job.status = IngestJobStatus.PROCESSING
        job.updated_at = datetime.now(timezone.utc)
        job_repo.update(job)
        db.commit()

        # Extract real metadata from media file
        filename = job.storage_key.split("/")[-1]
        mime_type = infer_mime_type_from_storage_key(job.storage_key)

        # Determine media type from mime_type
        media_type = "video"
        if mime_type.startswith("image/"):
            media_type = "image"
        elif mime_type.startswith("audio/"):
            media_type = "audio"

        # Extract metadata (returns empty dict on failure)
        storage_url = job.storage_key  # Assuming storage_key is usable as URL/path
        metadata = extract_media_metadata(storage_url, media_type)

        # Fill in defaults if metadata extraction failed
        if not metadata:
            metadata = {
                "duration": 0,
                "width": 0,
                "height": 0,
                "size_bytes": 0,
            }

        # Create Asset
        asset = Asset.create(
            project_id=job.project_id,
            library_id=job.library_id,
            name=filename,
            storage_key=job.storage_key,
            mime_type=mime_type,
            metadata=metadata,
            file_size=int(metadata.get("size_bytes", 0)),
            duration=float(metadata.get("duration", 0)),
            width=int(metadata.get("width", 0)),
            height=int(metadata.get("height", 0)),
            status=AssetStatus.READY,
            file_hash=job.file_hash,
        )
        asset_repo.create(asset)

        # Update job status to COMPLETED
        job.status = IngestJobStatus.COMPLETED
        job.result_asset_id = asset.id
        job.updated_at = datetime.now(timezone.utc)
        job_repo.update(job)

        db.commit()

        return {
            "status": "completed",
            "job_id": job.id,
            "asset_id": asset.id,
        }
    except Exception as e:
        db.rollback()
        logger.error(f"Failed to ingest asset {job_id}: {e}")

        # Update job status to FAILED
        try:
            job_repo = SQLAlchemyIngestJobRepository(db)
            job = job_repo.get(job_id)
            if job:
                job.status = IngestJobStatus.FAILED
                job.error_message = str(e)
                job.updated_at = datetime.now(timezone.utc)
                job_repo.update(job)
                db.commit()
        except Exception:
            db.rollback()

        return {
            "status": "failed",
            "job_id": job_id,
            "error": str(e),
        }
    finally:
        db.close()