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()