from sqlalchemy.orm import Session from packages.adapters.sqlalchemy_impl.models import GenerationTaskModel from packages.domain import GenerationTask from packages.domain.generation_task import GenerationTaskStatus def _to_domain(model: GenerationTaskModel) -> GenerationTask: """Convert ORM model to domain entity.""" return GenerationTask( id=model.id, project_id=model.project_id, strategy_id=model.strategy_id, asset_library_id=model.asset_library_id, voice_library_id=model.voice_library_id, template_id=model.template_id, asset_ids=list(model.asset_ids or []), title_ids=list(model.title_ids or []), voice_ids=list(model.voice_ids or []), status=GenerationTaskStatus(model.status) if model.status else GenerationTaskStatus.PENDING, progress=model.progress, result_count=int(model.result_count or 0), error_message=model.error_message, started_at=model.started_at, completed_at=model.completed_at, created_by_user_id=model.created_by_user_id, source_edit_plan_id=model.source_edit_plan_id or "", asset_select_mode=model.asset_select_mode or "", batch_id=model.batch_id or "", created_at=model.created_at, ) class SQLAlchemyGenerationTaskRepository: def __init__(self, session: Session): self.session = session def create(self, task: GenerationTask) -> GenerationTask: model = GenerationTaskModel( id=task.id, project_id=task.project_id, strategy_id=task.strategy_id, asset_library_id=task.asset_library_id, voice_library_id=task.voice_library_id, template_id=task.template_id, asset_ids=task.asset_ids, title_ids=task.title_ids, voice_ids=task.voice_ids, status=task.status, progress=task.progress, result_count=task.result_count, error_message=task.error_message, started_at=task.started_at, completed_at=task.completed_at, created_by_user_id=task.created_by_user_id, source_edit_plan_id=task.source_edit_plan_id or None, asset_select_mode=task.asset_select_mode or "", batch_id=task.batch_id or "", created_at=task.created_at, ) self.session.add(model) self.session.commit() return task def get(self, task_id: str) -> GenerationTask | None: model = self.session.query(GenerationTaskModel).filter(GenerationTaskModel.id == task_id).first() if model is None: return None return _to_domain(model) def list_by_project(self, project_id: str) -> list[GenerationTask]: models = ( self.session.query(GenerationTaskModel) .filter(GenerationTaskModel.project_id == project_id) .order_by(GenerationTaskModel.created_at.desc()) .all() ) return [_to_domain(m) for m in models] def list_by_user(self, user_id: str) -> list[GenerationTask]: models = ( self.session.query(GenerationTaskModel) .filter(GenerationTaskModel.created_by_user_id == user_id) .order_by(GenerationTaskModel.created_at.desc()) .all() ) return [_to_domain(m) for m in models] def count_by_user(self, user_id: str) -> int: return self.session.query(GenerationTaskModel).filter(GenerationTaskModel.created_by_user_id == user_id).count() def list_recent_by_user(self, user_id: str, limit: int = 5) -> list[GenerationTask]: models = ( self.session.query(GenerationTaskModel) .filter(GenerationTaskModel.created_by_user_id == user_id) .order_by(GenerationTaskModel.created_at.desc()) .limit(limit) .all() ) return [_to_domain(m) for m in models] def list_by_source_edit_plan(self, plan_id: str) -> list[GenerationTask]: models = ( self.session.query(GenerationTaskModel) .filter(GenerationTaskModel.source_edit_plan_id == plan_id) .order_by(GenerationTaskModel.created_at.desc()) .all() ) return [_to_domain(m) for m in models] def update(self, task: GenerationTask) -> GenerationTask: model = self.session.query(GenerationTaskModel).filter(GenerationTaskModel.id == task.id).first() if model is None: raise ValueError(f"GenerationTask {task.id} not found") model.project_id = task.project_id model.asset_library_id = task.asset_library_id model.strategy_id = task.strategy_id model.voice_library_id = task.voice_library_id model.template_id = task.template_id model.asset_ids = task.asset_ids model.title_ids = task.title_ids model.voice_ids = task.voice_ids model.status = task.status model.progress = task.progress model.result_count = task.result_count model.error_message = task.error_message model.started_at = task.started_at model.completed_at = task.completed_at model.source_edit_plan_id = task.source_edit_plan_id or None model.asset_select_mode = task.asset_select_mode or "" model.batch_id = task.batch_id or "" self.session.commit() return task