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

"""
Asset PostgreSQL Repository 实现
"""

import json

from sqlalchemy import and_, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from packages.adapters.sqlalchemy_impl.models import AssetModel
from packages.domain import Asset, AssetStatus, ClassificationStatus
from packages.ports.asset_repository import AssetRepository


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.library_id,
            name=asset.name,
            file_type=(asset.mime_type.split("/")[0] if "/" in asset.mime_type else asset.mime_type),
            file_size=asset.file_size,
            file_url=asset.storage_key,
            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=(json.dumps(asset.metadata) if asset.metadata else None),
            quality_score=asset.quality_score,
            uploaded_by_user_id=asset.uploaded_by_user_id or "system",
            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) -> Asset | None:
        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.file_size = asset.file_size
            model.file_url = asset.storage_key
            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.status = asset.status.value
            model.classification_status = asset.classification_status.value
            model.classification_result = json.dumps(asset.metadata) if asset.metadata else None
            model.quality_score = asset.quality_score
            model.uploaded_by_user_id = asset.uploaded_by_user_id or model.uploaded_by_user_id
            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:
        metadata = {}
        if model.classification_result:
            try:
                metadata = json.loads(model.classification_result)
            except Exception:
                metadata = {}
        mime_type = model.file_type
        if "/" not in mime_type:
            mime_type = {
                "video": "video/mp4",
                "audio": "audio/mpeg",
                "image": "image/jpeg",
            }.get(mime_type, mime_type)
        return Asset(
            id=model.id,
            workspace_id=model.workspace_id,
            project_id=model.project_id,
            library_id=model.asset_library_id,
            name=model.name,
            storage_key=model.file_url,
            mime_type=mime_type,
            file_size=int(model.file_size or 0),
            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),
            quality_score=model.quality_score,
            uploaded_by_user_id=model.uploaded_by_user_id,
            metadata=metadata,
            created_at=model.created_at,
            updated_at=model.updated_at,
        )