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