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

from datetime import datetime, timezone
from typing import Optional
import json

from celery import Celery

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


def extract_media_metadata(file_url: str, mime_type: str) -> dict:
    """
    Extract metadata from media file.
    
    Args:
        file_url: URL or path to the media file
        mime_type: MIME type of the file
        
    Returns:
        Dictionary containing metadata (duration, width, height, etc.)
    """
    metadata = {}
    
    if mime_type.startswith("video/"):
        try:
            import subprocess
            result = subprocess.run(
                ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams", file_url],
                capture_output=True, text=True, timeout=30
            )
            if result.returncode == 0:
                data = json.loads(result.stdout)
                video_stream = next((s for s in data.get("streams", []) if s["codec_type"] == "video"), None)
                if video_stream:
                    metadata["width"] = video_stream.get("width")
                    metadata["height"] = video_stream.get("height")
                    metadata["codec"] = video_stream.get("codec_name")
                format_info = data.get("format", {})
                metadata["duration"] = float(format_info.get("duration", 0))
                metadata["size_bytes"] = int(format_info.get("size", 0))
        except Exception:
            pass
    elif mime_type.startswith("image/"):
        try:
            from PIL import Image
            import requests
            from io import BytesIO
            response = requests.get(file_url, timeout=10)
            img = Image.open(BytesIO(response.content))
            metadata["width"] = img.width
            metadata["height"] = img.height
            metadata["format"] = img.format
            # Estimate size
            metadata["size_bytes"] = len(response.content)
        except Exception:
            pass
    elif mime_type.startswith("audio/"):
        try:
            import subprocess
            result = subprocess.run(
                ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", file_url],
                capture_output=True, text=True, timeout=30
            )
            if result.returncode == 0:
                data = json.loads(result.stdout)
                format_info = data.get("format", {})
                metadata["duration"] = float(format_info.get("duration", 0))
                metadata["size_bytes"] = int(format_info.get("size", 0))
        except Exception:
            pass
    
    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 storage
        filename = job.storage_key.split("/")[-1]
        mime_type = infer_mime_type_from_storage_key(job.storage_key)
        
        # Get file URL for metadata extraction
        # In production, this would be a presigned URL or internal storage path
        file_url = job.storage_key  # Use storage_key as path for ffprobe
        
        # Extract metadata using appropriate tool
        metadata = extract_media_metadata(file_url, mime_type)
        
        # Fallback for missing metadata
        if not metadata:
            # Log warning but continue with basic asset creation
            metadata = {
                "duration": 0,
                "width": None,
                "height": None,
                "size_bytes": 0,
                "extraction_failed": True
            }

        # Create Asset
        asset = Asset.create(
            workspace_id=job.workspace_id,
            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=metadata.get("size_bytes", 0),
            duration=metadata.get("duration", 0),
            width=metadata.get("width"),
            height=metadata.get("height"),
            status=AssetStatus.READY,
        )
        asset_repo.create(asset)
        db.commit()

        # 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()
        # Update job status to FAILED
        if job:
            job.status = IngestJobStatus.FAILED
            job.error_message = str(e)
            job.updated_at = datetime.now(timezone.utc)
            job_repo.update(job)
            db.commit()

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