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

import json

from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import GeneratedVideoModel
from packages.domain import GeneratedVideo


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

    def create(self, video: GeneratedVideo) -> GeneratedVideo:
        model = GeneratedVideoModel(
            id=video.id,
            project_id=video.project_id,
            generation_task_id=video.generation_task_id,
            name=video.name,
            file_url=video.file_url,
            file_size=video.file_size,
            duration=video.duration,
            thumbnail_url=video.thumbnail_url,
            width=video.width,
            height=video.height,
            fps=video.fps,
            status=video.status,
            review_status=video.review_status,
            generation_params=json.dumps(video.generation_params, ensure_ascii=False),
            video_fingerprint=json.dumps(video.video_fingerprint) if video.video_fingerprint else None,
            is_duplicate=video.is_duplicate,
            duplicate_of=video.duplicate_of,
            generated_at=video.generated_at,
            created_at=video.created_at,
        )
        self.session.add(model)
        self.session.commit()
        return video

    def get(self, video_id: str) -> GeneratedVideo | None:
        model = self.session.query(GeneratedVideoModel).filter(GeneratedVideoModel.id == video_id).first()
        if model is None:
            return None
        return GeneratedVideo(
            id=model.id,
            project_id=model.project_id,
            generation_task_id=model.generation_task_id,
            name=model.name,
            file_url=model.file_url,
            file_size=int(model.file_size or 0),
            duration=model.duration,
            thumbnail_url=model.thumbnail_url,
            width=int(model.width or 0),
            height=int(model.height or 0),
            fps=model.fps,
            status=getattr(model, "status", "completed"),
            review_status=getattr(model, "review_status", "pending_review"),
            generation_params=json.loads(getattr(model, "generation_params", "{}") or "{}"),
            video_fingerprint=json.loads(getattr(model, "video_fingerprint", "null") or "null"),
            is_duplicate=getattr(model, "is_duplicate", False),
            duplicate_of=getattr(model, "duplicate_of", None),
            generated_at=model.generated_at,
            created_at=model.created_at,
        )

    def update(self, video: GeneratedVideo) -> GeneratedVideo:
        model = self.session.query(GeneratedVideoModel).filter(GeneratedVideoModel.id == video.id).first()
        if model is None:
            return video
        model.status = video.status
        model.review_status = video.review_status
        model.generation_params = json.dumps(video.generation_params, ensure_ascii=False)
        model.video_fingerprint = json.dumps(video.video_fingerprint) if video.video_fingerprint else None
        model.is_duplicate = video.is_duplicate
        model.duplicate_of = video.duplicate_of
        self.session.add(model)
        self.session.commit()
        return video

    def list_by_project(self, project_id: str) -> list[GeneratedVideo]:
        models = self.session.query(GeneratedVideoModel).filter(GeneratedVideoModel.project_id == project_id).all()
        return [self._to_domain(model) for model in models]

    def list_by_generation_task(self, generation_task_id: str) -> list[GeneratedVideo]:
        models = (
            self.session.query(GeneratedVideoModel)
            .filter(GeneratedVideoModel.generation_task_id == generation_task_id)
            .all()
        )
        return [self._to_domain(model) for model in models]

    def list_by_batch(self, batch_id: str) -> list[GeneratedVideo]:
        """通过 batch_id 查找同批次生成的所有视频（跨 generation_task 关联查询）。"""
        from packages.adapters.sqlalchemy_impl.models import GenerationTaskModel

        task_ids = (
            self.session.query(GenerationTaskModel.id).filter(GenerationTaskModel.batch_id == batch_id).subquery()
        )
        models = (
            self.session.query(GeneratedVideoModel).filter(GeneratedVideoModel.generation_task_id.in_(task_ids)).all()
        )
        return [self._to_domain(model) for model in models]

    def list_paginated(
        self,
        *,
        project_id: str | None = None,
        status: str | None = None,
        review_status: str | None = None,
        page: int = 1,
        page_size: int = 20,
    ) -> tuple[list[GeneratedVideo], int]:
        """分页查询成片列表，支持按项目、状态、复核状态筛选。"""
        query = self.session.query(GeneratedVideoModel)

        if project_id:
            query = query.filter(GeneratedVideoModel.project_id == project_id)
        if status:
            query = query.filter(GeneratedVideoModel.status == status)
        if review_status:
            query = query.filter(GeneratedVideoModel.review_status == review_status)

        total = query.count()

        models = (
            query.order_by(GeneratedVideoModel.generated_at.desc())
            .offset((page - 1) * page_size)
            .limit(page_size)
            .all()
        )

        return [self._to_domain(model) for model in models], total

    def update_review_status(self, video_id: str, review_status: str) -> GeneratedVideo | None:
        """更新成片复核状态。"""
        model = self.session.query(GeneratedVideoModel).filter(GeneratedVideoModel.id == video_id).first()
        if model is None:
            return None
        model.review_status = review_status
        self.session.add(model)
        self.session.commit()
        return self._to_domain(model)

    def update_thumbnail(self, video_id: str, thumbnail_url: str) -> bool:
        """更新成片封面图URL。"""
        model = self.session.query(GeneratedVideoModel).filter(GeneratedVideoModel.id == video_id).first()
        if model is None:
            return False
        model.thumbnail_url = thumbnail_url
        self.session.add(model)
        self.session.commit()
        return True

    def get_by_ids(self, video_ids: list[str]) -> list[GeneratedVideo]:
        """批量获取成片记录。"""
        if not video_ids:
            return []
        models = self.session.query(GeneratedVideoModel).filter(GeneratedVideoModel.id.in_(video_ids)).all()
        return [self._to_domain(model) for model in models]

    @staticmethod
    def _to_domain(model: GeneratedVideoModel) -> GeneratedVideo:
        return GeneratedVideo(
            id=model.id,
            project_id=model.project_id,
            generation_task_id=model.generation_task_id,
            name=model.name,
            file_url=model.file_url,
            file_size=int(model.file_size or 0),
            duration=model.duration,
            thumbnail_url=model.thumbnail_url,
            width=int(model.width or 0),
            height=int(model.height or 0),
            fps=model.fps,
            status=getattr(model, "status", "completed"),
            review_status=getattr(model, "review_status", "pending_review"),
            generation_params=json.loads(getattr(model, "generation_params", "{}") or "{}"),
            video_fingerprint=json.loads(getattr(model, "video_fingerprint", "null") or "null"),
            is_duplicate=getattr(model, "is_duplicate", False),
            duplicate_of=getattr(model, "duplicate_of", None),
            generated_at=model.generated_at,
            created_at=model.created_at,
        )