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

"""JobRepository SQLAlchemy 适配器 — Phase 8 任务 2.10."""

from __future__ import annotations

from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import JobModel
from packages.domain.job import Job, JobStatus, JobType


def _to_domain(model: JobModel) -> Job:
    """ORM 模型 → 领域实体。"""
    return Job(
        id=model.id,
        project_id=model.project_id,
        job_type=JobType(model.job_type),
        status=JobStatus(model.status),
        progress=float(model.progress or 0.0),
        current_stage=model.current_stage or "",
        payload=dict(model.payload or {}),
        result=dict(model.result or {}),
        error_message=model.error_message or "",
        retry_count=int(model.retry_count or 0),
        max_retries=int(model.max_retries or 3),
        celery_task_id=model.celery_task_id or "",
        source_id=model.source_id or "",
        created_by_user_id=model.created_by_user_id or "",
        started_at=model.started_at,
        completed_at=model.completed_at,
        created_at=model.created_at,
        updated_at=model.updated_at,
    )


def _to_model(job: Job) -> JobModel:
    """领域实体 → ORM 模型。"""
    return JobModel(
        id=job.id,
        project_id=job.project_id,
        job_type=job.job_type.value,
        status=job.status.value,
        progress=job.progress,
        current_stage=job.current_stage,
        payload=job.payload,
        result=job.result,
        error_message=job.error_message,
        retry_count=job.retry_count,
        max_retries=job.max_retries,
        celery_task_id=job.celery_task_id,
        source_id=job.source_id,
        created_by_user_id=job.created_by_user_id,
        started_at=job.started_at,
        completed_at=job.completed_at,
        created_at=job.created_at,
        updated_at=job.updated_at,
    )


class SQLAlchemyJobRepository:
    """JobRepository 的 SQLAlchemy 实现。"""

    def __init__(self, session: Session):
        self.session = session

    def create(self, job: Job) -> Job:
        model = _to_model(job)
        self.session.add(model)
        self.session.commit()
        return job

    def get(self, job_id: str) -> Job | None:
        model = self.session.query(JobModel).filter(JobModel.id == job_id).first()
        if model is None:
            return None
        return _to_domain(model)

    def update(self, job: Job) -> Job:
        model = self.session.query(JobModel).filter(JobModel.id == job.id).first()
        if model is None:
            raise ValueError(f"Job {job.id} not found")
        model.project_id = job.project_id
        model.job_type = job.job_type.value
        model.status = job.status.value
        model.progress = job.progress
        model.current_stage = job.current_stage
        model.payload = job.payload
        model.result = job.result
        model.error_message = job.error_message
        model.retry_count = job.retry_count
        model.max_retries = job.max_retries
        model.celery_task_id = job.celery_task_id
        model.source_id = job.source_id
        model.created_by_user_id = job.created_by_user_id
        model.started_at = job.started_at
        model.completed_at = job.completed_at
        model.updated_at = job.updated_at
        self.session.commit()
        return job

    def list_by_project(
        self,
        project_id: str,
        *,
        job_type: JobType | str | None = None,
        status: JobStatus | str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[Job]:
        query = self.session.query(JobModel).filter(JobModel.project_id == project_id)
        if job_type is not None:
            jt = job_type.value if isinstance(job_type, JobType) else job_type
            query = query.filter(JobModel.job_type == jt)
        if status is not None:
            st = status.value if isinstance(status, JobStatus) else status
            query = query.filter(JobModel.status == st)
        models = query.order_by(JobModel.created_at.desc()).offset(offset).limit(limit).all()
        return [_to_domain(m) for m in models]

    def list_by_user(
        self,
        user_id: str,
        *,
        job_type: JobType | str | None = None,
        status: JobStatus | str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[Job]:
        query = self.session.query(JobModel).filter(JobModel.created_by_user_id == user_id)
        if job_type is not None:
            jt = job_type.value if isinstance(job_type, JobType) else job_type
            query = query.filter(JobModel.job_type == jt)
        if status is not None:
            st = status.value if isinstance(status, JobStatus) else status
            query = query.filter(JobModel.status == st)
        models = query.order_by(JobModel.created_at.desc()).offset(offset).limit(limit).all()
        return [_to_domain(m) for m in models]

    def count_by_project(
        self,
        project_id: str,
        *,
        status: JobStatus | str | None = None,
    ) -> int:
        query = self.session.query(JobModel).filter(JobModel.project_id == project_id)
        if status is not None:
            st = status.value if isinstance(status, JobStatus) else status
            query = query.filter(JobModel.status == st)
        return query.count()

    def find_active_by_source(
        self,
        source_id: str,
        job_type: JobType | str,
    ) -> Job | None:
        jt = job_type.value if isinstance(job_type, JobType) else job_type
        model = (
            self.session.query(JobModel)
            .filter(
                JobModel.source_id == source_id,
                JobModel.job_type == jt,
                JobModel.status.in_(["pending", "running"]),
            )
            .first()
        )
        if model is None:
            return None
        return _to_domain(model)