Little fixes, add exception handler

README.md

@@ -18,13 +18,17 @@ Run
 uv run python -m birthday_pool_bot -e .env run
 ```
 
-## TODO
+## For test
 
 * Поддержать рассылку пользователям о том, что появились новые сборы для их подписки (если они ещё
 не участвуют в сборе)
+
+## TODO
+
 * Добавить возможность добавлять ссылку на сбор денег Сбера, когда происходит оповещение о
 наступающем дне рождения
 * Добавить раздел с настройками времени оповещения о предстоящих событиях
 * Добавить возможность инициировать сбор денег по другому поводу (кроме ДР)
 * Добавить возможность отписаться во время нотификации о предстоящем ДР
 * Добавить возможность синхронизировать календарь с днями рождения
+* Добавить возможность организовывать поздравления для участников группы

birthday_pool_bot/dto.py

@@ -73,3 +73,4 @@ class SubscriptionFilter(BaseFilter):
     to_user_id: set[uuid.UUID]
     to_user_id__gt: uuid.UUID
     pool_id: set[uuid.UUID]
+    pool_id__null: bool

birthday_pool_bot/notifications/service.py

@@ -35,20 +35,22 @@ class NotificationsService(facet.AsyncioServiceMixin):
         users_repository = self._repositories_container.users
         subscriptions_repository = self._repositories_container.subscriptions
 
-        async with users_repository.transaction():
+        users_generator = users_repository.get_users_by_birthdays(
-            users = await users_repository.get_users_by_birthdays(
             birthday=birthday,
         )
+        async with users_repository.transaction():
+            users = [user async for user in users_generator]
         if not users:
             return
 
-        async with subscriptions_repository.transaction():
+        subscriptions_generator = subscriptions_repository.get_to_users_subscriptions(
-            subscriptions = await subscriptions_repository.get_to_users_subscriptions(
             to_users_ids={user.id for user in users},
             with_from_user=True,
             with_to_user=True,
             with_pool=True,
         )
+        async with subscriptions_repository.transaction():
+            subscriptions = [subscription async for subscription in subscriptions_generator]
         if not subscriptions:
             return
 

birthday_pool_bot/repositories/repositories/subscriptions.py

@@ -54,12 +54,16 @@ class SubscriptionsRepository(BaseRepository[Subscription, SubscriptionFilter]):
     async def get_to_users_subscriptions(
             self,
             to_users_ids: Iterable[uuid.UUID],
+            have_pool: bool | None = None,
+            pagination: BasePagination | None = None,
             with_from_user: bool = False,
             with_to_user: bool = False,
             with_pool: bool = False,
             with_pool_owner: bool = False,
-    ) -> list[Subscription]:
+    ) -> AsyncGenerator[Subscription, None]:
         filter_ = SubscriptionFilter(to_user_id=set(to_users_ids))
+        if have_pool is not None:
+            filter_.pool_id__null = not have_pool
         options = []
         if with_from_user:
             options.append(joinedload(DBSubscription.from_user))
@@ -72,10 +76,12 @@ class SubscriptionsRepository(BaseRepository[Subscription, SubscriptionFilter]):
 
         subscriptions_generator = self.get_items(
             filter_=filter_,
+            pagination=pagination,
             options=options,
         )
 
-        return [subscription async for subscription in subscriptions_generator]
+        async for subscription in subscriptions_generator:
+            yield subscription
 
     async def get_subscription(
             self,
@@ -134,3 +140,20 @@ class SubscriptionsRepository(BaseRepository[Subscription, SubscriptionFilter]):
             to_user_id={to_user_id},
         )
         await self.delete_items(filter_=filter_)
+
+    async def set_pool_for_subscription(
+            self,
+            from_user_id: uuid.UUID,
+            to_user_id: uuid.UUID,
+            pool_id: uuid.UUID,
+    ) -> Subscription | None:
+        filter_ = SubscriptionFilter(
+            from_user_id={from_user_id},
+            to_user_id={to_user_id},
+        )
+        subscriptions_generator = self.update_items(
+            filter_=filter_,
+            pool_id=pool_id,
+        )
+        async for subscription in subscriptions_generator:
+            return subscription

birthday_pool_bot/repositories/repositories/users.py

@@ -1,8 +1,8 @@
 import uuid
 from datetime import date
-from typing import Container, Sequence, TypeVar
+from typing import AsyncGenerator, Sequence, TypeVar
 
-from pydantic_filters import BasePagination, OffsetPagination
+from pydantic_filters import OffsetPagination
 from sqlalchemy import Column, func, inspect, or_, select, update
 from sqlmodel import SQLModel
 
@@ -38,22 +38,12 @@ class UsersRepository(BaseRepository[User, UserFilter]):
         async for user in self.get_items(filter_=filter_, pagination=pagination):
             return user
 
-    async def get_users_by_ids(
-            self,
-            user_ids: Container[uuid.UUID],
-            pagination: BasePagination | None = None,
-    ) -> list[User]:
-        filter_ = UserFilter(id=set(user_ids))
-        users_generator = self.get_items(filter_=filter_, pagination=pagination)
-        users = [user async for user in users_generator]
-        return users
-
     async def get_users_by_primary_keys(
             self,
             user_id: uuid.UUID | None = None,
             telegram_id: int | None = None,
             phone: str | None = None,            
-    ) -> list[User]:
+    ) -> AsyncGenerator[User, None]:
         filters = []
         if user_id is not None:
             filters.append(DBUser.id == user_id)
@@ -62,24 +52,24 @@ class UsersRepository(BaseRepository[User, UserFilter]):
         if phone is not None:
             filters.append(DBUser.phone == phone)
         if not filters:
-            return []
+            return
 
-        statement = select(DBUser).where(or_(filters))
+        statement = select(DBUser).where(or_(*filters))
         result = await self.session.exec(statement)
-
+        for db_user in result.all():
-        return [db_user.to_item() for db_user in result.all()]
+            yield db_user.to_item()
 
     async def get_users_by_birthdays(
             self,
             birthday: date | None = None,
-    ) -> list[User]:
+    ) -> AsyncGenerator[User, None]:
         statement = select(DBUser).where(
             func.extract("month", DBUser.birthday) == birthday.month,
             func.extract("day", DBUser.birthday) == birthday.day,
         )
         result = await self.session.execute(statement)
-
+        for (db_user,) in result.all():
-        return [db_user.to_item() for (db_user,) in result.all()]
+            yield db_user.to_item()
 
     async def create_user(self, user: User) -> User:
         return await self.create_item(item=user)

birthday_pool_bot/telegram_bot/storage/fabric.py

@@ -1,4 +1,4 @@
-import redis.asyncio import redis
+import redis.asyncio as redis
 from aiogram.fsm.storage.base import BaseStorage
 from aiogram.fsm.storage.memory import MemoryStorage
 from aiogram.fsm.storage.redis import RedisStorage
@@ -12,14 +12,12 @@ def get_telegram_bot_fsm_storage(settings: StorageSettings) -> BaseStorage:
         case StorageTypeEnum.MEMORY.value:
             return MemoryStorage()
         case StorageTypeEnum.REDIS.value:
-            redis_connection_pool = redis.ConnectionPool(
+            redis_client = redis.Redis(
                 host=settings.host,
                 port=settings.port,
                 db=settings.database,
                 password=settings.password,
-                max_connections=settings.max_connections,
             )
-            redis_client = redis.Redis(connection_pool=redis_connection_pool)
             return RedisStorage(redis=redis_client)
         case _:
             raise ValueError(f"Unsupported telegram bot storage type: {settings.type}")

birthday_pool_bot/telegram_bot/storage/settings.py

@@ -19,4 +19,3 @@ class RedisStorageSettings(StorageSettings):
     port: int = 6379
     database: int = 0
     password: str | None = None
-    max_connections: int = 10

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

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)
 

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
@@ -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 callback_query.answer(text=constants.ERROR_MESSAGE)
-        await logic.ask_new_subscription_user(message=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 callback_query.answer(text=constants.ERROR_MESSAGE)
-        await logic.ask_new_subscription_user(message=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()

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,
+    )

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",
 )

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,
+                )

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