xiaoxia-saas/packages/adapters/postgres/asset_repository.py

"""
Asset PostgreSQL Repository 实现
"""
from typing import List, Optional
from sqlalchemy import select, and_, func
from sqlalchemy.ext.asyncio import AsyncSession
from packages.ports.asset_repository import AssetRepository
from packages.domain.asset import Asset, AssetType, AssetStatus, ClassificationStatus
from packages.adapters.sqlalchemy_impl.models import AssetModel


class PostgresAssetRepository(AssetRepository):
    """素材 PostgreSQL 仓储实现"""

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

    async def create(self, asset: Asset) -> Asset:
        """创建素材"""
        model = AssetModel(
            id=asset.id,
            workspace_id=asset.workspace_id,
            project_id=asset.project_id,
            asset_library_id=asset.asset_library_id,
            name=asset.name,
            file_type=asset.file_type.value,
            file_size=asset.file_size,
            file_url=asset.file_url,
            thumbnail_url=asset.thumbnail_url,
            duration=asset.duration,
            width=asset.width,
            height=asset.height,
            fps=asset.fps,
            codec=asset.codec,
            status=asset.status.value,
            classification_status=asset.classification_status.value,
            classification_result=asset.classification_result,
            quality_score=asset.quality_score,
            uploaded_by_user_id=asset.uploaded_by_user_id,
            created_at=asset.created_at,
            updated_at=asset.updated_at,
        )
        self.session.add(model)
        await self.session.flush()
        return asset

    async def find_by_id(self, asset_id: str) -> Optional[Asset]:
        """根据 ID 查询素材"""
        result = await self.session.execute(
            select(AssetModel).where(AssetModel.id == asset_id)
        )
        model = result.scalar_one_or_none()
        return self._to_entity(model) if model else None

    async def find_by_project(
        self,
        project_id: str,
        workspace_id: str,
        skip: int = 0,
        limit: int = 100,
    ) -> List[Asset]:
        """根据项目查询素材列表"""
        result = await self.session.execute(
            select(AssetModel)
            .where(
                and_(
                    AssetModel.project_id == project_id,
                    AssetModel.workspace_id == workspace_id,
                )
            )
            .order_by(AssetModel.created_at.desc())
            .offset(skip)
            .limit(limit)
        )
        return [self._to_entity(model) for model in result.scalars().all()]

    async def find_by_library(
        self,
        library_id: str,
        workspace_id: str,
        skip: int = 0,
        limit: int = 100,
    ) -> List[Asset]:
        """根据素材库查询素材列表"""
        result = await self.session.execute(
            select(AssetModel)
            .where(
                and_(
                    AssetModel.asset_library_id == library_id,
                    AssetModel.workspace_id == workspace_id,
                )
            )
            .order_by(AssetModel.created_at.desc())
            .offset(skip)
            .limit(limit)
        )
        return [self._to_entity(model) for model in result.scalars().all()]

    async def update(self, asset: Asset) -> Asset:
        """更新素材"""
        result = await self.session.execute(
            select(AssetModel).where(AssetModel.id == asset.id)
        )
        model = result.scalar_one_or_none()
        if model:
            model.name = asset.name
            model.status = asset.status.value
            model.thumbnail_url = asset.thumbnail_url
            model.duration = asset.duration
            model.width = asset.width
            model.height = asset.height
            model.fps = asset.fps
            model.codec = asset.codec
            model.classification_status = asset.classification_status.value
            model.classification_result = asset.classification_result
            model.quality_score = asset.quality_score
            model.updated_at = asset.updated_at
            await self.session.flush()
        return asset

    async def delete(self, asset_id: str, workspace_id: str) -> bool:
        """删除素材"""
        result = await self.session.execute(
            select(AssetModel).where(
                and_(
                    AssetModel.id == asset_id,
                    AssetModel.workspace_id == workspace_id,
                )
            )
        )
        model = result.scalar_one_or_none()
        if model:
            await self.session.delete(model)
            await self.session.flush()
            return True
        return False

    async def count_by_project(self, project_id: str, workspace_id: str) -> int:
        """统计项目素材数量"""
        result = await self.session.execute(
            select(func.count(AssetModel.id)).where(
                and_(
                    AssetModel.project_id == project_id,
                    AssetModel.workspace_id == workspace_id,
                )
            )
        )
        return result.scalar() or 0

    def _to_entity(self, model: AssetModel) -> Asset:
        """模型转实体"""
        return Asset(
            id=model.id,
            workspace_id=model.workspace_id,
            project_id=model.project_id,
            asset_library_id=model.asset_library_id,
            name=model.name,
            file_type=AssetType(model.file_type),
            file_size=model.file_size,
            file_url=model.file_url,
            thumbnail_url=model.thumbnail_url,
            duration=model.duration,
            width=model.width,
            height=model.height,
            fps=model.fps,
            codec=model.codec,
            status=AssetStatus(model.status),
            classification_status=ClassificationStatus(model.classification_status),
            classification_result=model.classification_result,
            quality_score=model.quality_score,
            uploaded_by_user_id=model.uploaded_by_user_id,
            created_at=model.created_at,
            updated_at=model.updated_at,
        )