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

    def increment_asset_count(self, library_id: str, count_delta: int = 1, size_delta: int = 0) -> None:
        """原子递增素材计数（Issue #1776）。

        使用 SQL 级 UPDATE 保证并发安全，不单独 commit（由调用方统一事务提交）。
        """
        from sqlalchemy import func

        self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == library_id).update(
            {
                AssetLibraryModel.asset_count: func.coalesce(AssetLibraryModel.asset_count, 0) + count_delta,
                AssetLibraryModel.total_size: func.coalesce(AssetLibraryModel.total_size, 0) + size_delta,
            }
        )

    def decrement_asset_count(self, library_id: str, count_delta: int = 1, size_delta: int = 0) -> None:
        """原子递减素材计数（Issue #1776），下限为 0 防止负数。

        使用 SQL 级 UPDATE 保证并发安全，不单独 commit（由调用方统一事务提交）。
        使用 CASE WHEN 兼容 SQLite（测试）和 PostgreSQL（生产）。
        """
        from sqlalchemy import case, func

        self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == library_id).update(
            {
                AssetLibraryModel.asset_count: case(
                    (func.coalesce(AssetLibraryModel.asset_count, 0) - count_delta < 0, 0),
                    else_=func.coalesce(AssetLibraryModel.asset_count, 0) - count_delta,
                ),
                AssetLibraryModel.total_size: case(
                    (func.coalesce(AssetLibraryModel.total_size, 0) - size_delta < 0, 0),
                    else_=func.coalesce(AssetLibraryModel.total_size, 0) - size_delta,
                ),
            }
        )

    def recount_assets(self, library_id: str) -> int:
        """重算素材库计数（Issue #1776）。

        直接查询实际素材数量（排除已删除），更新 asset_count 和 total_size。
        返回重算后的实际计数。
        """
        from sqlalchemy import func

        from packages.adapters.sqlalchemy_impl.models import AssetModel

        # 查询实际计数（排除 deleted）
        actual_count = (
            self.session.query(func.count(AssetModel.id))
            .filter(
                AssetModel.asset_library_id == library_id,
                AssetModel.status != "deleted",
            )
            .scalar()
            or 0
        )
        # 查询实际总大小
        actual_size = (
            self.session.query(func.coalesce(func.sum(AssetModel.file_size), 0))
            .filter(
                AssetModel.asset_library_id == library_id,
                AssetModel.status != "deleted",
            )
            .scalar()
            or 0
        )
        # 更新素材库记录
        self.session.query(AssetLibraryModel).filter(AssetLibraryModel.id == library_id).update(
            {
                AssetLibraryModel.asset_count: actual_count,
                AssetLibraryModel.total_size: actual_size,
            }
        )
        return actual_count

    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,
        )