201 lines
9.2 KiB
Python
201 lines
9.2 KiB
Python
from datetime import datetime
|
||
from typing import Optional
|
||
from uuid import UUID
|
||
from sqlalchemy.future import select
|
||
from sqlalchemy.exc import SQLAlchemyError
|
||
from decimal import Decimal
|
||
from sqlalchemy import asc, desc, update
|
||
from sqlalchemy.orm import joinedload
|
||
from instance.model import Referral, User, Subscription, Transaction
|
||
|
||
|
||
class PostgresRepository:
|
||
def __init__(self, session_generator, logger):
|
||
self.session_generator = session_generator
|
||
self.logger = logger
|
||
|
||
async def create_user(self, telegram_id: int, username: str, invited_by: Optional[int]= None):
|
||
"""
|
||
Создаёт нового пользователя в PostgreSQL.
|
||
"""
|
||
async for session in self.session_generator():
|
||
try:
|
||
new_user = User(telegram_id=telegram_id, username=username, invited_by=invited_by)
|
||
session.add(new_user)
|
||
await session.commit()
|
||
return new_user
|
||
except SQLAlchemyError as e:
|
||
self.logger.error(f"Ошибка при создании пользователя {telegram_id}: {e}")
|
||
await session.rollback()
|
||
return None
|
||
|
||
async def get_active_subscription(self, telegram_id: int):
|
||
"""
|
||
Проверяет наличие активной подписки у пользователя.
|
||
"""
|
||
async for session in self.session_generator():
|
||
try:
|
||
result = await session.execute(
|
||
select(Subscription)
|
||
.join(User, Subscription.user_id == User.id)
|
||
.where(User.telegram_id == telegram_id, Subscription.expiry_date > datetime.utcnow())
|
||
)
|
||
result= result.scalars().first()
|
||
self.logger.info(f"Пользователь с id {telegram_id}, проверен и имеет {result}")
|
||
return result
|
||
except SQLAlchemyError as e:
|
||
self.logger.error(f"Ошибка проверки активной подписки для пользователя {telegram_id}: {e}")
|
||
return None
|
||
|
||
async def get_user_by_telegram_id(self, telegram_id: int):
|
||
"""
|
||
Возвращает пользователя по Telegram ID.
|
||
"""
|
||
async for session in self.session_generator():
|
||
try:
|
||
result = await session.execute(select(User).where(User.telegram_id == telegram_id))
|
||
if result:
|
||
return result.scalars().first()
|
||
return False
|
||
except SQLAlchemyError as e:
|
||
self.logger.error(f"Ошибка при получении пользователя {telegram_id}: {e}")
|
||
return False
|
||
|
||
|
||
async def update_balance(self, user: User, amount: float):
|
||
"""
|
||
Обновляет баланс пользователя.
|
||
|
||
:param user: Объект пользователя.
|
||
:param amount: Сумма для добавления/вычитания.
|
||
:return: True, если успешно, иначе False.
|
||
"""
|
||
self.logger.info(f"Обновление баланса пользователя: id={user.telegram_id}, current_balance={user.balance}, amount={amount}")
|
||
async for session in self.session_generator():
|
||
try:
|
||
user = await session.get(User, user.telegram_id) # Загружаем пользователя в той же сессии
|
||
if not user:
|
||
self.logger.warning(f"Пользователь с ID {user.telegram_id} не найден.")
|
||
return False
|
||
# Приведение amount к Decimal
|
||
user.balance += Decimal(amount)
|
||
await session.commit()
|
||
self.logger.info(f"Баланс пользователя id={user.telegram_id} успешно обновлен: new_balance={user.balance}")
|
||
return True
|
||
except SQLAlchemyError as e:
|
||
self.logger.error(f"Ошибка при обновлении баланса пользователя id={user.telegram_id}: {e}")
|
||
await session.rollback()
|
||
return False
|
||
|
||
async def get_last_transactions(self, user_telegram_id: int, limit: int = 10):
|
||
"""
|
||
Возвращает последние транзакции пользователя.
|
||
"""
|
||
async for session in self.session_generator():
|
||
try:
|
||
result = await session.execute(
|
||
select(Transaction)
|
||
.where(Transaction.user_id == user_telegram_id)
|
||
.order_by(desc(Transaction.created_at))
|
||
.limit(limit)
|
||
)
|
||
return result.scalars().all()
|
||
except SQLAlchemyError as e:
|
||
self.logger.error(f"Ошибка получения транзакций пользователя {user_telegram_id}: {e}")
|
||
return None
|
||
|
||
async def get_last_subscription_by_user_id(self, user_telegram_id: int, limit: int = 1):
|
||
"""
|
||
Извлекает последнюю подписку пользователя на основании user_id.
|
||
|
||
:param user_id: UUID пользователя.
|
||
:return: Объект Subscription или None.
|
||
"""
|
||
async for session in self.session_generator():
|
||
try:
|
||
result = await session.execute(
|
||
select(Subscription)
|
||
.where(Subscription.user_id == user_telegram_id)
|
||
.order_by(desc(Subscription.created_at))
|
||
.limit(limit)
|
||
)
|
||
subscriptions = list(result.scalars())
|
||
result.scalars()
|
||
self.logger.info(f"Найдены такие подписки: {subscriptions}")
|
||
return subscriptions
|
||
except SQLAlchemyError as e:
|
||
self.logger.error(f"Ошибка при получении подписки для пользователя {user_telegram_id}: {e}")
|
||
return None
|
||
|
||
async def add_record(self, record):
|
||
"""
|
||
Добавляет запись в базу данных.
|
||
|
||
:param record: Объект записи.
|
||
:return: Запись или None в случае ошибки.
|
||
"""
|
||
async for session in self.session_generator():
|
||
try:
|
||
session.add(record)
|
||
await session.commit()
|
||
return record
|
||
except SQLAlchemyError as e:
|
||
self.logger.error(f"Ошибка при добавлении записи: {record}: {e}")
|
||
await session.rollback()
|
||
return None
|
||
|
||
async def add_referral(self, referrer_id: int, referral_id: int):
|
||
"""
|
||
Добавление реферальной связи между пользователями.
|
||
"""
|
||
async for session in self.session_generator():
|
||
try:
|
||
# Проверить, существует ли уже такая реферальная связь
|
||
existing_referral = await session.execute(
|
||
select(Referral)
|
||
.where(
|
||
(Referral.inviter_id == referrer_id) &
|
||
(Referral.invited_id == referral_id)
|
||
)
|
||
)
|
||
existing_referral = existing_referral.scalars().first()
|
||
|
||
if existing_referral:
|
||
raise ValueError("Referral relationship already exists")
|
||
|
||
# Проверить, что пользователи существуют
|
||
referrer = await session.execute(
|
||
select(User).where(User.telegram_id == referrer_id)
|
||
)
|
||
referrer = referrer.scalars().first()
|
||
|
||
if not referrer:
|
||
raise ValueError("Referrer not found")
|
||
|
||
referral_user = await session.execute(
|
||
select(User).where(User.telegram_id == referral_id)
|
||
)
|
||
referral_user = referral_user.scalars().first()
|
||
|
||
if not referral_user:
|
||
raise ValueError("Referral user not found")
|
||
|
||
# Проверить, что пользователь не приглашает сам себя
|
||
if referrer_id == referral_id:
|
||
raise ValueError("User cannot refer themselves")
|
||
|
||
# Создать новую реферальную связь
|
||
new_referral = Referral(
|
||
inviter_id=referrer_id,
|
||
invited_id=referral_id
|
||
)
|
||
|
||
session.add(new_referral)
|
||
await session.commit()
|
||
|
||
self.logger.info(f"Реферальная связь создана: {referrer_id} -> {referral_id}")
|
||
|
||
except Exception as e:
|
||
await session.rollback()
|
||
self.logger.error(f"Ошибка при добавлении реферальной связи: {str(e)}")
|
||
raise |