GramCore/services/cookies/repositories.py
2023-08-05 10:41:12 +08:00

56 lines
2.0 KiB
Python

from typing import Optional, List
from sqlmodel import select
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
from gram_core.sqlmodel.session import AsyncSession
__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_by_region(self, region: RegionEnum) -> List[Cookies]:
async with AsyncSession(self.engine) as session:
statement = select(Cookies).where(Cookies.region == region)
results = await session.exec(statement)
cookies = results.all()
return cookies