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

import json
from datetime import datetime, timezone

from sqlalchemy.orm import Session

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


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

    def find_by_library(
        self,
        library_id: str,
        skip: int = 0,
        limit: int = 100,
    ) -> list[Asset]:
        models = self.session.query(AssetModel).filter(
            AssetModel.asset_library_id == library_id
        ).offset(skip).limit(limit).all()
        return [self._to_domain(model) for model in models]

    def find_by_project(
        self,
        project_id: str,
        skip: int = 0,
        limit: int = 100,
    ) -> list[Asset]:
        models = self.session.query(AssetModel).filter(
            AssetModel.project_id == project_id
        ).offset(skip).limit(limit).all()
        return [self._to_domain(model) for model in models]

    def find_by_id(self, asset_id: str) -> Asset | None:
        model = self.session.query(AssetModel).filter(AssetModel.id == asset_id).first()
        if model is None:
            return None
        return self._to_domain(model)

    def get(self, asset_id: str) -> Asset | None:
        return self.find_by_id(asset_id)

    def create(self, asset: Asset) -> Asset:
        now = datetime.now(timezone.utc)
        model = AssetModel(
            id=asset.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=now,
        )
        self.session.add(model)
        self.session.commit()
        return asset

    def update(self, asset: Asset) -> Asset:
        model = self.session.query(AssetModel).filter(AssetModel.id == asset.id).first()
        if model is None:
            raise ValueError(f"Asset {asset.id} not found")
        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 = datetime.now(timezone.utc)
        self.session.commit()
        return asset

    def delete(self, asset_id: str) -> bool:
        model = self.session.query(AssetModel).filter(AssetModel.id == asset_id).first()
        if model:
            self.session.delete(model)
            self.session.commit()
            return True
        return False

    def count_by_project(self, project_id: str) -> int:
        return self.session.query(AssetModel).filter(
            AssetModel.project_id == project_id
        ).count()

    def count_by_project_ids(self, project_ids: list[str]) -> int:
        if not project_ids:
            return 0
        return self.session.query(AssetModel).filter(
            AssetModel.project_id.in_(project_ids)
        ).count()

    def sum_storage_by_project_ids(self, project_ids: list[str]) -> int:
        if not project_ids:
            return 0
        from sqlalchemy import func
        result = self.session.query(func.coalesce(func.sum(AssetModel.file_size), 0)).filter(
            AssetModel.project_id.in_(project_ids)
        ).scalar()
        return int(result or 0)

    def _to_domain(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,
            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=int(model.width) if model.width is not None else None,
            height=int(model.height) if model.height is not None else None,
            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,
        )