доработка

tg_bot/domain/services/user_service.py

@@ -1,29 +1,67 @@
-from sqlalchemy.orm import Session
+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: Session):
+    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 = self.session.query(UserModel) \
-                .filter(UserModel.telegram_id == str(telegram_id)) \
-                .first()
+            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)
-                self.session.commit()
+                await self.session.commit()
                 return True
             else:
                 return False
         except Exception as e:
             print(f"Error activating premium: {e}")
-            self.session.rollback()
+            await self.session.rollback()
             return False