131 lines
3.7 KiB
Python
131 lines
3.7 KiB
Python
import datetime
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.sql.functions import count
|
|
|
|
from db.models import User
|
|
from db.models.transactions import BalanceTransaction, BalanceTxType
|
|
from schemas.users import NewUserDTO
|
|
|
|
|
|
class UserRepository:
|
|
async def get_users(self, session: AsyncSession) -> list[User]:
|
|
stmt = select(User)
|
|
res = await session.execute(stmt)
|
|
|
|
return list(res.scalars().all())
|
|
|
|
async def get_user_by_id(self, session: AsyncSession, user_id: int) -> User | None:
|
|
stmt = select(User).where(User.id == user_id)
|
|
result = await session.execute(stmt)
|
|
|
|
return result.scalar_one_or_none()
|
|
|
|
async def create_user(self, session: AsyncSession, user_data: NewUserDTO) -> User:
|
|
user = User(
|
|
id=user_data.id,
|
|
referal_id=user_data.referal,
|
|
subscription_link=user_data.link,
|
|
)
|
|
session.add(user)
|
|
await session.commit()
|
|
|
|
return user
|
|
|
|
async def update_user_topic(
|
|
self, session: AsyncSession, user_id: int, thread_id: int
|
|
) -> User | None:
|
|
user = await self.get_user_by_id(session, user_id=user_id)
|
|
if not user:
|
|
return None
|
|
user.support_thread_id = thread_id
|
|
|
|
await session.commit()
|
|
return user
|
|
|
|
async def update_user_link(self, session: AsyncSession, user_id: int, link: str) -> User | None:
|
|
user = await self.get_user_by_id(session, user_id=user_id)
|
|
if not user:
|
|
return None
|
|
user.subscription_link = link
|
|
|
|
await session.commit()
|
|
return user
|
|
|
|
async def update_user_balance(
|
|
self, session: AsyncSession, user_id: int, balance: int
|
|
) -> User | None:
|
|
user = await self.get_user_by_id(session, user_id=user_id)
|
|
if not user:
|
|
return None
|
|
user.balance = balance
|
|
|
|
await session.commit()
|
|
return user
|
|
|
|
async def increase_user_balance(
|
|
self,
|
|
session: AsyncSession,
|
|
user_id: int,
|
|
amount: int,
|
|
tx_type: BalanceTxType,
|
|
description: str,
|
|
) -> User | None:
|
|
user = await self.get_user_by_id(session, user_id=user_id)
|
|
if not user:
|
|
return None
|
|
|
|
created_at = datetime.datetime.now(datetime.UTC)
|
|
transaction = BalanceTransaction(
|
|
user_id=user_id,
|
|
amount=amount,
|
|
tx_type=tx_type,
|
|
balance_before=user.balance,
|
|
balance_after=user.balance + amount,
|
|
description=description,
|
|
created_at=created_at,
|
|
)
|
|
|
|
user.balance += amount
|
|
|
|
session.add(transaction)
|
|
|
|
await session.commit()
|
|
return user
|
|
|
|
async def decrease_user_balance(
|
|
self,
|
|
session: AsyncSession,
|
|
user_id: int,
|
|
amount: int,
|
|
tx_type: BalanceTxType,
|
|
description: str,
|
|
) -> User | None:
|
|
user = await self.get_user_by_id(session, user_id=user_id)
|
|
if not user or user.balance < amount:
|
|
return None
|
|
|
|
created_at = datetime.datetime.now(datetime.UTC)
|
|
transaction = BalanceTransaction(
|
|
user_id=user_id,
|
|
amount=amount,
|
|
tx_type=tx_type,
|
|
balance_before=user.balance,
|
|
balance_after=user.balance - amount,
|
|
description=description,
|
|
created_at=created_at,
|
|
)
|
|
|
|
user.balance -= amount
|
|
session.add(transaction)
|
|
|
|
await session.commit()
|
|
return user
|
|
|
|
async def user_count(self, session: AsyncSession) -> int:
|
|
stmt = select(count()).select_from(User)
|
|
res = await session.execute(stmt)
|
|
|
|
return int(res.scalar_one_or_none())
|