"""Tenant-aware base repository.""" from __future__ import annotations from typing import Generic, Sequence, TypeVar from uuid import UUID from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from app.core.database import Base ModelT = TypeVar("ModelT", bound=Base) class TenantBaseRepository(Generic[ModelT]): model: type[ModelT] def __init__(self, session: AsyncSession) -> None: self.session = session async def get(self, tenant_id: UUID, entity_id: UUID) -> ModelT | None: stmt = select(self.model).where( self.model.tenant_id == tenant_id, # type: ignore[attr-defined] self.model.id == entity_id, # type: ignore[attr-defined] ) result = await self.session.execute(stmt) return result.scalar_one_or_none() async def add(self, entity: ModelT) -> ModelT: self.session.add(entity) await self.session.flush() return entity async def delete(self, entity: ModelT) -> None: await self.session.delete(entity) await self.session.flush() async def list_by_tenant( self, tenant_id: UUID, *, offset: int = 0, limit: int = 20 ) -> Sequence[ModelT]: stmt = ( select(self.model) .where(self.model.tenant_id == tenant_id) # type: ignore[attr-defined] .offset(offset) .limit(limit) ) result = await self.session.execute(stmt) return result.scalars().all() async def count_by_tenant(self, tenant_id: UUID) -> int: stmt = ( select(func.count()) .select_from(self.model) .where(self.model.tenant_id == tenant_id) # type: ignore[attr-defined] ) result = await self.session.execute(stmt) return int(result.scalar_one())