Little fixes, add exception handler

2026-01-24 21:55:27 +03:00
parent 7758a3cf62
commit 239dd05992
14 changed files with 244 additions and 52 deletions
--- a/birthday_pool_bot/telegram_bot/ui/callback_data.py
+++ b/birthday_pool_bot/telegram_bot/ui/callback_data.py
@@ -57,3 +57,14 @@ class PoolActionEnum(str, enum.Enum):
 class PoolCallbackData(CallbackData, prefix="pool"):
    id: uuid.UUID
    action: PoolActionEnum = PoolActionEnum.SHOW
+
+
+class SubscribeToNewPoolAnswerEnum(str, enum.Enum):
+    YES = "yes"
+    NO = "no"
+
+
+class NotifyNewPoolCallbackData(CallbackData, prefix="snp"):
+    user_id: uuid.UUID
+    pool_id: uuid.UUID
+    answer: SubscribeToNewPoolAnswerEnum
--- a/birthday_pool_bot/telegram_bot/ui/dispatcher.py
+++ b/birthday_pool_bot/telegram_bot/ui/dispatcher.py
@@ -8,10 +8,12 @@ from .callback_data import (
    MenuCallbackData,
    NewSubscriptionCallbackData,
    NewSubscriptionConfirmCallbackData,
+    NotifyNewPoolCallbackData,
    PoolActionEnum,
    PoolCallbackData,
    PoolsBackCallbackData,
    PoolsCallbackData,
+    SubscribeToNewPoolAnswerEnum,
    SubscriptionActionEnum,
    SubscriptionCallbackData,
    SubscriptionsCallbackData,
@@ -19,6 +21,7 @@ from .callback_data import (
 from .middlewares import (
    AuthMiddleware,
    DependsMiddleware,
+    ExceptionHandlerMiddleware,
    # TypingMiddleware,
 )
 from .states import (
@@ -36,6 +39,7 @@ def setup_dispatcher(
        dispatcher: aiogram.Dispatcher,
        repositories_container: RepositoriesContainer,
 ):
+    exception_handler_middleware = ExceptionHandlerMiddleware()
    # typing_middleware = TypingMiddleware()
    auth_middleware = AuthMiddleware(users_repository=repositories_container.users)
    users_repository_middleware = DependsMiddleware(
@@ -51,6 +55,7 @@ def setup_dispatcher(
        object_=repositories_container.subscriptions,
    )
    middlewares = [
+        exception_handler_middleware,
        # typing_middleware,
        auth_middleware,
        users_repository_middleware,
@@ -333,6 +338,16 @@ def setup_router(router: aiogram.Router) -> aiogram.Router:
        PoolsCallbackData.filter(),
    )

+    # Notification about new pool
+    router.callback_query.register(
+        handlers.subscribe_to_new_pool_callback_handler,
+        NotifyNewPoolCallbackData.filter(aiogram.F.answer == SubscribeToNewPoolAnswerEnum.YES),
+    )
+    router.callback_query.register(
+        handlers.hide_new_pool_notification,
+        NotifyNewPoolCallbackData.filter(aiogram.F.answer == SubscribeToNewPoolAnswerEnum.NO),
+    )
+
    # Fallback
    router.message.register(handlers.fallback_message_handler)

--- a/birthday_pool_bot/telegram_bot/ui/handlers.py
+++ b/birthday_pool_bot/telegram_bot/ui/handlers.py
@@ -1,4 +1,6 @@
-from aiogram import types
+import asyncio
+
+from aiogram import Bot, types
 from aiogram.enums import MessageOriginType, ParseMode
 from aiogram.fsm.context import FSMContext

@@ -14,7 +16,11 @@ from birthday_pool_bot.repositories.repositories import (
    UsersRepository,
 )
 from . import constants, logic
-from .callback_data import PoolCallbackData, SubscriptionCallbackData
+from .callback_data import (
+    NotifyNewPoolCallbackData,
+    PoolCallbackData,
+    SubscriptionCallbackData,
+)
 from .exceptions import FlowInternalError
 from .states import (
    NewSubscriptionPoolState,
@@ -337,10 +343,11 @@ async def set_new_subscription_user_message_handler(

    async with users_repository.transaction():
        if message.contact is not None:
-            subscription_users = await users_repository.get_users_by_primary_keys(
+            subscription_users_generator = users_repository.get_users_by_primary_keys(
                phone=message.contact.phone_number,
                telegram_id=message.contact.user_id,
            )
+            subscription_users = [user async for user in subscription_users_generator]
            subscription_user = await users_repository.merge_users(*subscription_users)
            user_name = get_telegram_user_full_name(message.contact)
            user_phone = message.contact.phone_number
@@ -370,7 +377,7 @@ async def set_new_subscription_user_message_handler(
            if user_phone is None:
                await logic.ask_new_subscription_user(message=message)
                await state.set_state(NewSubscriptionState.WAITING_FOR_PHONE)
-                return                
+                return
            subscription_user = await users_repository.get_user_by_phone(phone=user_phone)

    if subscription_user is None:
@@ -401,7 +408,7 @@ async def set_new_subscription_user_message_handler(
        return

    await state.update_data(
-        subscription_user_id=subscription_user.id,
+        subscription_user_id=str(subscription_user.id),
    )

    if subscription_user.birthday is None:
@@ -577,7 +584,7 @@ async def ask_new_subscription_pool_decision_callback_handler(
            subscription_user_id=subscription_user_id,
        ) 
    except FlowInternalError:
-        await message.reply(text=constants.ERROR_MESSAGE)
+        await callback_query.answer(text=constants.ERROR_MESSAGE)
        await logic.ask_new_subscription_user(message=callback_query.message)
        await state.set_state(NewSubscriptionState.WAITING_FOR_PHONE)
        return
@@ -652,6 +659,8 @@ async def show_new_subscription_choosing_pools_callback_handler(
 async def show_new_subscription_choosing_pool_callback_handler(
        callback_query: types.CallbackQuery,
        state: FSMContext,
+        user: User,
+        users_repository: UsersRepository,
        pools_repository: PoolsRepository,
 ):
    callback_data = PoolCallbackData.unpack(callback_query.data)
@@ -668,7 +677,7 @@ async def show_new_subscription_choosing_pool_callback_handler(
            user=user,
            users_repository=users_repository,
            pools_repository=pools_repository,
-            subscription_user_id=subscription_user_id,
+            subscription_user_id=callback_data.user_id,
        )
        await state.set_state(NewSubscriptionState.WAITING_FOR_CHOOSE_POOL)
        return
@@ -679,6 +688,7 @@ async def show_new_subscription_choosing_pool_callback_handler(
 async def choose_new_subscription_pool_callback_handler(
        callback_query: types.CallbackQuery,
        state: FSMContext,
+        user: User,
        users_repository: UsersRepository,
        pools_repository: PoolsRepository,
 ):
@@ -686,8 +696,8 @@ async def choose_new_subscription_pool_callback_handler(
    subscription_user_id = state_data.get("subscription_user_id")
    subscription_name = state_data.get("subscription_name")
    if not all((subscription_user_id, subscription_name)):
-        await message.reply(text=constants.ERROR_MESSAGE)
-        await logic.ask_new_subscription_user(message=message)
+        await callback_query.answer(text=constants.ERROR_MESSAGE)
+        await logic.ask_new_subscription_user(message=callback_query.message)
        await state.set_state(NewSubscriptionState.WAITING_FOR_PHONE)
        return

@@ -696,8 +706,8 @@ async def choose_new_subscription_pool_callback_handler(
            user_id=subscription_user_id,
        )
    if subscription_user is None:
-        await message.reply(text=constants.ERROR_MESSAGE)
-        await logic.ask_new_subscription_user(message=message)
+        await callback_query.answer(text=constants.ERROR_MESSAGE)
+        await logic.ask_new_subscription_user(message=callback_query.message)
        await state.set_state(NewSubscriptionState.WAITING_FOR_PHONE)
        return

@@ -720,7 +730,7 @@ async def choose_new_subscription_pool_callback_handler(
        return

    await state.update_data(
-        subscription_pool_id=callback_data.id,
+        subscription_pool_id=str(callback_data.id),
    )

    try:
@@ -960,6 +970,7 @@ async def ask_new_subscription_confirmation_message_handler(
 async def confirm_new_subscription_callback_handler(
        callback_query: types.CallbackQuery,
        state: FSMContext,
+        bot: Bot,
        user: User,
        pools_repository: PoolsRepository,
        subscriptions_repository: SubscriptionsRepository,
@@ -1032,3 +1043,65 @@ async def confirm_new_subscription_callback_handler(
    )
    await state.clear()
    await state.set_state(MenuState.SUBSCRIPTIONS)
+
+    subscriptions_generator = subscriptions_repository.get_to_users_subscriptions(
+        to_users_ids={subscription_user_id},
+        have_pool=False,
+        with_from_user=True,
+        with_pool=True,
+        with_pool_owner=True,
+    )
+    async with subscriptions_repository.transaction():
+        coroutines = [
+            logic.notify_user_about_new_subscription_pool(
+                subscription=subscription,
+                user=subscription.from_user,
+                pool=subscription.pool,
+                pool_owner=subscription.pool.owner,
+                bot=bot,
+            )
+            async for subscription in subscriptions_generator
+        ]
+        await asyncio.gather(*coroutines, return_exceptions=True)
+
+
+async def subscribe_to_new_pool_callback_handler(
+        callback_query: types.CallbackQuery,
+        user: User,
+        subscriptions_repository: SubscriptionsRepository,
+        pools_repository: PoolsRepository,
+):
+    callback_data = NotifyNewPoolCallbackData.unpack(
+        callback_query.data,
+    )
+
+    async with subscriptions_repository.transaction():
+        subscription = await subscriptions_repository.get_subscription(
+            from_user_id=user.id,
+            to_user_id=callback_data.user_id,
+        )
+    async with pools_repository.transaction():
+        pool = await pools_repository.get_pool_by_id(
+            pool_id=callback_data.pool_id,
+        )
+    if not all((subscription, pool)):
+        await callback_query.answer(text=constants.ERROR_MESSAGE)
+        return
+
+    async with subscriptions_repository.transaction():
+        subscription = await subscriptions_repository.set_pool_for_subscription(
+            from_user_id=subscription.from_user_id,
+            to_user_id=subscription.to_user_id,
+            pool_id=pool.id,
+        )
+    if subscription is None:
+        await callback_query.answer(text=constants.ERROR_MESSAGE)
+        return
+
+    await callback_query.answer(text="Вы присоединились к сбору")
+
+
+async def hide_new_pool_notification(
+        callback_query: types.CallbackQuery,
+):
+    await callback_query.message.delete()
--- a/birthday_pool_bot/telegram_bot/ui/logic.py
+++ b/birthday_pool_bot/telegram_bot/ui/logic.py
@@ -1,11 +1,11 @@
 import uuid

-from aiogram import types
+from aiogram import Bot, types
 from aiogram.enums import ParseMode
 from aiogram.utils.keyboard import InlineKeyboardBuilder, ReplyKeyboardBuilder
 from pydantic_filters import PagePagination

-from birthday_pool_bot.dto import BankEnum, User
+from birthday_pool_bot.dto import BankEnum, Pool, Subscription, User
 from birthday_pool_bot.repositories.repositories import (
    PoolsRepository,
    SubscriptionsRepository,
@@ -17,12 +17,14 @@ from .callback_data import (
    MenuCallbackData,
    NewSubscriptionCallbackData,
    NewSubscriptionConfirmCallbackData,
+    NotifyNewPoolCallbackData,
    PoolActionEnum,
    PoolCallbackData,
    PoolsCallbackData,
    PoolsBackCallbackData,
    SubscriptionActionEnum,
    SubscriptionCallbackData,
+    SubscribeToNewPoolAnswerEnum,
    SubscriptionsCallbackData,
 )
 from .exceptions import FlowInternalError
@@ -759,3 +761,44 @@ async def ask_new_subscription_confirmation(
        parse_mode=ParseMode.MARKDOWN,
        reply_markup=reply_markup,
    )
+
+
+async def notify_user_about_new_subscription_pool(
+        subscription: Subscription,
+        user: User,
+        pool: Pool,
+        pool_owner: User,
+        bot: Bot,
+):
+    if user.telegram_id is None:
+        return
+
+    text = (
+        f"Создан новый сбор для {subscription.name}\n\n"
+        "Собирает пользователь: **\n\n"
+        "Сбор денег происходит на:\n"
+        f"_Телефон_: {pool.payment_data.phone}\n"
+        f"_Банк_: {pool.payment_data.bank}\n\n"
+    )
+    keyboard = InlineKeyboardBuilder()
+    keyboard.button(
+        text="Подписаться",
+        callback_data=NotifyNewPoolCallbackData(
+            action=SubscribeToNewPoolAnswerEnum.YES,
+        ).pack(),
+    )
+    keyboard.button(
+        text="Закрыть",
+        callback_data=NotifyNewPoolCallbackData(
+            action=SubscribeToNewPoolAnswerEnum.NO,
+        ).pack()
+    )
+    keyboard.adjust(2)
+    reply_markup = keyboard.as_markup()
+
+    await bot.send_message(
+        chat_id=user.telegram_id,
+        text=text,
+        parse_mode=ParseMode.MARKDOWN,
+        reply_markup=reply_markup,
+    )
--- a/birthday_pool_bot/telegram_bot/ui/middlewares/init.py
+++ b/birthday_pool_bot/telegram_bot/ui/middlewares/init.py
@@ -1,5 +1,6 @@
 from .auth import AuthMiddleware
 from .depends import DependsMiddleware
+from .exception_handler import ExceptionHandlerMiddleware
 from .typing import TypingMiddleware


@@ -8,6 +9,8 @@ __all__ = (
    "AuthMiddleware",
    # depends
    "DependsMiddleware",
+    # exception_handler,
+    "ExceptionHandlerMiddleware",
    # typing
    "TypingMiddleware",
 )
--- a/birthday_pool_bot/telegram_bot/ui/middlewares/exception_handler.py
+++ b/birthday_pool_bot/telegram_bot/ui/middlewares/exception_handler.py
@@ -0,0 +1,31 @@
+import traceback
+from typing import Any, Awaitable, Callable
+
+import aiogram
+from aiogram import types
+from aiogram.enums import ParseMode
+
+from birthday_pool_bot.telegram_bot.ui import constants
+
+
+class ExceptionHandlerMiddleware(aiogram.BaseMiddleware):
+    async def __call__(
+            self,
+            handler: Callable[[types.TelegramObject, dict[str, Any]], Awaitable[Any]],
+            event: types.TelegramObject,
+            data: dict[str, Any],
+    ):
+        try:
+            await handler(event, data)
+        except Exception:
+            traceback.print_exc()
+            if isinstance(event, types.Message):
+                await event.reply(
+                    text=constants.ERROR_MESSAGE,
+                    parse_mode=ParseMode.MARKDOWN,
+                )
+            elif isinstance(event, types.CallbackQuery) is not None:
+                await event.answer(
+                    text=constants.ERROR_MESSAGE,
+                    parse_mode=ParseMode.MARKDOWN,
+                )
--- a/birthday_pool_bot/telegram_bot/ui/utils.py
+++ b/birthday_pool_bot/telegram_bot/ui/utils.py
@@ -1,7 +1,6 @@
 import datetime
 import re

-from aiogram import types
 from aiogram.types import User as TelegramUser
 from pydantic_core import PydanticCustomError
 from pydantic_extra_types.phone_numbers import PhoneNumber