diff --git a/bot.py b/bot.py index ac37e8bf..f37d24e2 100644 --- a/bot.py +++ b/bot.py @@ -2,6 +2,7 @@ from aiogram import Bot, Dispatcher, Router from aiogram.fsm.storage.memory import MemoryStorage from config import API_TOKEN, CRYPTO_BOT_ENABLE, FREEKASSA_ENABLE, ROBOKASSA_ENABLE, STARS_ENABLE, YOOKASSA_ENABLE +from middlewares.user import UserMiddleware from middlewares.admin import AdminMiddleware from middlewares.database import DatabaseMiddleware from middlewares.logging import LoggingMiddleware @@ -51,3 +52,6 @@ dp.callback_query.middleware(DatabaseMiddleware()) dp.message.middleware(AdminMiddleware()) dp.callback_query.middleware(AdminMiddleware()) + +dp.message.middleware(UserMiddleware()) +dp.callback_query.middleware(UserMiddleware()) diff --git a/database.py b/database.py index 186c77ea..2f6f5952 100644 --- a/database.py +++ b/database.py @@ -9,6 +9,22 @@ from logger import logger async def init_db(): conn = await asyncpg.connect(DATABASE_URL) + # Таблица для хранения основной информации о пользователях из Telegram + await conn.execute( + """ + CREATE TABLE IF NOT EXISTS users ( + tg_id BIGINT PRIMARY KEY NOT NULL, + username TEXT, + first_name TEXT, + last_name TEXT, + language_code TEXT, + is_bot BOOLEAN DEFAULT FALSE, + created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP + ) + """ + ) + # Таблица для хранения информации о пользователях await conn.execute( """ @@ -420,3 +436,46 @@ async def get_tg_id_by_client_id(client_id: str): return result["tg_id"] if result else None finally: await conn.close() + + +async def upsert_user(tg_id: int, username: str = None, first_name: str = None, last_name: str = None, language_code: str = None, is_bot: bool = False): + """ + Создает или обновляет информацию о пользователе в базе данных. + + Args: + tg_id (int): Уникальный идентификатор пользователя в Telegram + username (str, optional): Никнейм пользователя + first_name (str, optional): Имя пользователя + last_name (str, optional): Фамилия пользователя + language_code (str, optional): Код языка пользователя + is_bot (bool, optional): Флаг, указывающий является ли пользователь ботом + """ + conn = await asyncpg.connect(DATABASE_URL) + try: + await conn.execute( + """ + INSERT INTO users (tg_id, username, first_name, last_name, language_code, is_bot, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + ON CONFLICT (tg_id) DO UPDATE + SET + username = COALESCE(EXCLUDED.username, users.username), + first_name = COALESCE(EXCLUDED.first_name, users.first_name), + last_name = COALESCE(EXCLUDED.last_name, users.last_name), + language_code = COALESCE(EXCLUDED.language_code, users.language_code), + is_bot = EXCLUDED.is_bot, + updated_at = CURRENT_TIMESTAMP + """, + tg_id, username, first_name, last_name, language_code, is_bot + ) + + # Создаем запись в connections, если ее еще нет + await conn.execute( + """ + INSERT INTO connections (tg_id, balance, trial) + VALUES ($1, 0.0, 0) + ON CONFLICT (tg_id) DO NOTHING + """, + tg_id + ) + finally: + await conn.close() diff --git a/middlewares/user.py b/middlewares/user.py new file mode 100644 index 00000000..7fcbfc14 --- /dev/null +++ b/middlewares/user.py @@ -0,0 +1,28 @@ +from typing import Any, Awaitable, Callable, Dict + +from aiogram import BaseMiddleware +from aiogram.types import TelegramObject, User + +from database import upsert_user + + +class UserMiddleware(BaseMiddleware): + async def __call__( + self, + handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]], + event: TelegramObject, + data: Dict[str, Any], + ) -> Any: + user: User = data.get("event_from_user") + + if user: + await upsert_user( + tg_id=user.id, + username=user.username, + first_name=user.first_name, + last_name=user.last_name, + language_code=user.language_code, + is_bot=user.is_bot + ) + + return await handler(event, data)