68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy import select
|
|
from datetime import datetime, timedelta
|
|
from typing import Optional
|
|
from tg_bot.infrastructure.database.models import UserModel
|
|
|
|
|
|
class UserService:
|
|
|
|
def __init__(self, session: AsyncSession):
|
|
self.session = session
|
|
|
|
async def get_user_by_telegram_id(self, telegram_id: int) -> Optional[UserModel]:
|
|
result = await self.session.execute(
|
|
select(UserModel).filter_by(telegram_id=str(telegram_id))
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def get_or_create_user(
|
|
self,
|
|
telegram_id: int,
|
|
username: str = "",
|
|
first_name: str = "",
|
|
last_name: str = ""
|
|
) -> UserModel:
|
|
user = await self.get_user_by_telegram_id(telegram_id)
|
|
if not user:
|
|
user = UserModel(
|
|
telegram_id=str(telegram_id),
|
|
username=username,
|
|
first_name=first_name,
|
|
last_name=last_name
|
|
)
|
|
self.session.add(user)
|
|
await self.session.commit()
|
|
else:
|
|
user.username = username
|
|
user.first_name = first_name
|
|
user.last_name = last_name
|
|
await self.session.commit()
|
|
return user
|
|
|
|
async def update_user_questions(self, telegram_id: int) -> bool:
|
|
user = await self.get_user_by_telegram_id(telegram_id)
|
|
if user:
|
|
user.questions_used += 1
|
|
await self.session.commit()
|
|
return True
|
|
return False
|
|
|
|
async def activate_premium(self, telegram_id: int) -> bool:
|
|
try:
|
|
user = await self.get_user_by_telegram_id(telegram_id)
|
|
if user:
|
|
user.is_premium = True
|
|
if user.premium_until and user.premium_until > datetime.now():
|
|
user.premium_until = user.premium_until + timedelta(days=30)
|
|
else:
|
|
user.premium_until = datetime.now() + timedelta(days=30)
|
|
await self.session.commit()
|
|
return True
|
|
else:
|
|
return False
|
|
except Exception as e:
|
|
print(f"Error activating premium: {e}")
|
|
await self.session.rollback()
|
|
return False
|