GramCore/services/cookies/repositories.py

from typing import Optional, List, Tuple

from sqlmodel import select
from sqlmodel.ext.asyncio.session import AsyncSession

from gram_core.base_service import BaseService
from gram_core.basemodel import RegionEnum
from gram_core.dependence.database import Database
from gram_core.services.cookies.models import CookiesDataBase as Cookies, CookiesStatusEnum
from gram_core.services.devices.models import DevicesDataBase as Devices

__all__ = ("CookiesRepository",)


class CookiesRepository(BaseService.Component):
    def __init__(self, database: Database):
        self.engine = database.engine

    async def get(
        self,
        user_id: int,
        account_id: Optional[int] = None,
        region: Optional[RegionEnum] = None,
    ) -> Optional[Cookies]:
        async with AsyncSession(self.engine) as session:
            statement = select(Cookies).where(Cookies.user_id == user_id)
            if account_id is not None:
                statement = statement.where(Cookies.account_id == account_id)
            if region is not None:
                statement = statement.where(Cookies.region == region)
            results = await session.exec(statement)
            return results.first()

    async def add(self, cookies: Cookies) -> None:
        async with AsyncSession(self.engine) as session:
            session.add(cookies)
            await session.commit()

    async def update(self, cookies: Cookies) -> Cookies:
        async with AsyncSession(self.engine) as session:
            session.add(cookies)
            await session.commit()
            await session.refresh(cookies)
            return cookies

    async def delete(self, cookies: Cookies) -> None:
        async with AsyncSession(self.engine) as session:
            await session.delete(cookies)
            await session.commit()

    async def get_all(
        self,
        user_id: Optional[int] = None,
        account_id: Optional[int] = None,
        region: Optional[RegionEnum] = None,
        status: Optional[CookiesStatusEnum] = None,
    ) -> List[Cookies]:
        async with AsyncSession(self.engine) as session:
            statement = select(Cookies)
            if user_id is not None:
                statement = statement.where(Cookies.user_id == user_id)
            if account_id is not None:
                statement = statement.where(Cookies.account_id == account_id)
            if region is not None:
                statement = statement.where(Cookies.region == region)
            if status is not None:
                statement = statement.where(Cookies.status == status)
            results = await session.exec(statement)
            return results.all()

    async def get_by_devices(self, is_valid: bool = None) -> List[Tuple[Cookies, Devices]]:
        async with AsyncSession(self.engine) as session:
            statement = select(Cookies, Devices).where(Cookies.account_id == Devices.account_id)
            if is_valid is not None:
                statement = statement.where(Devices.is_valid == is_valid)
            results = await session.exec(statement)
            return results.all()