xiaoxia-saas/packages/adapters/sqlalchemy_impl/tts_job_repository.py

"""SQLAlchemy implementation of TTSJobRepository."""

from __future__ import annotations

from typing import List, Optional

from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import TTSJobModel
from packages.domain.tts_job import TTSJob, TTSJobStatus


class SQLAlchemyTTSJobRepository:
    """SQLAlchemy TTS 任务仓储。"""

    def __init__(self, session: Session) -> None:
        self.session = session

    def create(self, job: TTSJob) -> TTSJob:
        model = TTSJobModel(
            id=job.id,
            user_id=job.user_id,
            input_text=job.input_text,
            voice_id=job.voice_id,
            voice_model=job.voice_model,
            project_id=job.project_id,
            voice_clone_profile_id=job.voice_clone_profile_id,
            status=job.status,
            output_audio_url=job.output_audio_url,
            output_audio_key=job.output_audio_key,
            duration=job.duration,
            file_size=job.file_size,
            sample_rate=job.sample_rate,
            format=job.format,
            error_message=job.error_message,
            retry_count=job.retry_count,
            max_retries=job.max_retries,
            metadata_=job.metadata,
            started_at=job.started_at,
            completed_at=job.completed_at,
        )
        self.session.add(model)
        self.session.commit()
        self.session.refresh(model)
        return self._model_to_entity(model)

    def get(self, job_id: str) -> Optional[TTSJob]:
        model = (
            self.session.query(TTSJobModel)
            .filter(
                TTSJobModel.id == job_id,
                TTSJobModel.status != "deleted",
            )
            .first()
        )
        if model is None:
            return None
        return self._model_to_entity(model)

    def update(self, job: TTSJob) -> TTSJob:
        model = self.session.query(TTSJobModel).filter(TTSJobModel.id == job.id).first()
        if model is None:
            raise ValueError(f"TTSJob {job.id} not found")
        model.input_text = job.input_text
        model.voice_id = job.voice_id
        model.voice_model = job.voice_model
        model.project_id = job.project_id
        model.voice_clone_profile_id = job.voice_clone_profile_id
        model.status = job.status
        model.output_audio_url = job.output_audio_url
        model.output_audio_key = job.output_audio_key
        model.duration = job.duration
        model.file_size = job.file_size
        model.sample_rate = job.sample_rate
        model.format = job.format
        model.error_message = job.error_message
        model.retry_count = job.retry_count
        model.max_retries = job.max_retries
        model.metadata_ = job.metadata
        model.started_at = job.started_at
        model.completed_at = job.completed_at
        self.session.commit()
        self.session.refresh(model)
        return self._model_to_entity(model)

    def delete(self, job_id: str) -> bool:
        model = self.session.query(TTSJobModel).filter(TTSJobModel.id == job_id).first()
        if model is None:
            return False
        model.status = "deleted"
        self.session.commit()
        return True

    def list_by_user(
        self,
        user_id: str,
        *,
        status: Optional[str] = None,
        limit: int = 50,
        offset: int = 0,
    ) -> List[TTSJob]:
        query = self.session.query(TTSJobModel).filter(
            TTSJobModel.user_id == user_id,
            TTSJobModel.status != "deleted",
        )
        if status:
            query = query.filter(TTSJobModel.status == status)
        query = query.order_by(TTSJobModel.created_at.desc())
        models = query.offset(offset).limit(limit).all()
        return [self._model_to_entity(m) for m in models]

    def count_by_user(self, user_id: str, *, status: Optional[str] = None) -> int:
        query = self.session.query(TTSJobModel).filter(
            TTSJobModel.user_id == user_id,
            TTSJobModel.status != "deleted",
        )
        if status:
            query = query.filter(TTSJobModel.status == status)
        return query.count()

    def list_by_profile(
        self,
        voice_clone_profile_id: str,
        *,
        status: Optional[str] = None,
        limit: int = 50,
        offset: int = 0,
    ) -> List[TTSJob]:
        query = self.session.query(TTSJobModel).filter(
            TTSJobModel.voice_clone_profile_id == voice_clone_profile_id,
            TTSJobModel.status != "deleted",
        )
        if status:
            query = query.filter(TTSJobModel.status == status)
        query = query.order_by(TTSJobModel.created_at.desc())
        models = query.offset(offset).limit(limit).all()
        return [self._model_to_entity(m) for m in models]

    @staticmethod
    def _model_to_entity(model: TTSJobModel) -> TTSJob:
        return TTSJob(
            id=model.id,
            user_id=model.user_id,
            input_text=model.input_text or "",
            voice_id=model.voice_id or "",
            voice_model=model.voice_model or "",
            project_id=model.project_id or "",
            voice_clone_profile_id=model.voice_clone_profile_id or "",
            status=TTSJobStatus(model.status) if model.status else TTSJobStatus.PENDING,
            output_audio_url=model.output_audio_url or "",
            output_audio_key=model.output_audio_key or "",
            duration=model.duration or 0.0,
            file_size=model.file_size or 0,
            sample_rate=model.sample_rate or 22050,
            format=model.format or "mp3",
            error_message=model.error_message or "",
            retry_count=model.retry_count or 0,
            max_retries=model.max_retries or 3,
            metadata=model.metadata_ or {},
            started_at=model.started_at,
            completed_at=model.completed_at,
            created_at=model.created_at,
            updated_at=model.updated_at,
        )