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

from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import AssetLibraryModel
from packages.domain import AssetLibrary, AssetLibraryKind


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

    def get(self, library_id: str) -> AssetLibrary | None:
        model = self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == library_id).first()
        if model is None:
            return None
        return AssetLibrary(
            id=model.id,
            project_id=model.project_id,
            name=model.name,
            kind=AssetLibraryKind(model.kind),
            asset_count=int(model.asset_count or 0),
            total_size=int(model.total_size or 0),
            created_at=model.created_at,
            updated_at=model.updated_at,
        )

    def find_by_id(self, library_id: str) -> AssetLibrary | None:
        return self.get(library_id)

    def find_by_project(self, project_id: str) -> list[AssetLibrary]:
        models = self.session.query(AssetLibraryModel).filter(AssetLibraryModel.project_id == project_id).all()
        return [
            AssetLibrary(
                id=model.id,
                project_id=model.project_id,
                name=model.name,
                kind=AssetLibraryKind(model.kind),
                asset_count=int(model.asset_count or 0),
                total_size=int(model.total_size or 0),
                created_at=model.created_at,
                updated_at=model.updated_at,
            )
            for model in models
        ]

    def create(self, library: AssetLibrary) -> AssetLibrary:
        model = AssetLibraryModel(
            id=library.id,
            project_id=library.project_id,
            name=library.name,
            kind=library.kind.value,
            asset_count=library.asset_count,
            total_size=library.total_size,
            created_at=library.created_at,
            updated_at=library.updated_at,
        )
        self.session.add(model)
        self.session.commit()
        return library

    def update(self, library: AssetLibrary) -> AssetLibrary:
        model = self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == library.id).first()
        if model:
            model.project_id = library.project_id
            model.name = library.name
            model.kind = library.kind.value
            model.asset_count = library.asset_count
            model.total_size = library.total_size
            model.updated_at = library.updated_at
            self.session.commit()
        return library

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

    async def increment_asset_count(self, library_id: str, size_delta: int) -> None:
        model = self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == library_id).first()
        if model:
            model.asset_count = (model.asset_count or 0) + 1
            model.total_size = (model.total_size or 0) + size_delta
            self.session.commit()

    async def decrement_asset_count(self, library_id: str, size_delta: int) -> None:
        model = self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == library_id).first()
        if model:
            model.asset_count = max(0, (model.asset_count or 0) - 1)
            model.total_size = max(0, (model.total_size or 0) - size_delta)
            self.session.commit()

    def get_or_create_default_library(
        self,
        project_id: str,
        kind: AssetLibraryKind,
        *,
        name: str | None = None,
    ) -> AssetLibrary:
        """幂等获取/创建项目下指定 kind 的默认素材库（Issue #1775）。

        依赖唯一约束 uq_asset_libraries_project_kind(project_id, kind)：
        并发创建只有一个成功，其余 IntegrityError 后回滚重查，
        保证同一项目同 kind 永远只有一个素材库。
        """
        from sqlalchemy.exc import IntegrityError

        default_names = {
            AssetLibraryKind.VIDEO: "视频素材库",
            AssetLibraryKind.VOICE: "配音素材库",
            AssetLibraryKind.IMAGE: "图片素材库",
        }
        library_name = name or default_names.get(kind, f"{kind.value}素材库")

        # 快速路径
        existing = (
            self.session.query(AssetLibraryModel)
            .filter(AssetLibraryModel.project_id == project_id, AssetLibraryModel.kind == kind.value)
            .first()
        )
        if existing:
            return self._to_entity(existing)

        library = AssetLibrary.create(project_id=project_id, name=library_name, kind=kind)
        model = AssetLibraryModel(
            id=library.id,
            project_id=library.project_id,
            name=library.name,
            kind=library.kind.value,
            asset_count=0,
            total_size=0,
            created_at=library.created_at,
            updated_at=library.updated_at,
        )
        try:
            self.session.add(model)
            self.session.commit()
            return library
        except IntegrityError:
            self.session.rollback()
            existing = (
                self.session.query(AssetLibraryModel)
                .filter(
                    AssetLibraryModel.project_id == project_id,
                    AssetLibraryModel.kind == kind.value,
                )
                .first()
            )
            if existing:
                return self._to_entity(existing)
            raise

    def _to_entity(self, model: AssetLibraryModel) -> AssetLibrary:
        return AssetLibrary(
            id=model.id,
            project_id=model.project_id,
            name=model.name,
            kind=AssetLibraryKind(model.kind),
            asset_count=int(model.asset_count or 0),
            total_size=int(model.total_size or 0),
            created_at=model.created_at,
            updated_at=model.updated_at,
        )