From 15be932b3b5d2972086d84aad86e5e34ca340332 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 26 Nov 2023 22:30:50 +0100 Subject: [PATCH 01/44] use thread id in storage --- src/aiogram_dialog/context/intent_middleware.py | 2 ++ src/aiogram_dialog/context/storage.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index 49978260..bcd7d5eb 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -37,6 +37,8 @@ def storage_proxy(self, data: dict): storage=data["fsm_storage"], user_id=data["event_from_user"].id, chat_id=data["event_chat"].id, + chat_type=data["event_chat"].type, + thread_id=data.get("event_thread_id"), state_groups=self.registry.state_groups(), ) return proxy diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 47d30d7b..c8705e5d 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -17,6 +17,8 @@ def __init__( storage: BaseStorage, user_id: int, chat_id: int, + chat_type: str, + thread_id: Optional[int], bot: Bot, state_groups: Dict[str, Type[StatesGroup]], ): @@ -24,6 +26,8 @@ def __init__( self.state_groups = state_groups self.user_id = user_id self.chat_id = chat_id + self.thread_id = thread_id + self.chat_type = chat_type self.bot = bot async def load_context(self, intent_id: str) -> Context: @@ -87,6 +91,7 @@ def _context_key(self, intent_id: str) -> StorageKey: bot_id=self.bot.id, chat_id=self.chat_id, user_id=self.user_id, + thread_id=self.thread_id, destiny=f"aiogd:context:{intent_id}", ) From 2a87bb85b803211c20df8b442d2d9712134125ba Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 26 Nov 2023 22:47:27 +0100 Subject: [PATCH 02/44] save access settings --- src/aiogram_dialog/api/entities/__init__.py | 4 +-- src/aiogram_dialog/api/entities/stack.py | 11 +++++- src/aiogram_dialog/context/storage.py | 37 +++++++++++++++++++-- 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/src/aiogram_dialog/api/entities/__init__.py b/src/aiogram_dialog/api/entities/__init__.py index 97c8b650..cf88fe8b 100644 --- a/src/aiogram_dialog/api/entities/__init__.py +++ b/src/aiogram_dialog/api/entities/__init__.py @@ -5,7 +5,7 @@ "MediaAttachment", "MediaId", "ShowMode", "StartMode", "MarkupVariant", "NewMessage", "OldMessage", "UnknownText", - "DEFAULT_STACK_ID", "Stack", + "AccessSettings", "DEFAULT_STACK_ID", "Stack", "DIALOG_EVENT_NAME", "DialogAction", "DialogUpdateEvent", "DialogStartEvent", "DialogSwitchEvent", "DialogUpdate", ] @@ -16,7 +16,7 @@ from .media import MediaAttachment, MediaId from .modes import ShowMode, StartMode from .new_message import MarkupVariant, NewMessage, OldMessage, UnknownText -from .stack import DEFAULT_STACK_ID, Stack +from .stack import AccessSettings, DEFAULT_STACK_ID, Stack from .update_event import ( DIALOG_EVENT_NAME, DialogAction, DialogStartEvent, DialogSwitchEvent, DialogUpdate, DialogUpdateEvent, diff --git a/src/aiogram_dialog/api/entities/stack.py b/src/aiogram_dialog/api/entities/stack.py index 21440771..a2ef9dda 100644 --- a/src/aiogram_dialog/api/entities/stack.py +++ b/src/aiogram_dialog/api/entities/stack.py @@ -2,8 +2,9 @@ import string import time from dataclasses import dataclass, field -from typing import List, Optional +from typing import Any, List, Optional +from aiogram.enums import ChatMemberStatus from aiogram.fsm.state import State from aiogram_dialog.api.exceptions import DialogStackOverflow @@ -33,6 +34,13 @@ def new_id(): return id_to_str(new_int_id()) +@dataclass +class AccessSettings: + user_ids: List[int] + member_status: Optional[ChatMemberStatus] + custom: Any = None + + @dataclass(unsafe_hash=True) class Stack: _id: str = field(compare=True, default_factory=new_id) @@ -44,6 +52,7 @@ class Stack: last_income_media_group_id: Optional[str] = field( compare=False, default=None, ) + access_settings: Optional[AccessSettings] = None @property def id(self): diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index c8705e5d..55ee2441 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -2,11 +2,12 @@ from typing import Dict, Optional, Type from aiogram import Bot +from aiogram.enums import ChatMemberStatus from aiogram.fsm.state import State, StatesGroup from aiogram.fsm.storage.base import BaseStorage, StorageKey from aiogram_dialog.api.entities import ( - Context, DEFAULT_STACK_ID, Stack, + AccessSettings, Context, DEFAULT_STACK_ID, Stack, ) from aiogram_dialog.api.exceptions import UnknownIntent, UnknownState @@ -47,7 +48,11 @@ async def load_stack(self, stack_id: str = DEFAULT_STACK_ID) -> Stack: ) if not data: return Stack(_id=stack_id) - return Stack(**data) + + access_settings = self._parse_access_settings( + data.pop("access_settings"), + ) + return Stack(access_settings=access_settings, **data) async def save_context(self, context: Optional[Context]) -> None: if not context: @@ -81,6 +86,7 @@ async def save_stack(self, stack: Optional[Stack]) -> None: ) else: data = copy(vars(stack)) + data["access_settings"] = self._dump_access_settings(stack.access_settings) await self.storage.set_data( key=self._stack_key(stack.id), data=data, @@ -100,6 +106,7 @@ def _stack_key(self, stack_id: str) -> StorageKey: bot_id=self.bot.id, chat_id=self.chat_id, user_id=self.user_id, + thread_id=self.thread_id, destiny=f"aiogd:stack:{stack_id}", ) @@ -112,3 +119,29 @@ def _state(self, state: str) -> State: except KeyError: raise UnknownState(f"Unknown state group {group}") raise UnknownState(f"Unknown state {state}") + + def _parse_access_settings( + self, raw: Optional[Dict], + ) -> Optional[AccessSettings]: + if not raw: + return None + if raw_member_status := raw.get("member_status"): + member_status = ChatMemberStatus(raw_member_status) + else: + member_status = None + return AccessSettings( + user_ids=raw.get("user_ids") or [], + member_status=member_status, + custom=raw.get("custom") + ) + + def _dump_access_settings( + self, access_settings: Optional[AccessSettings], + ) -> Optional[Dict]: + if not access_settings: + return None + return { + "user_ids": access_settings.user_ids, + "member_status": access_settings.member_status, + "custom": access_settings.custom, + } From 9913b417429947d73f353c8d0e09d8b690fbed71 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 26 Nov 2023 22:49:50 +0100 Subject: [PATCH 03/44] add thread id to error middleware --- src/aiogram_dialog/context/intent_middleware.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index bcd7d5eb..5eda732f 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -267,6 +267,8 @@ async def __call__( storage=data["fsm_storage"], user_id=user.id, chat_id=chat.id, + chat_type=chat.type, + thread_id=data.get("event_thread_id"), state_groups=self.registry.state_groups(), ) data[STORAGE_KEY] = proxy From b11af4afe57643b6cd25f1d5f638e6f9cf82f91e Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 26 Nov 2023 23:02:44 +0100 Subject: [PATCH 04/44] Access validator --- .../context/intent_middleware.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index 5eda732f..7afbf4cd 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -3,6 +3,7 @@ from aiogram import Router from aiogram.dispatcher.middlewares.base import BaseMiddleware +from aiogram.enums import ChatType from aiogram.types import CallbackQuery, Chat, Message, User from aiogram.types.error_event import ErrorEvent @@ -23,6 +24,20 @@ logger = getLogger(__name__) +class AccessValidator: + def is_allowed( + self, stack: Stack, user: User, chat: Chat, + ) -> bool: + if not stack.access_settings: + return True + if chat.type is ChatType.PRIVATE: + return True + if stack.access_settings.user_ids: + if user.id not in stack.access_settings.user_ids: + return False + return True + + class IntentMiddlewareFactory: def __init__( self, @@ -30,6 +45,7 @@ def __init__( ): super().__init__() self.registry = registry + self.access_validator = AccessValidator() # TODO: inject def storage_proxy(self, data: dict): proxy = StorageProxy( @@ -88,6 +104,11 @@ async def _load_context( raise InvalidStackIdError( f"Both stack id and intent id are None: {event}", ) + + if not self.access_validator.is_allowed( + stack, event.from_user, data["event_chat"], + ): + return data[STORAGE_KEY] = proxy data[STACK_KEY] = stack data[CONTEXT_KEY] = context From 2283cfea4fe846bc7a110a5e193b94174d5e2f1e Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Tue, 28 Nov 2023 23:06:15 +0100 Subject: [PATCH 05/44] Stack access validator protocol --- src/aiogram_dialog/api/protocols/__init__.py | 2 ++ .../api/protocols/stack_access.py | 17 +++++++++++++++++ src/aiogram_dialog/context/intent_middleware.py | 9 +++++---- 3 files changed, 24 insertions(+), 4 deletions(-) create mode 100644 src/aiogram_dialog/api/protocols/stack_access.py diff --git a/src/aiogram_dialog/api/protocols/__init__.py b/src/aiogram_dialog/api/protocols/__init__.py index 77571aea..21610eaf 100644 --- a/src/aiogram_dialog/api/protocols/__init__.py +++ b/src/aiogram_dialog/api/protocols/__init__.py @@ -4,6 +4,7 @@ "MediaIdStorageProtocol", "MessageManagerProtocol", "MessageNotModified", "DialogProtocol", "DialogRegistryProtocol", + "StackAccessValidator", ] from .dialog import DialogProtocol @@ -11,3 +12,4 @@ from .media import MediaIdStorageProtocol from .message_manager import MessageManagerProtocol, MessageNotModified from .registry import DialogRegistryProtocol +from .stack_access import StackAccessValidator diff --git a/src/aiogram_dialog/api/protocols/stack_access.py b/src/aiogram_dialog/api/protocols/stack_access.py new file mode 100644 index 00000000..3a5fd4b3 --- /dev/null +++ b/src/aiogram_dialog/api/protocols/stack_access.py @@ -0,0 +1,17 @@ +from abc import abstractmethod +from typing import Protocol + +from aiogram.types import User, Chat + +from aiogram_dialog.api.entities import Stack + + +class StackAccessValidator(Protocol): + @abstractmethod + async def is_allowed( + self, stack: Stack, user: User, chat: Chat, + ) -> bool: + raise NotImplementedError + + @abstractmethod + async def \ No newline at end of file diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index 7afbf4cd..a40e03d6 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -17,7 +17,7 @@ CALLBACK_DATA_KEY, CONTEXT_KEY, EVENT_SIMULATED, STACK_KEY, STORAGE_KEY, ) -from aiogram_dialog.api.protocols import DialogRegistryProtocol +from aiogram_dialog.api.protocols import DialogRegistryProtocol, StackAccessValidator from aiogram_dialog.utils import remove_indent_id, split_reply_callback from .storage import StorageProxy @@ -25,7 +25,7 @@ class AccessValidator: - def is_allowed( + async def is_allowed( self, stack: Stack, user: User, chat: Chat, ) -> bool: if not stack.access_settings: @@ -45,7 +45,7 @@ def __init__( ): super().__init__() self.registry = registry - self.access_validator = AccessValidator() # TODO: inject + self.access_validator: StackAccessValidator = AccessValidator() # TODO: inject def storage_proxy(self, data: dict): proxy = StorageProxy( @@ -82,6 +82,7 @@ async def _load_context( data: dict, ) -> None: proxy = self.storage_proxy(data) + chat = data["event_chat"] logger.debug( "Loading context for intent: `%s`, " "stack: `%s`, user: `%s`, chat: `%s`", @@ -106,7 +107,7 @@ async def _load_context( ) if not self.access_validator.is_allowed( - stack, event.from_user, data["event_chat"], + stack, event.from_user, chat, ): return data[STORAGE_KEY] = proxy From 00f00ef8c43f833de95abf30eb0fe1e5332ec003 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 29 Nov 2023 00:42:55 +0100 Subject: [PATCH 06/44] refactor intent middleware --- .../api/protocols/stack_access.py | 5 +- .../context/intent_middleware.py | 126 +++++++++++++----- src/aiogram_dialog/context/storage.py | 10 +- 3 files changed, 96 insertions(+), 45 deletions(-) diff --git a/src/aiogram_dialog/api/protocols/stack_access.py b/src/aiogram_dialog/api/protocols/stack_access.py index 3a5fd4b3..5f156907 100644 --- a/src/aiogram_dialog/api/protocols/stack_access.py +++ b/src/aiogram_dialog/api/protocols/stack_access.py @@ -1,7 +1,7 @@ from abc import abstractmethod from typing import Protocol -from aiogram.types import User, Chat +from aiogram.types import Chat, User from aiogram_dialog.api.entities import Stack @@ -12,6 +12,3 @@ async def is_allowed( self, stack: Stack, user: User, chat: Chat, ) -> bool: raise NotImplementedError - - @abstractmethod - async def \ No newline at end of file diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index a40e03d6..e2730258 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -17,7 +17,9 @@ CALLBACK_DATA_KEY, CONTEXT_KEY, EVENT_SIMULATED, STACK_KEY, STORAGE_KEY, ) -from aiogram_dialog.api.protocols import DialogRegistryProtocol, StackAccessValidator +from aiogram_dialog.api.protocols import ( + DialogRegistryProtocol, StackAccessValidator, +) from aiogram_dialog.utils import remove_indent_id, split_reply_callback from .storage import StorageProxy @@ -42,10 +44,14 @@ class IntentMiddlewareFactory: def __init__( self, registry: DialogRegistryProtocol, + access_validator: Optional[StackAccessValidator] = None, ): super().__init__() self.registry = registry - self.access_validator: StackAccessValidator = AccessValidator() # TODO: inject + if access_validator: + self.access_validator = access_validator + else: + self.access_validator = AccessValidator() def storage_proxy(self, data: dict): proxy = StorageProxy( @@ -53,7 +59,6 @@ def storage_proxy(self, data: dict): storage=data["fsm_storage"], user_id=data["event_from_user"].id, chat_id=data["event_chat"].id, - chat_type=data["event_chat"].type, thread_id=data.get("event_thread_id"), state_groups=self.registry.state_groups(), ) @@ -74,46 +79,73 @@ def _check_outdated(self, intent_id: str, stack: Stack): f"for stack ({stack.id})", ) - async def _load_context( + async def _load_stack( self, - event: ChatEvent, - intent_id: Optional[str], + stack_id: Optional[str], + proxy: StorageProxy, + user: User, + chat: Chat, + ) -> Optional[Stack]: + if stack_id is None: + raise InvalidStackIdError("Both stack id and intent id are None") + stack = await proxy.load_stack(stack_id) + if not await self.access_validator.is_allowed(stack, user, chat): + return + return stack + + async def _load_context_by_stack( + self, + proxy: StorageProxy, stack_id: Optional[str], data: dict, ) -> None: - proxy = self.storage_proxy(data) + user = data["event_from_user"] chat = data["event_chat"] logger.debug( - "Loading context for intent: `%s`, " - "stack: `%s`, user: `%s`, chat: `%s`", - intent_id, - stack_id, - event.from_user.id, - proxy.chat_id, + "Loading context for stack: `%s`, user: `%s`, chat: `%s`", + stack_id, user.id, chat.id, ) - if intent_id is not None: - context = await proxy.load_context(intent_id) - stack = await proxy.load_stack(context.stack_id) - self._check_outdated(intent_id, stack) - elif stack_id is not None: - stack = await proxy.load_stack(stack_id) - if stack.empty(): - context = None - else: - context = await proxy.load_context(stack.last_intent_id()) + stack = await self._load_stack(stack_id, proxy, user, chat) + if not stack: + return + if stack.empty(): + context = None else: - raise InvalidStackIdError( - f"Both stack id and intent id are None: {event}", - ) + context = await proxy.load_context(stack.last_intent_id()) + data[STORAGE_KEY] = proxy + data[STACK_KEY] = stack + data[CONTEXT_KEY] = context - if not self.access_validator.is_allowed( - stack, event.from_user, chat, - ): + async def _load_context_by_intent( + self, + proxy: StorageProxy, + intent_id: Optional[str], + data: dict, + ) -> None: + user = data["event_from_user"] + chat = data["event_chat"] + logger.debug( + "Loading context for intent: `%s`, user: `%s`, chat: `%s`", + intent_id, user.id, chat.id, + ) + context = await proxy.load_context(intent_id) + stack = await self._load_stack(context.stack_id, proxy, user, chat) + if not stack: return + self._check_outdated(intent_id, stack) + data[STORAGE_KEY] = proxy data[STACK_KEY] = stack data[CONTEXT_KEY] = context + async def _load_default_context( + self, event: ChatEvent, data: dict, + ) -> None: + proxy = self.storage_proxy(data) + return await self._load_context_by_stack( + proxy=proxy, stack_id=DEFAULT_STACK_ID, data=data, + ) + def _intent_id_from_reply( self, event: Message, data: dict, ) -> Optional[str]: @@ -156,9 +188,13 @@ async def process_message( ) if intent_id := self._intent_id_from_reply(event, data): - await self._load_context(event, intent_id, DEFAULT_STACK_ID, data) + await self._load_context_by_intent( + proxy=self.storage_proxy(data), + intent_id=intent_id, + data=data, + ) else: - await self._load_context(event, None, DEFAULT_STACK_ID, data) + await self._load_default_context(event, data) return await handler(event, data) async def process_my_chat_member( @@ -167,7 +203,7 @@ async def process_my_chat_member( event: Message, data: dict, ) -> None: - await self._load_context(event, None, DEFAULT_STACK_ID, data) + await self._load_default_context(event, data) return await handler(event, data) async def process_chat_join_request( @@ -176,7 +212,7 @@ async def process_chat_join_request( event: Message, data: dict, ) -> None: - await self._load_context(event, None, DEFAULT_STACK_ID, data) + await self._load_default_context(event, data) return await handler(event, data) async def process_aiogd_update( @@ -185,7 +221,18 @@ async def process_aiogd_update( event: DialogUpdateEvent, data: dict, ): - await self._load_context(event, event.intent_id, event.stack_id, data) + if event.intent_id: + await self._load_context_by_intent( + proxy=self.storage_proxy(data), + intent_id=event.intent_id, + data=data, + ) + else: + await self._load_context_by_stack( + proxy=self.storage_proxy(data), + stack_id=event.stack_id, + data=data, + ) return await handler(event, data) async def process_callback_query( @@ -202,10 +249,17 @@ async def process_callback_query( original_data = event.data if event.data: intent_id, callback_data = remove_indent_id(event.data) - await self._load_context(event, intent_id, DEFAULT_STACK_ID, data) + if intent_id: + await self._load_context_by_intent( + proxy=self.storage_proxy(data), + intent_id=intent_id, + data=data, + ) + else: + await self._load_default_context(event, data) data[CALLBACK_DATA_KEY] = original_data else: - await self._load_context(event, None, DEFAULT_STACK_ID, data) + await self._load_default_context(event, data) return await handler(event, data) diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 55ee2441..05d8b8ca 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -16,9 +16,8 @@ class StorageProxy: def __init__( self, storage: BaseStorage, - user_id: int, + user_id: Optional[int], chat_id: int, - chat_type: str, thread_id: Optional[int], bot: Bot, state_groups: Dict[str, Type[StatesGroup]], @@ -28,7 +27,6 @@ def __init__( self.user_id = user_id self.chat_id = chat_id self.thread_id = thread_id - self.chat_type = chat_type self.bot = bot async def load_context(self, intent_id: str) -> Context: @@ -86,7 +84,9 @@ async def save_stack(self, stack: Optional[Stack]) -> None: ) else: data = copy(vars(stack)) - data["access_settings"] = self._dump_access_settings(stack.access_settings) + data["access_settings"] = self._dump_access_settings( + stack.access_settings, + ) await self.storage.set_data( key=self._stack_key(stack.id), data=data, @@ -132,7 +132,7 @@ def _parse_access_settings( return AccessSettings( user_ids=raw.get("user_ids") or [], member_status=member_status, - custom=raw.get("custom") + custom=raw.get("custom"), ) def _dump_access_settings( From 6450c84481be7fd09e557662fe7773d5e68e36e9 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Tue, 5 Dec 2023 17:58:57 +0100 Subject: [PATCH 07/44] special default stack id generation --- example/mega/bot.py | 6 +- pyproject.toml | 2 +- src/aiogram_dialog/api/entities/stack.py | 2 +- .../context/intent_middleware.py | 15 ++-- src/aiogram_dialog/context/storage.py | 19 ++++- src/aiogram_dialog/manager/manager.py | 3 + tests/test_group.py | 82 +++++++++++++++++++ 7 files changed, 116 insertions(+), 13 deletions(-) create mode 100644 tests/test_group.py diff --git a/example/mega/bot.py b/example/mega/bot.py index e1d1eabd..05f3f30d 100644 --- a/example/mega/bot.py +++ b/example/mega/bot.py @@ -24,7 +24,11 @@ async def start(message: Message, dialog_manager: DialogManager): # it is important to reset stack because user wants to restart everything - await dialog_manager.start(states.Main.MAIN, mode=StartMode.RESET_STACK) + await dialog_manager.start( + states.Main.MAIN, + mode=StartMode.RESET_STACK, + show_mode=ShowMode.SEND, + ) async def on_unknown_intent(event: ErrorEvent, dialog_manager: DialogManager): diff --git a/pyproject.toml b/pyproject.toml index fed3d5bb..11284d31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - 'aiogram>=3.0.0', + 'aiogram>=3.2.0', 'jinja2', 'cachetools>=4.0.0,<6.0.0', 'magic_filter', diff --git a/src/aiogram_dialog/api/entities/stack.py b/src/aiogram_dialog/api/entities/stack.py index a2ef9dda..37bf2da3 100644 --- a/src/aiogram_dialog/api/entities/stack.py +++ b/src/aiogram_dialog/api/entities/stack.py @@ -37,7 +37,7 @@ def new_id(): @dataclass class AccessSettings: user_ids: List[int] - member_status: Optional[ChatMemberStatus] + member_status: Optional[ChatMemberStatus] = None custom: Any = None diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index e2730258..4f1118e6 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -57,10 +57,10 @@ def storage_proxy(self, data: dict): proxy = StorageProxy( bot=data["bot"], storage=data["fsm_storage"], + state_groups=self.registry.state_groups(), user_id=data["event_from_user"].id, chat_id=data["event_chat"].id, thread_id=data.get("event_thread_id"), - state_groups=self.registry.state_groups(), ) return proxy @@ -90,6 +90,10 @@ async def _load_stack( raise InvalidStackIdError("Both stack id and intent id are None") stack = await proxy.load_stack(stack_id) if not await self.access_validator.is_allowed(stack, user, chat): + logger.debug( + "Stack %s is not allowed for user %s", + stack.id, user.id, + ) return return stack @@ -141,9 +145,10 @@ async def _load_context_by_intent( async def _load_default_context( self, event: ChatEvent, data: dict, ) -> None: - proxy = self.storage_proxy(data) return await self._load_context_by_stack( - proxy=proxy, stack_id=DEFAULT_STACK_ID, data=data, + proxy=self.storage_proxy(data), + stack_id=DEFAULT_STACK_ID, + data=data, ) def _intent_id_from_reply( @@ -243,9 +248,6 @@ async def process_callback_query( ): if "event_chat" not in data: return await handler(event, data) - proxy = self.storage_proxy(data) - data[STORAGE_KEY] = proxy - original_data = event.data if event.data: intent_id, callback_data = remove_indent_id(event.data) @@ -343,7 +345,6 @@ async def __call__( storage=data["fsm_storage"], user_id=user.id, chat_id=chat.id, - chat_type=chat.type, thread_id=data.get("event_thread_id"), state_groups=self.registry.state_groups(), ) diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 05d8b8ca..1dc27453 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -41,11 +41,16 @@ async def load_context(self, intent_id: str) -> Context: return Context(**data) async def load_stack(self, stack_id: str = DEFAULT_STACK_ID) -> Stack: + stack_id = self._fixed_stack_id(stack_id) data = await self.storage.get_data( key=self._stack_key(stack_id), ) if not data: - return Stack(_id=stack_id) + if self.user_id: + default_access = AccessSettings(user_ids=[self.user_id]) + else: + default_access = AccessSettings(user_ids=[]) + return Stack(_id=stack_id, access_settings=default_access) access_settings = self._parse_access_settings( data.pop("access_settings"), @@ -96,16 +101,24 @@ def _context_key(self, intent_id: str) -> StorageKey: return StorageKey( bot_id=self.bot.id, chat_id=self.chat_id, - user_id=self.user_id, + user_id=self.chat_id, thread_id=self.thread_id, destiny=f"aiogd:context:{intent_id}", ) + def _fixed_stack_id(self, stack_id: str) -> str: + if stack_id != DEFAULT_STACK_ID: + return stack_id + if self.user_id in (None, self.chat_id): + return stack_id + return f"<{self.user_id}>" + def _stack_key(self, stack_id: str) -> StorageKey: + stack_id = self._fixed_stack_id(stack_id) return StorageKey( bot_id=self.bot.id, chat_id=self.chat_id, - user_id=self.user_id, + user_id=self.chat_id, thread_id=self.thread_id, destiny=f"aiogd:stack:{stack_id}", ) diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index 99e614c1..e18bed54 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -2,6 +2,7 @@ from typing import Any, Dict, Optional from aiogram import Router +from aiogram.enums import ChatType from aiogram.fsm.state import State from aiogram.types import ( CallbackQuery, Chat, ErrorEvent, Message, ReplyKeyboardMarkup, User, @@ -404,6 +405,8 @@ def _save_last_message(self, message: OldMessage): def _calc_show_mode(self) -> ShowMode: if self.show_mode is not ShowMode.AUTO: return self.show_mode + if self.middleware_data["event_chat"].type == ChatType.GROUP: + return ShowMode.EDIT if self.current_stack().last_reply_keyboard: return ShowMode.DELETE_AND_SEND if self.current_stack().id != DEFAULT_STACK_ID: diff --git a/tests/test_group.py b/tests/test_group.py new file mode 100644 index 00000000..971d07a5 --- /dev/null +++ b/tests/test_group.py @@ -0,0 +1,82 @@ +from typing import Any + +import pytest +from aiogram import Dispatcher +from aiogram.filters import CommandStart +from aiogram.fsm.state import State, StatesGroup + +from aiogram_dialog import ( + Dialog, DialogManager, setup_dialogs, StartMode, Window, +) +from aiogram_dialog.test_tools import BotClient, MockMessageManager +from aiogram_dialog.test_tools.keyboard import InlineButtonTextLocator +from aiogram_dialog.widgets.kbd import Button +from aiogram_dialog.widgets.text import Format, Const + + +class MainSG(StatesGroup): + start = State() + + +window = Window( + Format("stub"), + Button(Const("Button"), id="btn"), + state=MainSG.start, +) + + +async def start(event: Any, dialog_manager: DialogManager): + await dialog_manager.start(MainSG.start, mode=StartMode.RESET_STACK) + + +@pytest.fixture() +def message_manager(): + return MockMessageManager() + + +@pytest.fixture() +def dp(message_manager): + dp = Dispatcher() + dp.include_router(Dialog(window)) + setup_dialogs(dp, message_manager=message_manager) + return dp + + +@pytest.fixture() +def client(dp): + return BotClient(dp, chat_id=-1, user_id=1, chat_type="group") + + +@pytest.fixture() +def second_client(dp): + return BotClient(dp, chat_id=-1, user_id=2, chat_type="group") + + +@pytest.mark.asyncio +async def test_second_user(dp, client, second_client, message_manager): + dp.message.register(start, CommandStart()) + await client.send("/start") + first_message = message_manager.one_message() + assert first_message.text == "stub" + message_manager.reset_history() + await second_client.send("test") + assert not message_manager.sent_messages + await second_client.click( + first_message, InlineButtonTextLocator("Button"), + ) + assert not message_manager.sent_messages + +@pytest.mark.asyncio +async def test_same_user(dp, client, message_manager): + dp.message.register(start, CommandStart()) + await client.send("/start") + first_message = message_manager.one_message() + assert first_message.text == "stub" + message_manager.reset_history() + await client.send("test") + assert first_message.text == "stub" + message_manager.reset_history() + await client.click( + first_message, InlineButtonTextLocator("Button"), + ) + assert first_message.text == "stub" From b6821e44fd6840b926a09d9f36a4a6cfe5b68b34 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 6 Dec 2023 00:17:05 +0100 Subject: [PATCH 08/44] default group stack id --- example/subdialog.py | 1 + src/aiogram_dialog/api/entities/__init__.py | 4 +-- src/aiogram_dialog/api/entities/stack.py | 1 + src/aiogram_dialog/context/storage.py | 19 ++++++++------- tests/test_group.py | 27 +++++++++++++++++++++ 5 files changed, 41 insertions(+), 11 deletions(-) diff --git a/example/subdialog.py b/example/subdialog.py index c54dbb5f..02527c31 100644 --- a/example/subdialog.py +++ b/example/subdialog.py @@ -13,6 +13,7 @@ Data, Dialog, DialogManager, Window, StartMode, setup_dialogs, ) +from aiogram_dialog.api.entities import GROUP_STACK_ID from aiogram_dialog.tools import render_transitions, render_preview from aiogram_dialog.widgets.input import MessageInput from aiogram_dialog.widgets.kbd import ( diff --git a/src/aiogram_dialog/api/entities/__init__.py b/src/aiogram_dialog/api/entities/__init__.py index cf88fe8b..01c89236 100644 --- a/src/aiogram_dialog/api/entities/__init__.py +++ b/src/aiogram_dialog/api/entities/__init__.py @@ -5,7 +5,7 @@ "MediaAttachment", "MediaId", "ShowMode", "StartMode", "MarkupVariant", "NewMessage", "OldMessage", "UnknownText", - "AccessSettings", "DEFAULT_STACK_ID", "Stack", + "AccessSettings", "DEFAULT_STACK_ID", "GROUP_STACK_ID", "Stack", "DIALOG_EVENT_NAME", "DialogAction", "DialogUpdateEvent", "DialogStartEvent", "DialogSwitchEvent", "DialogUpdate", ] @@ -16,7 +16,7 @@ from .media import MediaAttachment, MediaId from .modes import ShowMode, StartMode from .new_message import MarkupVariant, NewMessage, OldMessage, UnknownText -from .stack import AccessSettings, DEFAULT_STACK_ID, Stack +from .stack import AccessSettings, DEFAULT_STACK_ID, GROUP_STACK_ID, Stack from .update_event import ( DIALOG_EVENT_NAME, DialogAction, DialogStartEvent, DialogSwitchEvent, DialogUpdate, DialogUpdateEvent, diff --git a/src/aiogram_dialog/api/entities/stack.py b/src/aiogram_dialog/api/entities/stack.py index 37bf2da3..cd5d46f5 100644 --- a/src/aiogram_dialog/api/entities/stack.py +++ b/src/aiogram_dialog/api/entities/stack.py @@ -11,6 +11,7 @@ from .context import Context, Data DEFAULT_STACK_ID = "" +GROUP_STACK_ID = "<->" _STACK_LIMIT = 100 _ID_SYMS = string.digits + string.ascii_letters diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 1dc27453..e839a156 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -40,17 +40,18 @@ async def load_context(self, intent_id: str) -> Context: data["state"] = self._state(data["state"]) return Context(**data) + def _default_access_settings(self, stack_id: str) -> AccessSettings: + if stack_id == DEFAULT_STACK_ID: + return AccessSettings(user_ids=[self.user_id]) + else: + return AccessSettings(user_ids=[]) + async def load_stack(self, stack_id: str = DEFAULT_STACK_ID) -> Stack: - stack_id = self._fixed_stack_id(stack_id) - data = await self.storage.get_data( - key=self._stack_key(stack_id), - ) + fixed_stack_id = self._fixed_stack_id(stack_id) + data = await self.storage.get_data(self._stack_key(fixed_stack_id)) if not data: - if self.user_id: - default_access = AccessSettings(user_ids=[self.user_id]) - else: - default_access = AccessSettings(user_ids=[]) - return Stack(_id=stack_id, access_settings=default_access) + access_settings = self._default_access_settings(stack_id) + return Stack(_id=fixed_stack_id, access_settings=access_settings) access_settings = self._parse_access_settings( data.pop("access_settings"), diff --git a/tests/test_group.py b/tests/test_group.py index 971d07a5..1b4d88af 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -1,3 +1,4 @@ +import asyncio from typing import Any import pytest @@ -8,6 +9,7 @@ from aiogram_dialog import ( Dialog, DialogManager, setup_dialogs, StartMode, Window, ) +from aiogram_dialog.api.entities import GROUP_STACK_ID from aiogram_dialog.test_tools import BotClient, MockMessageManager from aiogram_dialog.test_tools.keyboard import InlineButtonTextLocator from aiogram_dialog.widgets.kbd import Button @@ -29,6 +31,11 @@ async def start(event: Any, dialog_manager: DialogManager): await dialog_manager.start(MainSG.start, mode=StartMode.RESET_STACK) +async def start_shared(event: Any, dialog_manager: DialogManager): + dialog_manager = dialog_manager.bg(stack_id=GROUP_STACK_ID) + await dialog_manager.start(MainSG.start, mode=StartMode.RESET_STACK) + + @pytest.fixture() def message_manager(): return MockMessageManager() @@ -66,6 +73,7 @@ async def test_second_user(dp, client, second_client, message_manager): ) assert not message_manager.sent_messages + @pytest.mark.asyncio async def test_same_user(dp, client, message_manager): dp.message.register(start, CommandStart()) @@ -74,9 +82,28 @@ async def test_same_user(dp, client, message_manager): assert first_message.text == "stub" message_manager.reset_history() await client.send("test") + first_message = message_manager.one_message() assert first_message.text == "stub" message_manager.reset_history() await client.click( first_message, InlineButtonTextLocator("Button"), ) + first_message = message_manager.one_message() + assert first_message.text == "stub" + + +@pytest.mark.asyncio +async def test_shared_stack(dp, client, second_client, message_manager): + dp.message.register(start_shared, CommandStart()) + await client.send("/start") + await asyncio.sleep(0.01) # synchronization workaround, fixme + first_message = message_manager.one_message() assert first_message.text == "stub" + message_manager.reset_history() + await second_client.send("test") + assert not message_manager.sent_messages + await second_client.click( + first_message, InlineButtonTextLocator("Button"), + ) + second_message = message_manager.one_message() + assert second_message.text == "stub" From e60798af1bf1154eb6f54dad1908edfa4592bed7 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 6 Dec 2023 01:18:21 +0100 Subject: [PATCH 09/44] Add access settings to start method --- .../api/entities/update_event.py | 2 + src/aiogram_dialog/api/protocols/manager.py | 2 + src/aiogram_dialog/manager/bg_manager.py | 3 ++ src/aiogram_dialog/manager/manager.py | 44 +++++++++++++------ src/aiogram_dialog/manager/sub_manager.py | 16 +++++-- src/aiogram_dialog/manager/update_handler.py | 9 ++-- src/aiogram_dialog/tools/preview.py | 2 + 7 files changed, 58 insertions(+), 20 deletions(-) diff --git a/src/aiogram_dialog/api/entities/update_event.py b/src/aiogram_dialog/api/entities/update_event.py index 3c189858..f0e20bd5 100644 --- a/src/aiogram_dialog/api/entities/update_event.py +++ b/src/aiogram_dialog/api/entities/update_event.py @@ -14,6 +14,7 @@ ShowMode, StartMode, ) +from .stack import AccessSettings DIALOG_EVENT_NAME = "aiogd_update" @@ -42,6 +43,7 @@ class DialogUpdateEvent(TelegramObject): class DialogStartEvent(DialogUpdateEvent): new_state: State mode: StartMode + access_settings: Optional[AccessSettings] = None class DialogSwitchEvent(DialogUpdateEvent): diff --git a/src/aiogram_dialog/api/protocols/manager.py b/src/aiogram_dialog/api/protocols/manager.py index 824ebe52..ce164806 100644 --- a/src/aiogram_dialog/api/protocols/manager.py +++ b/src/aiogram_dialog/api/protocols/manager.py @@ -5,6 +5,7 @@ from aiogram.fsm.state import State from aiogram_dialog.api.entities import ( + AccessSettings, ChatEvent, Context, Data, ShowMode, Stack, StartMode, ) @@ -25,6 +26,7 @@ async def start( data: Data = None, mode: StartMode = StartMode.NORMAL, show_mode: Optional[ShowMode] = None, + access_settings: Optional[AccessSettings] = None, ) -> None: raise NotImplementedError diff --git a/src/aiogram_dialog/manager/bg_manager.py b/src/aiogram_dialog/manager/bg_manager.py index 20156ec0..d45e78e9 100644 --- a/src/aiogram_dialog/manager/bg_manager.py +++ b/src/aiogram_dialog/manager/bg_manager.py @@ -6,6 +6,7 @@ from aiogram.types import Chat, User from aiogram_dialog.api.entities import ( + AccessSettings, Data, DEFAULT_STACK_ID, DialogAction, @@ -130,6 +131,7 @@ async def start( data: Data = None, mode: StartMode = StartMode.NORMAL, show_mode: Optional[ShowMode] = None, + access_settings: Optional[AccessSettings] = None, ) -> None: await self._load() await self._notify( @@ -139,6 +141,7 @@ async def start( new_state=state, mode=mode, show_mode=show_mode, + access_settings=access_settings, **self._base_event_params(), ), ) diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index e18bed54..a2098a37 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -9,6 +9,7 @@ ) from aiogram_dialog.api.entities import ( + AccessSettings, ChatEvent, Context, Data, DEFAULT_STACK_ID, LaunchMode, MediaId, NewMessage, ShowMode, Stack, StartMode, ) @@ -198,16 +199,17 @@ async def start( data: Data = None, mode: StartMode = StartMode.NORMAL, show_mode: Optional[ShowMode] = None, + access_settings: Optional[AccessSettings] = None, ) -> None: self.check_disabled() self.show_mode = show_mode or self.show_mode if mode is StartMode.NORMAL: - await self._start_normal(state, data) + await self._start_normal(state, data, access_settings) elif mode is StartMode.RESET_STACK: await self.reset_stack(remove_keyboard=False) - await self._start_normal(state, data) + await self._start_normal(state, data, access_settings) elif mode is StartMode.NEW_STACK: - await self._start_new_stack(state, data) + await self._start_new_stack(state, data, access_settings) else: raise ValueError(f"Unknown start mode: {mode}") @@ -222,14 +224,25 @@ async def reset_stack(self, remove_keyboard: bool = True) -> None: await self._remove_kbd() self._data[CONTEXT_KEY] = None - async def _start_new_stack(self, state: State, data: Data = None) -> None: + async def _start_new_stack( + self, state: State, data: Data, + access_settings: Optional[AccessSettings], + ) -> None: stack = Stack() await self.bg(stack_id=stack.id).start( - state, data, StartMode.NORMAL, self.show_mode, + state, data, + mode=StartMode.NORMAL, + show_mode=self.show_mode, + access_settings=access_settings, ) - async def _start_normal(self, state: State, data: Data = None) -> None: + async def _start_normal( + self, state: State, data: Data, + access_settings: Optional[AccessSettings], + ) -> None: stack = self.current_stack() + if access_settings is not None: + stack.access_settings = access_settings old_dialog: Optional[DialogProtocol] = None if not stack.empty(): old_dialog = self.dialog() @@ -240,13 +253,7 @@ async def _start_normal(self, state: State, data: Data = None) -> None: ) new_dialog = self._registry.find_dialog(state) - launch_mode = new_dialog.launch_mode - if launch_mode in (LaunchMode.EXCLUSIVE, LaunchMode.ROOT): - await self.reset_stack(remove_keyboard=False) - if launch_mode is LaunchMode.SINGLE_TOP: - if new_dialog is old_dialog: - await self.storage().remove_context(stack.pop()) - + await self._process_launch_mode(old_dialog, new_dialog) if self.has_context(): await self.storage().save_context(self.current_context()) context = stack.push(state, data) @@ -256,6 +263,17 @@ async def _start_normal(self, state: State, data: Data = None) -> None: if new_context and context.id == new_context.id: await self.show() + async def _process_launch_mode( + self, + old_dialog: Optional[DialogProtocol], + new_dialog: DialogProtocol, + ): + if new_dialog.launch_mode in (LaunchMode.EXCLUSIVE, LaunchMode.ROOT): + await self.reset_stack(remove_keyboard=False) + if new_dialog.launch_mode is LaunchMode.SINGLE_TOP: + if new_dialog is old_dialog: + await self.storage().remove_context(self.current_stack().pop()) + async def next(self) -> None: context = self.current_context() states = self.dialog().states() diff --git a/src/aiogram_dialog/manager/sub_manager.py b/src/aiogram_dialog/manager/sub_manager.py index 592f2991..9e09054f 100644 --- a/src/aiogram_dialog/manager/sub_manager.py +++ b/src/aiogram_dialog/manager/sub_manager.py @@ -5,6 +5,7 @@ from aiogram.types import Message from aiogram_dialog.api.entities import ( + AccessSettings, ChatEvent, Data, ShowMode, StartMode, ) from aiogram_dialog.api.entities import Context, Stack @@ -107,11 +108,18 @@ async def done( async def mark_closed(self) -> None: await self.manager.mark_closed() - async def start(self, state: State, data: Data = None, - mode: StartMode = StartMode.NORMAL, - show_mode: Optional[ShowMode] = None) -> None: + async def start( + self, + state: State, + data: Data = None, + mode: StartMode = StartMode.NORMAL, + show_mode: Optional[ShowMode] = None, + access_settings: Optional[AccessSettings] = None, + ) -> None: await self.manager.start( - state=state, data=data, mode=mode, show_mode=show_mode, + state=state, data=data, + mode=mode, show_mode=show_mode, + access_settings=access_settings, ) async def switch_to( diff --git a/src/aiogram_dialog/manager/update_handler.py b/src/aiogram_dialog/manager/update_handler.py index 19ecd117..be4b2554 100644 --- a/src/aiogram_dialog/manager/update_handler.py +++ b/src/aiogram_dialog/manager/update_handler.py @@ -5,14 +5,16 @@ DialogStartEvent, DialogSwitchEvent, DialogUpdateEvent, + ShowMode, ) -from .manager import ManagerImpl -from .. import ShowMode +from aiogram_dialog.api.protocols import DialogManager logger = getLogger(__name__) -async def handle_update(event: DialogUpdateEvent, dialog_manager: ManagerImpl): +async def handle_update( + event: DialogUpdateEvent, dialog_manager: DialogManager, +) -> None: dialog_manager.show_mode = event.show_mode or ShowMode.AUTO if isinstance(event, DialogStartEvent): await dialog_manager.start( @@ -20,6 +22,7 @@ async def handle_update(event: DialogUpdateEvent, dialog_manager: ManagerImpl): data=event.data, mode=event.mode, show_mode=event.show_mode, + access_settings=event.access_settings, ) elif isinstance(event, DialogSwitchEvent): await dialog_manager.switch_to(state=event.new_state) diff --git a/src/aiogram_dialog/tools/preview.py b/src/aiogram_dialog/tools/preview.py index db22caae..ef9751b9 100644 --- a/src/aiogram_dialog/tools/preview.py +++ b/src/aiogram_dialog/tools/preview.py @@ -16,6 +16,7 @@ BaseDialogManager, Dialog, DialogManager, DialogProtocol, ) from aiogram_dialog.api.entities import ( + AccessSettings, ChatEvent, Context, Data, @@ -151,6 +152,7 @@ async def start( data: Data = None, mode: StartMode = StartMode.NORMAL, show_mode: ShowMode = ShowMode.AUTO, + access_settings: Optional[AccessSettings] = None, ) -> None: self.set_state(state) From 95ecf884d9f694f91e30bc348dc4f50c4c61b380 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 6 Dec 2023 23:31:14 +0100 Subject: [PATCH 10/44] Fix imports order --- tests/test_group.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_group.py b/tests/test_group.py index 1b4d88af..db08463a 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -13,7 +13,7 @@ from aiogram_dialog.test_tools import BotClient, MockMessageManager from aiogram_dialog.test_tools.keyboard import InlineButtonTextLocator from aiogram_dialog.widgets.kbd import Button -from aiogram_dialog.widgets.text import Format, Const +from aiogram_dialog.widgets.text import Const, Format class MainSG(StatesGroup): From c3cd24c510714dc3f7bee53ffd08e502d53193bf Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 6 Dec 2023 23:33:43 +0100 Subject: [PATCH 11/44] 2.2.0a1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 176343c7..86987486 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ where = ["src"] [project] name = "aiogram_dialog" -version = "2.1.0" +version = "2.2.0a1" readme = "README.md" authors = [ { name = "Andrey Tikhonov", email = "17@itishka.org" }, From a423292e7d1d19db6c4d491bf244b0cd1a881bb4 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 9 Dec 2023 23:14:21 +0100 Subject: [PATCH 12/44] fix loading access settings --- src/aiogram_dialog/context/storage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index e839a156..29c9ee9f 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -54,7 +54,7 @@ async def load_stack(self, stack_id: str = DEFAULT_STACK_ID) -> Stack: return Stack(_id=fixed_stack_id, access_settings=access_settings) access_settings = self._parse_access_settings( - data.pop("access_settings"), + data.pop("access_settings", None), ) return Stack(access_settings=access_settings, **data) From 0c5f87adb51555d2ecb56b5b423ff5f871106cec Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 9 Dec 2023 23:14:37 +0100 Subject: [PATCH 13/44] fix checking group for show mode --- src/aiogram_dialog/manager/manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index a2098a37..4929a460 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -423,7 +423,7 @@ def _save_last_message(self, message: OldMessage): def _calc_show_mode(self) -> ShowMode: if self.show_mode is not ShowMode.AUTO: return self.show_mode - if self.middleware_data["event_chat"].type == ChatType.GROUP: + if self.middleware_data["event_chat"].type != ChatType.PRIVATE: return ShowMode.EDIT if self.current_stack().last_reply_keyboard: return ShowMode.DELETE_AND_SEND From ace7637ebd600ac524bee29b4fa7e16a05499f61 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 9 Dec 2023 23:24:03 +0100 Subject: [PATCH 14/44] support thread id in message manager --- src/aiogram_dialog/api/entities/new_message.py | 1 + src/aiogram_dialog/manager/message_manager.py | 2 ++ src/aiogram_dialog/window.py | 1 + 3 files changed, 4 insertions(+) diff --git a/src/aiogram_dialog/api/entities/new_message.py b/src/aiogram_dialog/api/entities/new_message.py index 8fd4ffa2..52499591 100644 --- a/src/aiogram_dialog/api/entities/new_message.py +++ b/src/aiogram_dialog/api/entities/new_message.py @@ -31,6 +31,7 @@ class OldMessage: @dataclass class NewMessage: chat: Chat + thread_id: Optional[int] = None text: Optional[str] = None reply_markup: Optional[MarkupVariant] = None parse_mode: Optional[str] = None diff --git a/src/aiogram_dialog/manager/message_manager.py b/src/aiogram_dialog/manager/message_manager.py index c8a12bb3..2dc907ef 100644 --- a/src/aiogram_dialog/manager/message_manager.py +++ b/src/aiogram_dialog/manager/message_manager.py @@ -353,6 +353,7 @@ async def send_text(self, bot: Bot, new_message: NewMessage) -> Message: return await bot.send_message( new_message.chat.id, text=new_message.text, + message_thread_id=new_message.thread_id, disable_web_page_preview=new_message.disable_web_page_preview, reply_markup=new_message.reply_markup, parse_mode=new_message.parse_mode, @@ -372,6 +373,7 @@ async def send_media(self, bot: Bot, new_message: NewMessage) -> Message: return await method( new_message.chat.id, await self.get_media_source(new_message.media, bot), + message_thread_id=new_message.thread_id, caption=new_message.text, reply_markup=new_message.reply_markup, parse_mode=new_message.parse_mode, diff --git a/src/aiogram_dialog/window.py b/src/aiogram_dialog/window.py index 10fb2612..9735f76b 100644 --- a/src/aiogram_dialog/window.py +++ b/src/aiogram_dialog/window.py @@ -113,6 +113,7 @@ async def render( try: return NewMessage( chat=chat, + thread_id=manager.middleware_data.get("event_thread_id"), text=await self.render_text(current_data, manager), reply_markup=await self.render_kbd(current_data, manager), parse_mode=self.parse_mode, From a05930585c73e8a8e92d8673e1c6a1eeb1bfc78c Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 9 Dec 2023 23:24:13 +0100 Subject: [PATCH 15/44] cleanup window protocol --- src/aiogram_dialog/api/internal/window.py | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/src/aiogram_dialog/api/internal/window.py b/src/aiogram_dialog/api/internal/window.py index f5fbff9e..861abf99 100644 --- a/src/aiogram_dialog/api/internal/window.py +++ b/src/aiogram_dialog/api/internal/window.py @@ -1,7 +1,6 @@ from abc import abstractmethod from typing import ( Any, - Dict, Protocol, ) @@ -11,29 +10,9 @@ from aiogram_dialog.api.entities import NewMessage from aiogram_dialog.api.protocols import DialogProtocol from .manager import DialogManager -from .widgets import RawKeyboard class WindowProtocol(Protocol): - @abstractmethod - async def render_text(self, data: Dict, - manager: DialogManager) -> str: - raise NotImplementedError - - @abstractmethod - async def render_kbd( - self, data: Dict, manager: DialogManager, - ) -> RawKeyboard: - raise NotImplementedError - - @abstractmethod - async def load_data( - self, - dialog: "DialogProtocol", - manager: DialogManager, - ) -> Dict: - raise NotImplementedError - @abstractmethod async def process_message( self, From 09675a7795113a3bd0446565a9104e9c577b5f95 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 9 Dec 2023 23:38:37 +0100 Subject: [PATCH 16/44] add stack access validator --- .../context/access_validator.py | 24 +++++++++++++++++++ .../context/intent_middleware.py | 22 ++--------------- src/aiogram_dialog/setup.py | 23 +++++++++++++++--- 3 files changed, 46 insertions(+), 23 deletions(-) create mode 100644 src/aiogram_dialog/context/access_validator.py diff --git a/src/aiogram_dialog/context/access_validator.py b/src/aiogram_dialog/context/access_validator.py new file mode 100644 index 00000000..94e9596a --- /dev/null +++ b/src/aiogram_dialog/context/access_validator.py @@ -0,0 +1,24 @@ +from logging import getLogger + +from aiogram.enums import ChatType +from aiogram.types import Chat, User + +from aiogram_dialog.api.entities import ( + Stack, +) + +logger = getLogger(__name__) + + +class DefaultAccessValidator: + async def is_allowed( + self, stack: Stack, user: User, chat: Chat, + ) -> bool: + if not stack.access_settings: + return True + if chat.type is ChatType.PRIVATE: + return True + if stack.access_settings.user_ids: + if user.id not in stack.access_settings.user_ids: + return False + return True diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index e3bbdf32..8c68f005 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -3,7 +3,6 @@ from aiogram import Router from aiogram.dispatcher.middlewares.base import BaseMiddleware -from aiogram.enums import ChatType from aiogram.types import CallbackQuery, Chat, Message, User from aiogram.types.error_event import ErrorEvent @@ -26,32 +25,15 @@ logger = getLogger(__name__) -class AccessValidator: - async def is_allowed( - self, stack: Stack, user: User, chat: Chat, - ) -> bool: - if not stack.access_settings: - return True - if chat.type is ChatType.PRIVATE: - return True - if stack.access_settings.user_ids: - if user.id not in stack.access_settings.user_ids: - return False - return True - - class IntentMiddlewareFactory: def __init__( self, registry: DialogRegistryProtocol, - access_validator: Optional[StackAccessValidator] = None, + access_validator: StackAccessValidator, ): super().__init__() self.registry = registry - if access_validator: - self.access_validator = access_validator - else: - self.access_validator = AccessValidator() + self.access_validator = access_validator def storage_proxy(self, data: dict): proxy = StorageProxy( diff --git a/src/aiogram_dialog/setup.py b/src/aiogram_dialog/setup.py index c07a895f..5463692f 100644 --- a/src/aiogram_dialog/setup.py +++ b/src/aiogram_dialog/setup.py @@ -9,7 +9,7 @@ from aiogram_dialog.api.internal import DialogManagerFactory from aiogram_dialog.api.protocols import ( BgManagerFactory, DialogProtocol, DialogRegistryProtocol, - MediaIdStorageProtocol, MessageManagerProtocol, + MediaIdStorageProtocol, MessageManagerProtocol, StackAccessValidator, ) from aiogram_dialog.context.intent_middleware import ( context_saver_middleware, @@ -25,6 +25,7 @@ from aiogram_dialog.manager.message_manager import MessageManager from aiogram_dialog.manager.update_handler import handle_update from .about import about_dialog +from .context.access_validator import DefaultAccessValidator def _setup_event_observer(router: Router) -> None: @@ -86,6 +87,7 @@ def _register_middleware( router: Router, dialog_manager_factory: DialogManagerFactory, bg_manager_factory: BgManagerFactory, + stack_access_validator: StackAccessValidator, ): registry = DialogRegistry(router) manager_middleware = ManagerMiddleware( @@ -93,7 +95,9 @@ def _register_middleware( router=router, registry=registry, ) - intent_middleware = IntentMiddlewareFactory(registry=registry) + intent_middleware = IntentMiddlewareFactory( + registry=registry, access_validator=stack_access_validator, + ) # delayed configuration of middlewares router.startup.register(_startup_callback(registry)) update_handler = router.observers[DIALOG_EVENT_NAME] @@ -149,6 +153,15 @@ def _prepare_dialog_manager_factory( ) +def _prepare_stack_access_validator( + stack_access_validator: Optional[StackAccessValidator], +) -> StackAccessValidator: + if stack_access_validator: + return stack_access_validator + else: + return DefaultAccessValidator() + + def collect_dialogs(router: Router) -> Iterable[DialogProtocol]: if isinstance(router, DialogProtocol): yield router @@ -166,6 +179,7 @@ def setup_dialogs( dialog_manager_factory: Optional[DialogManagerFactory] = None, message_manager: Optional[MessageManagerProtocol] = None, media_id_storage: Optional[MediaIdStorageProtocol] = None, + stack_access_validator: Optional[StackAccessValidator] = None, ) -> BgManagerFactory: _setup_event_observer(router) _register_event_handler(router, handle_update) @@ -176,11 +190,14 @@ def setup_dialogs( message_manager=message_manager, media_id_storage=media_id_storage, ) + stack_access_validator = _prepare_stack_access_validator( + stack_access_validator, + ) bg_manager_factory = BgManagerFactoryImpl(router) _register_middleware( router=router, dialog_manager_factory=dialog_manager_factory, bg_manager_factory=bg_manager_factory, - + stack_access_validator=stack_access_validator, ) return bg_manager_factory From 6081d59d2b475f9dce1f0668be165fdf9bd9d791 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 10 Dec 2023 12:08:48 +0100 Subject: [PATCH 17/44] Events isolation --- .../context/intent_middleware.py | 24 +++++++++++++++---- src/aiogram_dialog/context/storage.py | 20 ++++++++++++++-- src/aiogram_dialog/setup.py | 24 +++++++++++++++++-- 3 files changed, 59 insertions(+), 9 deletions(-) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index 8c68f005..06dc4f66 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -3,6 +3,7 @@ from aiogram import Router from aiogram.dispatcher.middlewares.base import BaseMiddleware +from aiogram.fsm.storage.base import BaseEventIsolation from aiogram.types import CallbackQuery, Chat, Message, User from aiogram.types.error_event import ErrorEvent @@ -30,15 +31,18 @@ def __init__( self, registry: DialogRegistryProtocol, access_validator: StackAccessValidator, + events_isolation: BaseEventIsolation, ): super().__init__() self.registry = registry self.access_validator = access_validator + self.events_isolation = events_isolation def storage_proxy(self, data: dict): proxy = StorageProxy( bot=data["bot"], storage=data["fsm_storage"], + events_isolation=self.events_isolation, state_groups=self.registry.state_groups(), user_id=data["event_from_user"].id, chat_id=data["event_chat"].id, @@ -76,6 +80,7 @@ async def _load_stack( "Stack %s is not allowed for user %s", stack.id, user.id, ) + await proxy.unlock() return return stack @@ -257,11 +262,16 @@ async def process_callback_query( async def context_saver_middleware(handler, event, data): - result = await handler(event, data) - proxy: StorageProxy = data.pop(STORAGE_KEY, None) - if proxy: - await proxy.save_context(data.pop(CONTEXT_KEY)) - await proxy.save_stack(data.pop(STACK_KEY)) + proxy: StorageProxy = data.get(STORAGE_KEY, None) + try: + result = await handler(event, data) + proxy: StorageProxy = data.pop(STORAGE_KEY, None) + if proxy: + await proxy.save_context(data.pop(CONTEXT_KEY)) + await proxy.save_stack(data.pop(STACK_KEY)) + finally: + if proxy: + await proxy.unlock() return result @@ -269,9 +279,11 @@ class IntentErrorMiddleware(BaseMiddleware): def __init__( self, registry: DialogRegistryProtocol, + events_isolation: BaseEventIsolation, ): super().__init__() self.registry = registry + self.events_isolation = events_isolation def _is_error_supported( self, event: ErrorEvent, data: Dict[str, Any], @@ -326,6 +338,7 @@ async def __call__( proxy = StorageProxy( bot=data["bot"], storage=data["fsm_storage"], + events_isolation=self.events_isolation, user_id=user.id, chat_id=chat.id, thread_id=data.get("event_thread_id"), @@ -348,6 +361,7 @@ async def __call__( finally: proxy: StorageProxy = data.pop(STORAGE_KEY, None) if proxy: + await proxy.unlock() context = data.pop(CONTEXT_KEY) if context is not None: await proxy.save_context(context) diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 29c9ee9f..b3f543e3 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -1,10 +1,12 @@ +from contextlib import AsyncExitStack from copy import copy from typing import Dict, Optional, Type from aiogram import Bot from aiogram.enums import ChatMemberStatus from aiogram.fsm.state import State, StatesGroup -from aiogram.fsm.storage.base import BaseStorage, StorageKey +from aiogram.fsm.storage.base import BaseStorage, StorageKey, \ + BaseEventIsolation from aiogram_dialog.api.entities import ( AccessSettings, Context, DEFAULT_STACK_ID, Stack, @@ -12,10 +14,12 @@ from aiogram_dialog.api.exceptions import UnknownIntent, UnknownState + class StorageProxy: def __init__( self, storage: BaseStorage, + events_isolation: BaseEventIsolation, user_id: Optional[int], chat_id: int, thread_id: Optional[int], @@ -23,11 +27,21 @@ def __init__( state_groups: Dict[str, Type[StatesGroup]], ): self.storage = storage + self.events_isolation = events_isolation self.state_groups = state_groups self.user_id = user_id self.chat_id = chat_id self.thread_id = thread_id self.bot = bot + self.lock_stack = AsyncExitStack() + + async def lock(self, key: StorageKey): + await self.lock_stack.enter_async_context( + self.events_isolation.lock(key), + ) + + async def unlock(self): + await self.lock_stack.aclose() async def load_context(self, intent_id: str) -> Context: data = await self.storage.get_data( @@ -48,7 +62,9 @@ def _default_access_settings(self, stack_id: str) -> AccessSettings: async def load_stack(self, stack_id: str = DEFAULT_STACK_ID) -> Stack: fixed_stack_id = self._fixed_stack_id(stack_id) - data = await self.storage.get_data(self._stack_key(fixed_stack_id)) + key = self._stack_key(fixed_stack_id) + await self.lock(key) + data = await self.storage.get_data(key) if not data: access_settings = self._default_access_settings(stack_id) return Stack(_id=fixed_stack_id, access_settings=access_settings) diff --git a/src/aiogram_dialog/setup.py b/src/aiogram_dialog/setup.py index 5463692f..548b20d5 100644 --- a/src/aiogram_dialog/setup.py +++ b/src/aiogram_dialog/setup.py @@ -3,6 +3,8 @@ from aiogram import Router from aiogram.dispatcher.event.telegram import TelegramEventObserver from aiogram.fsm.state import any_state, State, StatesGroup +from aiogram.fsm.storage.base import BaseEventIsolation +from aiogram.fsm.storage.memory import SimpleEventIsolation from aiogram_dialog.api.entities import DIALOG_EVENT_NAME from aiogram_dialog.api.exceptions import UnregisteredDialogError @@ -88,6 +90,7 @@ def _register_middleware( dialog_manager_factory: DialogManagerFactory, bg_manager_factory: BgManagerFactory, stack_access_validator: StackAccessValidator, + events_isolation: BaseEventIsolation, ): registry = DialogRegistry(router) manager_middleware = ManagerMiddleware( @@ -96,13 +99,18 @@ def _register_middleware( registry=registry, ) intent_middleware = IntentMiddlewareFactory( - registry=registry, access_validator=stack_access_validator, + registry=registry, + access_validator=stack_access_validator, + events_isolation=events_isolation, ) # delayed configuration of middlewares router.startup.register(_startup_callback(registry)) update_handler = router.observers[DIALOG_EVENT_NAME] - router.errors.middleware(IntentErrorMiddleware(registry=registry)) + router.errors.middleware(IntentErrorMiddleware( + registry=registry, + events_isolation=events_isolation, + )) router.message.middleware(manager_middleware) router.callback_query.middleware(manager_middleware) @@ -162,6 +170,15 @@ def _prepare_stack_access_validator( return DefaultAccessValidator() +def _prepare_events_isolation( + events_isolation: Optional[BaseEventIsolation], +) -> BaseEventIsolation: + if events_isolation: + return events_isolation + else: + return SimpleEventIsolation() + + def collect_dialogs(router: Router) -> Iterable[DialogProtocol]: if isinstance(router, DialogProtocol): yield router @@ -180,6 +197,7 @@ def setup_dialogs( message_manager: Optional[MessageManagerProtocol] = None, media_id_storage: Optional[MediaIdStorageProtocol] = None, stack_access_validator: Optional[StackAccessValidator] = None, + events_isolation: Optional[BaseEventIsolation] = None, ) -> BgManagerFactory: _setup_event_observer(router) _register_event_handler(router, handle_update) @@ -193,11 +211,13 @@ def setup_dialogs( stack_access_validator = _prepare_stack_access_validator( stack_access_validator, ) + events_isolation = _prepare_events_isolation(events_isolation) bg_manager_factory = BgManagerFactoryImpl(router) _register_middleware( router=router, dialog_manager_factory=dialog_manager_factory, bg_manager_factory=bg_manager_factory, stack_access_validator=stack_access_validator, + events_isolation=events_isolation ) return bg_manager_factory From 621f1a55b8f12eeddce46c07e37fd51d0e6d6784 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 10 Dec 2023 12:13:03 +0100 Subject: [PATCH 18/44] 2.2.0a2 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 86987486..58c20f78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ where = ["src"] [project] name = "aiogram_dialog" -version = "2.2.0a1" +version = "2.2.0a2" readme = "README.md" authors = [ { name = "Andrey Tikhonov", email = "17@itishka.org" }, From d7c396a1f66ef512e39e2622f6775d68987e07f0 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 30 Dec 2023 01:16:48 +0100 Subject: [PATCH 19/44] change access validator interface --- .../api/protocols/stack_access.py | 5 ++--- src/aiogram_dialog/context/access_validator.py | 9 ++++++--- .../context/intent_middleware.py | 18 +++++++++++++----- src/aiogram_dialog/context/storage.py | 6 +++--- src/aiogram_dialog/setup.py | 2 +- 5 files changed, 25 insertions(+), 15 deletions(-) diff --git a/src/aiogram_dialog/api/protocols/stack_access.py b/src/aiogram_dialog/api/protocols/stack_access.py index 5f156907..30091553 100644 --- a/src/aiogram_dialog/api/protocols/stack_access.py +++ b/src/aiogram_dialog/api/protocols/stack_access.py @@ -1,14 +1,13 @@ from abc import abstractmethod from typing import Protocol -from aiogram.types import Chat, User - +from aiogram_dialog import ChatEvent from aiogram_dialog.api.entities import Stack class StackAccessValidator(Protocol): @abstractmethod async def is_allowed( - self, stack: Stack, user: User, chat: Chat, + self, stack: Stack, event: ChatEvent, data: dict, ) -> bool: raise NotImplementedError diff --git a/src/aiogram_dialog/context/access_validator.py b/src/aiogram_dialog/context/access_validator.py index 94e9596a..f6a33e51 100644 --- a/src/aiogram_dialog/context/access_validator.py +++ b/src/aiogram_dialog/context/access_validator.py @@ -1,24 +1,27 @@ from logging import getLogger from aiogram.enums import ChatType -from aiogram.types import Chat, User +from aiogram_dialog import ChatEvent from aiogram_dialog.api.entities import ( Stack, ) +from aiogram_dialog.api.protocols import StackAccessValidator logger = getLogger(__name__) -class DefaultAccessValidator: +class DefaultAccessValidator(StackAccessValidator): async def is_allowed( - self, stack: Stack, user: User, chat: Chat, + self, stack: Stack, event: ChatEvent, data: dict, ) -> bool: if not stack.access_settings: return True + chat = data["event_chat"] if chat.type is ChatType.PRIVATE: return True if stack.access_settings.user_ids: + user = data["event_from_user"] if user.id not in stack.access_settings.user_ids: return False return True diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index 06dc4f66..3b2f7e78 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -67,15 +67,16 @@ def _check_outdated(self, intent_id: str, stack: Stack): async def _load_stack( self, + event: ChatEvent, stack_id: Optional[str], proxy: StorageProxy, - user: User, - chat: Chat, + data: dict, ) -> Optional[Stack]: if stack_id is None: raise InvalidStackIdError("Both stack id and intent id are None") stack = await proxy.load_stack(stack_id) - if not await self.access_validator.is_allowed(stack, user, chat): + if not await self.access_validator.is_allowed(stack, event, data): + user = data["event_from_user"] logger.debug( "Stack %s is not allowed for user %s", stack.id, user.id, @@ -86,6 +87,7 @@ async def _load_stack( async def _load_context_by_stack( self, + event: ChatEvent, proxy: StorageProxy, stack_id: Optional[str], data: dict, @@ -96,7 +98,7 @@ async def _load_context_by_stack( "Loading context for stack: `%s`, user: `%s`, chat: `%s`", stack_id, user.id, chat.id, ) - stack = await self._load_stack(stack_id, proxy, user, chat) + stack = await self._load_stack(event, stack_id, proxy, data) if not stack: return if stack.empty(): @@ -109,6 +111,7 @@ async def _load_context_by_stack( async def _load_context_by_intent( self, + event: ChatEvent, proxy: StorageProxy, intent_id: Optional[str], data: dict, @@ -120,7 +123,7 @@ async def _load_context_by_intent( intent_id, user.id, chat.id, ) context = await proxy.load_context(intent_id) - stack = await self._load_stack(context.stack_id, proxy, user, chat) + stack = await self._load_stack(event, context.stack_id, proxy, data) if not stack: return self._check_outdated(intent_id, stack) @@ -133,6 +136,7 @@ async def _load_default_context( self, event: ChatEvent, data: dict, ) -> None: return await self._load_context_by_stack( + event=event, proxy=self.storage_proxy(data), stack_id=DEFAULT_STACK_ID, data=data, @@ -182,6 +186,7 @@ async def process_message( if intent_id := self._intent_id_from_reply(event, data): await self._load_context_by_intent( + event=event, proxy=self.storage_proxy(data), intent_id=intent_id, data=data, @@ -216,12 +221,14 @@ async def process_aiogd_update( ): if event.intent_id: await self._load_context_by_intent( + event=event, proxy=self.storage_proxy(data), intent_id=event.intent_id, data=data, ) else: await self._load_context_by_stack( + event=event, proxy=self.storage_proxy(data), stack_id=event.stack_id, data=data, @@ -241,6 +248,7 @@ async def process_callback_query( intent_id, callback_data = remove_indent_id(event.data) if intent_id: await self._load_context_by_intent( + event=event, proxy=self.storage_proxy(data), intent_id=intent_id, data=data, diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index b3f543e3..5884691b 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -5,8 +5,9 @@ from aiogram import Bot from aiogram.enums import ChatMemberStatus from aiogram.fsm.state import State, StatesGroup -from aiogram.fsm.storage.base import BaseStorage, StorageKey, \ - BaseEventIsolation +from aiogram.fsm.storage.base import ( + BaseEventIsolation, BaseStorage, StorageKey, +) from aiogram_dialog.api.entities import ( AccessSettings, Context, DEFAULT_STACK_ID, Stack, @@ -14,7 +15,6 @@ from aiogram_dialog.api.exceptions import UnknownIntent, UnknownState - class StorageProxy: def __init__( self, diff --git a/src/aiogram_dialog/setup.py b/src/aiogram_dialog/setup.py index 548b20d5..8c119e22 100644 --- a/src/aiogram_dialog/setup.py +++ b/src/aiogram_dialog/setup.py @@ -218,6 +218,6 @@ def setup_dialogs( dialog_manager_factory=dialog_manager_factory, bg_manager_factory=bg_manager_factory, stack_access_validator=stack_access_validator, - events_isolation=events_isolation + events_isolation=events_isolation, ) return bg_manager_factory From ca91d4caec3be18f9ddc9093b937384cb35929c9 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 30 Dec 2023 12:53:19 +0100 Subject: [PATCH 20/44] Test event isolation --- requirements_dev.txt | 1 + tests/test_isolation.py | 46 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) create mode 100644 tests/test_isolation.py diff --git a/requirements_dev.txt b/requirements_dev.txt index 3437218f..052269fd 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -15,3 +15,4 @@ flake8-print pytest pytest-asyncio +pytest-repeat diff --git a/tests/test_isolation.py b/tests/test_isolation.py new file mode 100644 index 00000000..1e8a5d4f --- /dev/null +++ b/tests/test_isolation.py @@ -0,0 +1,46 @@ +import asyncio +from asyncio import Event + +import pytest +from aiogram import Dispatcher +from aiogram.filters import CommandStart +from aiogram.fsm.storage.memory import MemoryStorage +from aiogram.types import Message + +from aiogram_dialog import ( + setup_dialogs, ) +from aiogram_dialog.test_tools import BotClient, MockMessageManager + + +async def start( + message: Message, data: list, event_common: Event, +): + data.append(1) + await event_common.wait() + + +@pytest.mark.asyncio +@pytest.mark.repeat(10) +async def test_concurrent_events(): + event_common = Event() + data = [] + dp = Dispatcher( + event_common=event_common, + data=data, + storage=MemoryStorage(), + ) + dp.message.register(start, CommandStart()) + + client = BotClient(dp) + message_manager = MockMessageManager() + setup_dialogs(dp, message_manager=message_manager) + + # start + t1 = asyncio.create_task(client.send("/start")) + t2 = asyncio.create_task(client.send("/start")) + await asyncio.sleep(0.1) + assert len(data) == 1 # "Only single event expected to be processed" + event_common.set() + await t1 + await t2 + assert len(data) == 2 From 9320f9534bfa492ea0e1830f794fc5011692cf13 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 30 Dec 2023 13:01:35 +0100 Subject: [PATCH 21/44] 2.2.0a3 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 58c20f78..f253c033 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ where = ["src"] [project] name = "aiogram_dialog" -version = "2.2.0a2" +version = "2.2.0a3" readme = "README.md" authors = [ { name = "Andrey Tikhonov", email = "17@itishka.org" }, From e4be20864c9dc80ff26d7ba495afd8046c43b31b Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Mon, 3 Jun 2024 18:22:48 +0200 Subject: [PATCH 22/44] 2.2.0a4 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d8eb1cce..16926fc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ where = ["src"] [project] name = "aiogram_dialog" -version = "2.2.0a3" +version = "2.2.0a4" readme = "README.md" authors = [ { name = "Andrey Tikhonov", email = "17@itishka.org" }, From 95da31ba151b4409e78b26698d9e9afa7f5d5c5b Mon Sep 17 00:00:00 2001 From: KuroAngel <145038102+KuroKoka551@users.noreply.github.com> Date: Fri, 7 Jun 2024 00:09:04 +0500 Subject: [PATCH 23/44] Fix not released lock --- .../context/intent_middleware.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index 220ca5cc..bb9852f7 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -74,16 +74,18 @@ async def _load_stack( ) -> Optional[Stack]: if stack_id is None: raise InvalidStackIdError("Both stack id and intent id are None") - stack = await proxy.load_stack(stack_id) - if not await self.access_validator.is_allowed(stack, event, data): - user = data["event_from_user"] - logger.debug( - "Stack %s is not allowed for user %s", - stack.id, user.id, - ) + try: + stack = await proxy.load_stack(stack_id) + if not await self.access_validator.is_allowed(stack, event, data): + user = data["event_from_user"] + logger.debug( + "Stack %s is not allowed for user %s", + stack.id, user.id, + ) + return + return stack + finally: await proxy.unlock() - return - return stack async def _load_context_by_stack( self, From 9e799031c0beea2dfc11970b8274caf73b2ed2f1 Mon Sep 17 00:00:00 2001 From: KuroAngel <145038102+KuroKoka551@users.noreply.github.com> Date: Fri, 7 Jun 2024 21:23:42 +0500 Subject: [PATCH 24/44] release lock on Outdated error --- .../context/intent_middleware.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index bb9852f7..b16bac49 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -74,18 +74,16 @@ async def _load_stack( ) -> Optional[Stack]: if stack_id is None: raise InvalidStackIdError("Both stack id and intent id are None") - try: - stack = await proxy.load_stack(stack_id) - if not await self.access_validator.is_allowed(stack, event, data): - user = data["event_from_user"] - logger.debug( - "Stack %s is not allowed for user %s", - stack.id, user.id, - ) - return - return stack - finally: + stack = await proxy.load_stack(stack_id) + if not await self.access_validator.is_allowed(stack, event, data): + user = data["event_from_user"] + logger.debug( + "Stack %s is not allowed for user %s", + stack.id, user.id, + ) await proxy.unlock() + return + return stack async def _load_context_by_stack( self, @@ -128,7 +126,11 @@ async def _load_context_by_intent( stack = await self._load_stack(event, context.stack_id, proxy, data) if not stack: return - self._check_outdated(intent_id, stack) + try: + self._check_outdated(intent_id, stack) + except OutdatedIntent: + await proxy.unlock() + raise data[STORAGE_KEY] = proxy data[STACK_KEY] = stack From 9dfbd8ba0e3765d41abf03f90239c07c6e80c52f Mon Sep 17 00:00:00 2001 From: KuroAngel <145038102+KuroKoka551@users.noreply.github.com> Date: Fri, 7 Jun 2024 21:31:23 +0500 Subject: [PATCH 25/44] More exception caught --- src/aiogram_dialog/context/intent_middleware.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index b16bac49..df3d77c2 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -104,7 +104,11 @@ async def _load_context_by_stack( if stack.empty(): context = None else: - context = await proxy.load_context(stack.last_intent_id()) + try: + context = await proxy.load_context(stack.last_intent_id()) + except: + await proxy.unlock() + raise data[STORAGE_KEY] = proxy data[STACK_KEY] = stack data[CONTEXT_KEY] = context @@ -128,7 +132,7 @@ async def _load_context_by_intent( return try: self._check_outdated(intent_id, stack) - except OutdatedIntent: + except: await proxy.unlock() raise From 66d1a260a0d3e1c9c83bec5adbabd8ac9e5f3eef Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 8 Jun 2024 22:43:41 +0200 Subject: [PATCH 26/44] add thread id to bg manager --- src/aiogram_dialog/api/entities/update_event.py | 1 + src/aiogram_dialog/api/protocols/manager.py | 2 ++ src/aiogram_dialog/manager/bg_manager.py | 10 ++++++++++ src/aiogram_dialog/manager/manager.py | 7 ++++++- src/aiogram_dialog/manager/message_manager.py | 2 +- src/aiogram_dialog/manager/sub_manager.py | 17 +++++++++++++---- src/aiogram_dialog/manager/updater.py | 1 + 7 files changed, 34 insertions(+), 6 deletions(-) diff --git a/src/aiogram_dialog/api/entities/update_event.py b/src/aiogram_dialog/api/entities/update_event.py index f0e20bd5..9b8d8dd6 100644 --- a/src/aiogram_dialog/api/entities/update_event.py +++ b/src/aiogram_dialog/api/entities/update_event.py @@ -37,6 +37,7 @@ class DialogUpdateEvent(TelegramObject): data: Any intent_id: Optional[str] stack_id: Optional[str] + thread_id: Optional[int] show_mode: Optional[ShowMode] = None diff --git a/src/aiogram_dialog/api/protocols/manager.py b/src/aiogram_dialog/api/protocols/manager.py index f62548db..41a59e23 100644 --- a/src/aiogram_dialog/api/protocols/manager.py +++ b/src/aiogram_dialog/api/protocols/manager.py @@ -52,6 +52,7 @@ def bg( user_id: Optional[int] = None, chat_id: Optional[int] = None, stack_id: Optional[str] = None, + thread_id: Optional[int] = None, load: bool = False, # load chat and user ) -> "BaseDialogManager": raise NotImplementedError @@ -65,6 +66,7 @@ def bg( user_id: int, chat_id: int, stack_id: Optional[str] = None, + thread_id: Optional[int] = None, load: bool = False, # load chat and user ) -> "BaseDialogManager": raise NotImplementedError diff --git a/src/aiogram_dialog/manager/bg_manager.py b/src/aiogram_dialog/manager/bg_manager.py index d45e78e9..d1955f87 100644 --- a/src/aiogram_dialog/manager/bg_manager.py +++ b/src/aiogram_dialog/manager/bg_manager.py @@ -36,6 +36,7 @@ def __init__( router: Router, intent_id: Optional[str], stack_id: Optional[str], + thread_id: Optional[int], load: bool = False, ): self.user = user @@ -45,6 +46,7 @@ def __init__( self._updater = Updater(router) self.intent_id = intent_id self.stack_id = stack_id + self.thread_id = thread_id self.load = load def bg( @@ -52,6 +54,7 @@ def bg( user_id: Optional[int] = None, chat_id: Optional[int] = None, stack_id: Optional[str] = None, + thread_id: Optional[int] = None, load: bool = False, ) -> "BaseDialogManager": if chat_id in (None, self.chat.id): @@ -75,6 +78,9 @@ def bg( else: intent_id = None + if thread_id is None and same_chat: + thread_id = self.thread_id + return BgManager( user=user, chat=chat, @@ -82,6 +88,7 @@ def bg( router=self._router, intent_id=intent_id, stack_id=stack_id, + thread_id=thread_id, load=load, ) @@ -91,6 +98,7 @@ def _base_event_params(self): "chat": self.chat, "intent_id": self.intent_id, "stack_id": self.stack_id, + "thread_id": self.thread_id, } async def _notify(self, event: DialogUpdateEvent): @@ -186,6 +194,7 @@ def bg( user_id: int, chat_id: int, stack_id: Optional[str] = None, + thread_id: Optional[int] = None, load: bool = False, ) -> "BaseDialogManager": chat = FakeChat(id=chat_id, type="") @@ -200,5 +209,6 @@ def bg( router=self._router, intent_id=None, stack_id=stack_id, + thread_id=thread_id, load=load, ) diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index 73468344..ff80960a 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -487,18 +487,22 @@ def bg( user_id: Optional[int] = None, chat_id: Optional[int] = None, stack_id: Optional[str] = None, + thread_id: Optional[int] = None, load: bool = False, ) -> BaseDialogManager: user = self._get_fake_user(user_id) chat = self._get_fake_chat(chat_id) intent_id = None + same_chat = self.is_same_chat(user, chat) if stack_id is None: - if self.is_same_chat(user, chat): + if same_chat: stack_id = self.current_stack().id if self.has_context(): intent_id = self.current_context().id else: stack_id = DEFAULT_STACK_ID + if thread_id is None and same_chat: + thread_id = self.middleware_data.get("event_thread_id") return BgManager( user=user, @@ -507,6 +511,7 @@ def bg( router=self._router, intent_id=intent_id, stack_id=stack_id, + thread_id=thread_id, load=load, ) diff --git a/src/aiogram_dialog/manager/message_manager.py b/src/aiogram_dialog/manager/message_manager.py index 2dc907ef..ba2e5a7c 100644 --- a/src/aiogram_dialog/manager/message_manager.py +++ b/src/aiogram_dialog/manager/message_manager.py @@ -349,7 +349,7 @@ async def send_message(self, bot: Bot, new_message: NewMessage) -> Message: return await self.send_text(bot, new_message) async def send_text(self, bot: Bot, new_message: NewMessage) -> Message: - logger.debug("send_text to %s", new_message.chat) + logger.debug("send_text to chat %s, thread %s", new_message.chat.id, new_message.thread_id) return await bot.send_message( new_message.chat.id, text=new_message.text, diff --git a/src/aiogram_dialog/manager/sub_manager.py b/src/aiogram_dialog/manager/sub_manager.py index 9bd534dd..04dfa75b 100644 --- a/src/aiogram_dialog/manager/sub_manager.py +++ b/src/aiogram_dialog/manager/sub_manager.py @@ -137,9 +137,18 @@ async def update( self.current_context().dialog_data.update(data) await self.show(show_mode) - def bg(self, user_id: Optional[int] = None, chat_id: Optional[int] = None, - stack_id: Optional[str] = None, - load: bool = False) -> BaseDialogManager: + def bg( + self, + user_id: Optional[int] = None, + chat_id: Optional[int] = None, + stack_id: Optional[str] = None, + thread_id: Optional[int] = None, + load: bool = False, + ) -> BaseDialogManager: return self.manager.bg( - user_id=user_id, chat_id=chat_id, stack_id=stack_id, load=load, + user_id=user_id, + chat_id=chat_id, + stack_id=stack_id, + thread_id=thread_id, + load=load, ) diff --git a/src/aiogram_dialog/manager/updater.py b/src/aiogram_dialog/manager/updater.py index b85a340b..df8926f7 100644 --- a/src/aiogram_dialog/manager/updater.py +++ b/src/aiogram_dialog/manager/updater.py @@ -30,5 +30,6 @@ async def _process_update(self, bot: Bot, update: DialogUpdate) -> None: bot=bot, event_from_user=event.from_user, event_chat=event.chat, + event_thread_id=event.thread_id, **self.dp.workflow_data, ) From ddc430ac2190591df200b531e99b7b98b79ec6d0 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sat, 8 Jun 2024 22:58:45 +0200 Subject: [PATCH 27/44] context unlocker middleware, fix tools --- src/aiogram_dialog/context/intent_middleware.py | 13 +++++++++---- src/aiogram_dialog/setup.py | 7 +++++++ src/aiogram_dialog/tools/preview.py | 1 + 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index df3d77c2..0864cede 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -278,13 +278,18 @@ async def process_callback_query( async def context_saver_middleware(handler, event, data): + result = await handler(event, data) + proxy: StorageProxy = data.pop(STORAGE_KEY, None) + if proxy: + await proxy.save_context(data.pop(CONTEXT_KEY)) + await proxy.save_stack(data.pop(STACK_KEY)) + return result + + +async def context_unlocker_middleware(handler, event, data): proxy: StorageProxy = data.get(STORAGE_KEY, None) try: result = await handler(event, data) - proxy: StorageProxy = data.pop(STORAGE_KEY, None) - if proxy: - await proxy.save_context(data.pop(CONTEXT_KEY)) - await proxy.save_stack(data.pop(STACK_KEY)) finally: if proxy: await proxy.unlock() diff --git a/src/aiogram_dialog/setup.py b/src/aiogram_dialog/setup.py index aaab5eba..e9761ac2 100644 --- a/src/aiogram_dialog/setup.py +++ b/src/aiogram_dialog/setup.py @@ -15,6 +15,7 @@ ) from aiogram_dialog.context.intent_middleware import ( context_saver_middleware, + context_unlocker_middleware, IntentErrorMiddleware, IntentMiddlewareFactory, ) @@ -133,6 +134,12 @@ def _register_middleware( intent_middleware.process_chat_join_request, ) + router.message.outer_middleware(context_unlocker_middleware) + router.callback_query.outer_middleware(context_unlocker_middleware) + update_handler.outer_middleware(context_unlocker_middleware) + router.my_chat_member.outer_middleware(context_unlocker_middleware) + router.chat_join_request.outer_middleware(context_unlocker_middleware) + router.message.middleware(context_saver_middleware) router.callback_query.middleware(context_saver_middleware) update_handler.middleware(context_saver_middleware) diff --git a/src/aiogram_dialog/tools/preview.py b/src/aiogram_dialog/tools/preview.py index 2a436548..5605127f 100644 --- a/src/aiogram_dialog/tools/preview.py +++ b/src/aiogram_dialog/tools/preview.py @@ -70,6 +70,7 @@ def __init__(self): data={}, intent_id=None, stack_id=None, + thread_id=None, ) self._context: Optional[Context] = None self._dialog: Optional[DialogProtocol] = None From 58f54804496edd3fdfbc24c9179b663bdb174220 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Fri, 21 Jun 2024 23:01:40 +0200 Subject: [PATCH 28/44] business id --- pyproject.toml | 2 +- src/aiogram_dialog/__init__.py | 2 + .../api/entities/update_event.py | 1 + src/aiogram_dialog/api/protocols/__init__.py | 6 ++- src/aiogram_dialog/api/protocols/manager.py | 11 +++++- .../context/intent_middleware.py | 35 ++++++++++-------- src/aiogram_dialog/context/storage.py | 4 ++ src/aiogram_dialog/manager/bg_manager.py | 37 +++++++++++++++---- src/aiogram_dialog/manager/manager.py | 25 +++++++++++-- src/aiogram_dialog/manager/sub_manager.py | 10 +++-- src/aiogram_dialog/tools/preview.py | 11 ++++-- 11 files changed, 107 insertions(+), 37 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 16926fc7..3bf4b7a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - 'aiogram>=3.2.0', + 'aiogram>=3.5.0', 'jinja2', 'cachetools>=4.0.0,<6.0.0', 'magic_filter', diff --git a/src/aiogram_dialog/__init__.py b/src/aiogram_dialog/__init__.py index f7f27faf..54fb11cb 100644 --- a/src/aiogram_dialog/__init__.py +++ b/src/aiogram_dialog/__init__.py @@ -9,6 +9,7 @@ "BgManagerFactory", "DialogManager", "DialogProtocol", + "UnsetId", "setup_dialogs", "ShowMode", "SubManager", @@ -22,6 +23,7 @@ ) from .api.protocols import ( BaseDialogManager, BgManagerFactory, DialogManager, DialogProtocol, + UnsetId, ) from .dialog import Dialog from .manager.sub_manager import SubManager diff --git a/src/aiogram_dialog/api/entities/update_event.py b/src/aiogram_dialog/api/entities/update_event.py index 9b8d8dd6..17315d09 100644 --- a/src/aiogram_dialog/api/entities/update_event.py +++ b/src/aiogram_dialog/api/entities/update_event.py @@ -38,6 +38,7 @@ class DialogUpdateEvent(TelegramObject): intent_id: Optional[str] stack_id: Optional[str] thread_id: Optional[int] + business_connection_id: Optional[str] show_mode: Optional[ShowMode] = None diff --git a/src/aiogram_dialog/api/protocols/__init__.py b/src/aiogram_dialog/api/protocols/__init__.py index 21610eaf..b83cef41 100644 --- a/src/aiogram_dialog/api/protocols/__init__.py +++ b/src/aiogram_dialog/api/protocols/__init__.py @@ -1,6 +1,6 @@ __all__ = [ "DialogProtocol", - "BaseDialogManager", "BgManagerFactory", "DialogManager", + "BaseDialogManager", "BgManagerFactory", "DialogManager", "UnsetId", "MediaIdStorageProtocol", "MessageManagerProtocol", "MessageNotModified", "DialogProtocol", "DialogRegistryProtocol", @@ -8,7 +8,9 @@ ] from .dialog import DialogProtocol -from .manager import BaseDialogManager, BgManagerFactory, DialogManager +from .manager import ( + BaseDialogManager, BgManagerFactory, DialogManager, UnsetId, +) from .media import MediaIdStorageProtocol from .message_manager import MessageManagerProtocol, MessageNotModified from .registry import DialogRegistryProtocol diff --git a/src/aiogram_dialog/api/protocols/manager.py b/src/aiogram_dialog/api/protocols/manager.py index 41a59e23..686513bd 100644 --- a/src/aiogram_dialog/api/protocols/manager.py +++ b/src/aiogram_dialog/api/protocols/manager.py @@ -1,5 +1,6 @@ from abc import abstractmethod -from typing import Any, Dict, Optional, Protocol +from enum import Enum +from typing import Any, Dict, Optional, Protocol, Union from aiogram import Bot from aiogram.fsm.state import State @@ -10,6 +11,10 @@ ) +class UnsetId(Enum): + UNSET = "UNSET" + + class BaseDialogManager(Protocol): @abstractmethod async def done( @@ -52,7 +57,8 @@ def bg( user_id: Optional[int] = None, chat_id: Optional[int] = None, stack_id: Optional[str] = None, - thread_id: Optional[int] = None, + thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, # load chat and user ) -> "BaseDialogManager": raise NotImplementedError @@ -67,6 +73,7 @@ def bg( chat_id: int, stack_id: Optional[str] = None, thread_id: Optional[int] = None, + business_connection_id: Optional[str] = None, load: bool = False, # load chat and user ) -> "BaseDialogManager": raise NotImplementedError diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index 0864cede..d618ff6a 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -1,11 +1,14 @@ from logging import getLogger -from typing import Any, Awaitable, Callable, Dict, Optional +from typing import Any, Awaitable, Callable, Dict, Optional, cast from aiogram import Router from aiogram.dispatcher.middlewares.base import BaseMiddleware from aiogram.fsm.storage.base import BaseEventIsolation from aiogram.types import CallbackQuery, Chat, Message, User from aiogram.types.error_event import ErrorEvent +from aiogram.dispatcher.middlewares.user_context import ( + EVENT_CONTEXT_KEY, EventContext +) from aiogram_dialog.api.entities import ( ChatEvent, Context, DEFAULT_STACK_ID, DialogUpdateEvent, Stack, @@ -39,14 +42,16 @@ def __init__( self.events_isolation = events_isolation def storage_proxy(self, data: dict): + event_context = cast(EventContext, data.get(EVENT_CONTEXT_KEY)) proxy = StorageProxy( bot=data["bot"], storage=data["fsm_storage"], events_isolation=self.events_isolation, state_groups=self.registry.states_groups(), - user_id=data["event_from_user"].id, - chat_id=data["event_chat"].id, - thread_id=data.get("event_thread_id"), + user_id=event_context.user_id, + chat_id=event_context.chat_id, + thread_id=event_context.thread_id, + business_connection_id=event_context.business_connection_id, ) return proxy @@ -94,9 +99,10 @@ async def _load_context_by_stack( ) -> None: user = data["event_from_user"] chat = data["event_chat"] + thread_id = data.get("event_thread_id") logger.debug( - "Loading context for stack: `%s`, user: `%s`, chat: `%s`", - stack_id, user.id, chat.id, + "Loading context for stack: `%s`, user: `%s`, chat: `%s`, thread: `%s`", + stack_id, user.id, chat.id, thread_id, ) stack = await self._load_stack(event, stack_id, proxy, data) if not stack: @@ -328,14 +334,13 @@ async def _fix_broken_stack( async def _load_last_context( self, storage: StorageProxy, stack: Stack, - chat: Chat, user: User, ) -> Optional[Context]: try: return await storage.load_context(stack.last_intent_id()) except (UnknownIntent, OutdatedIntent): logger.warning( "Stack is broken for user %s, chat %s, resetting", - user.id, chat.id, + storage.user_id, storage.chat_id, ) await self._fix_broken_stack(storage, stack) return None @@ -353,17 +358,16 @@ async def __call__( return await handler(event, data) try: - chat = data["event_chat"] - user = data["event_from_user"] - + event_context = cast(EventContext, data.get(EVENT_CONTEXT_KEY)) proxy = StorageProxy( bot=data["bot"], storage=data["fsm_storage"], events_isolation=self.events_isolation, - user_id=user.id, - chat_id=chat.id, - thread_id=data.get("event_thread_id"), state_groups=self.registry.states_groups(), + user_id=event_context.user_id, + chat_id=event_context.chat_id, + thread_id=event_context.thread_id, + business_connection_id=event_context.business_connection_id, ) data[STORAGE_KEY] = proxy if isinstance(error, OutdatedIntent): @@ -374,7 +378,8 @@ async def __call__( context = None else: context = await self._load_last_context( - storage=proxy, stack=stack, chat=chat, user=user, + storage=proxy, + stack=stack, ) data[STACK_KEY] = stack data[CONTEXT_KEY] = context diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 5884691b..75d57d83 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -23,6 +23,7 @@ def __init__( user_id: Optional[int], chat_id: int, thread_id: Optional[int], + business_connection_id: Optional[str], bot: Bot, state_groups: Dict[str, Type[StatesGroup]], ): @@ -32,6 +33,7 @@ def __init__( self.user_id = user_id self.chat_id = chat_id self.thread_id = thread_id + self.business_connection_id = business_connection_id self.bot = bot self.lock_stack = AsyncExitStack() @@ -120,6 +122,7 @@ def _context_key(self, intent_id: str) -> StorageKey: chat_id=self.chat_id, user_id=self.chat_id, thread_id=self.thread_id, + business_connection_id=self.business_connection_id, destiny=f"aiogd:context:{intent_id}", ) @@ -137,6 +140,7 @@ def _stack_key(self, stack_id: str) -> StorageKey: chat_id=self.chat_id, user_id=self.chat_id, thread_id=self.thread_id, + business_connection_id=self.business_connection_id, destiny=f"aiogd:stack:{stack_id}", ) diff --git a/src/aiogram_dialog/manager/bg_manager.py b/src/aiogram_dialog/manager/bg_manager.py index d1955f87..9054f6ca 100644 --- a/src/aiogram_dialog/manager/bg_manager.py +++ b/src/aiogram_dialog/manager/bg_manager.py @@ -1,5 +1,5 @@ from logging import getLogger -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Union from aiogram import Bot, Router from aiogram.fsm.state import State @@ -20,7 +20,9 @@ from aiogram_dialog.api.internal import ( FakeChat, FakeUser, ) -from aiogram_dialog.api.protocols import BaseDialogManager, BgManagerFactory +from aiogram_dialog.api.protocols import ( + BaseDialogManager, BgManagerFactory, UnsetId, +) from aiogram_dialog.manager.updater import Updater from aiogram_dialog.utils import is_chat_loaded, is_user_loaded @@ -37,6 +39,7 @@ def __init__( intent_id: Optional[str], stack_id: Optional[str], thread_id: Optional[int], + business_connection_id: Optional[str], load: bool = False, ): self.user = user @@ -47,6 +50,7 @@ def __init__( self.intent_id = intent_id self.stack_id = stack_id self.thread_id = thread_id + self.business_connection_id = business_connection_id self.load = load def bg( @@ -54,7 +58,8 @@ def bg( user_id: Optional[int] = None, chat_id: Optional[int] = None, stack_id: Optional[str] = None, - thread_id: Optional[int] = None, + thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, ) -> "BaseDialogManager": if chat_id in (None, self.chat.id): @@ -67,7 +72,11 @@ def bg( else: user = FakeUser(id=user_id, is_bot=False, first_name="") - same_chat = user.id == self.user.id and chat.id == self.chat.id + same_chat = ( + user.id == self.user.id and + chat.id == self.chat.id and + business_connection_id == self.business_connection_id + ) if stack_id is None: if same_chat: stack_id = self.stack_id @@ -78,8 +87,16 @@ def bg( else: intent_id = None - if thread_id is None and same_chat: - thread_id = self.thread_id + if thread_id is UnsetId.UNSET: + if same_chat: + thread_id = self.thread_id + else: + thread_id = None + if business_connection_id is UnsetId.UNSET: + if same_chat: + business_connection_id = self.business_connection_id + else: + business_connection_id = None return BgManager( user=user, @@ -89,6 +106,7 @@ def bg( intent_id=intent_id, stack_id=stack_id, thread_id=thread_id, + business_connection_id=business_connection_id, load=load, ) @@ -99,6 +117,7 @@ def _base_event_params(self): "intent_id": self.intent_id, "stack_id": self.stack_id, "thread_id": self.thread_id, + "business_connection_id": self.business_connection_id, } async def _notify(self, event: DialogUpdateEvent): @@ -128,7 +147,9 @@ async def done( await self._load() await self._notify( DialogUpdateEvent( - action=DialogAction.DONE, data=result, show_mode=show_mode, + action=DialogAction.DONE, + data=result, + show_mode=show_mode, **self._base_event_params(), ), ) @@ -195,6 +216,7 @@ def bg( chat_id: int, stack_id: Optional[str] = None, thread_id: Optional[int] = None, + business_connection_id: Optional[str] = None, load: bool = False, ) -> "BaseDialogManager": chat = FakeChat(id=chat_id, type="") @@ -210,5 +232,6 @@ def bg( intent_id=None, stack_id=stack_id, thread_id=thread_id, + business_connection_id=business_connection_id, load=load, ) diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index ff80960a..71722459 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -1,7 +1,9 @@ from logging import getLogger -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Union, cast from aiogram import Router +from aiogram.dispatcher.middlewares.user_context import EVENT_CONTEXT_KEY, \ + EventContext from aiogram.enums import ChatType from aiogram.fsm.state import State from aiogram.types import ( @@ -24,6 +26,7 @@ from aiogram_dialog.api.protocols import ( BaseDialogManager, DialogManager, DialogProtocol, DialogRegistryProtocol, MediaIdStorageProtocol, MessageManagerProtocol, MessageNotModified, + UnsetId, ) from aiogram_dialog.context.storage import StorageProxy from aiogram_dialog.utils import get_media_id @@ -487,7 +490,8 @@ def bg( user_id: Optional[int] = None, chat_id: Optional[int] = None, stack_id: Optional[str] = None, - thread_id: Optional[int] = None, + thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, ) -> BaseDialogManager: user = self._get_fake_user(user_id) @@ -501,8 +505,20 @@ def bg( intent_id = self.current_context().id else: stack_id = DEFAULT_STACK_ID - if thread_id is None and same_chat: - thread_id = self.middleware_data.get("event_thread_id") + + event_context = cast( + self.middleware_data.get(EVENT_CONTEXT_KEY), EventContext, + ) + if thread_id is UnsetId.UNSET: + if same_chat: + thread_id = event_context.thread_id + else: + thread_id = None + if business_connection_id is UnsetId.UNSET: + if same_chat: + business_connection_id = event_context.business_connection_id + else: + business_connection_id = None return BgManager( user=user, @@ -512,6 +528,7 @@ def bg( intent_id=intent_id, stack_id=stack_id, thread_id=thread_id, + business_connection_id=business_connection_id, load=load, ) diff --git a/src/aiogram_dialog/manager/sub_manager.py b/src/aiogram_dialog/manager/sub_manager.py index 04dfa75b..6d2f11f3 100644 --- a/src/aiogram_dialog/manager/sub_manager.py +++ b/src/aiogram_dialog/manager/sub_manager.py @@ -1,5 +1,5 @@ import dataclasses -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Union from aiogram.fsm.state import State from aiogram.types import Message @@ -10,7 +10,9 @@ ) from aiogram_dialog.api.entities import Context, Stack from aiogram_dialog.api.internal import Widget -from aiogram_dialog.api.protocols import BaseDialogManager, DialogManager +from aiogram_dialog.api.protocols import ( + BaseDialogManager, DialogManager, UnsetId, +) class SubManager(DialogManager): @@ -142,7 +144,8 @@ def bg( user_id: Optional[int] = None, chat_id: Optional[int] = None, stack_id: Optional[str] = None, - thread_id: Optional[int] = None, + thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, ) -> BaseDialogManager: return self.manager.bg( @@ -150,5 +153,6 @@ def bg( chat_id=chat_id, stack_id=stack_id, thread_id=thread_id, + business_connection_id=business_connection_id, load=load, ) diff --git a/src/aiogram_dialog/tools/preview.py b/src/aiogram_dialog/tools/preview.py index 5605127f..f58dc554 100644 --- a/src/aiogram_dialog/tools/preview.py +++ b/src/aiogram_dialog/tools/preview.py @@ -2,7 +2,7 @@ import logging from dataclasses import dataclass from datetime import datetime -from typing import Any, Dict, List, Optional, Type +from typing import Any, Dict, List, Optional, Type, Union from aiogram import Router from aiogram.fsm.state import State, StatesGroup @@ -29,6 +29,7 @@ StartMode, ) from aiogram_dialog.api.exceptions import NoContextError +from aiogram_dialog.api.protocols import UnsetId from aiogram_dialog.setup import collect_dialogs from aiogram_dialog.utils import split_reply_callback @@ -211,8 +212,12 @@ async def update( def bg( self, - user_id: Optional[int] = None, chat_id: Optional[int] = None, - stack_id: Optional[str] = None, load: bool = False, + user_id: Optional[int] = None, + chat_id: Optional[int] = None, + stack_id: Optional[str] = None, + thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, + load: bool = False, ) -> BaseDialogManager: return self From d07dea28735d47e96cb143c65f850fbf929bc959 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Fri, 21 Jun 2024 23:47:59 +0200 Subject: [PATCH 29/44] use business connection id --- example/mega/bot.py | 1 + src/aiogram_dialog/api/entities/new_message.py | 2 ++ src/aiogram_dialog/manager/manager.py | 17 ++++++++++++----- src/aiogram_dialog/manager/message_manager.py | 16 +++++++++++++++- src/aiogram_dialog/setup.py | 4 ++++ .../test_tools/mock_message_manager.py | 1 + src/aiogram_dialog/window.py | 10 ++++++++-- 7 files changed, 43 insertions(+), 8 deletions(-) diff --git a/example/mega/bot.py b/example/mega/bot.py index 0971e724..d373b80b 100644 --- a/example/mega/bot.py +++ b/example/mega/bot.py @@ -75,6 +75,7 @@ def setup_dp(): storage = MemoryStorage() dp = Dispatcher(storage=storage) dp.message.register(start, F.text == "/start") + dp.business_message.register(start, F.text == "/start") dp.errors.register( on_unknown_intent, ExceptionTypeFilter(UnknownIntent), diff --git a/src/aiogram_dialog/api/entities/new_message.py b/src/aiogram_dialog/api/entities/new_message.py index 52499591..a0f16221 100644 --- a/src/aiogram_dialog/api/entities/new_message.py +++ b/src/aiogram_dialog/api/entities/new_message.py @@ -26,12 +26,14 @@ class OldMessage: media_uniq_id: Optional[str] text: Union[str, None, UnknownText] = None has_reply_keyboard: bool = False + business_connection_id: Optional[str] = None @dataclass class NewMessage: chat: Chat thread_id: Optional[int] = None + business_connection_id: Optional[str] = None text: Optional[str] = None reply_markup: Optional[MarkupVariant] = None parse_mode: Optional[str] = None diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index 71722459..7e4acd9f 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -373,7 +373,9 @@ def _get_message_from_callback( ) -> Optional[OldMessage]: current_message = event.message stack = self.current_stack() - chat = self.middleware_data["event_chat"] + event_context = cast( + EventContext, self.middleware_data.get(EVENT_CONTEXT_KEY), + ) if current_message: media_id = get_media_id(current_message) return OldMessage( @@ -381,8 +383,9 @@ def _get_message_from_callback( media_uniq_id=(media_id.file_unique_id if media_id else None), text=current_message.text, has_reply_keyboard=self.is_event_simulated(), - chat=chat, + chat=event_context.chat, message_id=current_message.message_id, + business_connection_id=event_context.business_connection_id, ) elif not stack or not stack.last_message_id: return None @@ -392,8 +395,9 @@ def _get_message_from_callback( media_uniq_id=None, text=UnknownText.UNKNOWN, has_reply_keyboard=self.is_event_simulated(), - chat=chat, + chat=event_context.chat, message_id=stack.last_message_id, + business_connection_id=event_context.business_connection_id, ) def _get_last_message(self) -> Optional[OldMessage]: @@ -405,16 +409,19 @@ def _get_last_message(self) -> Optional[OldMessage]: return self._get_message_from_callback(event) stack = self.current_stack() - chat = self.middleware_data["event_chat"] if not stack or not stack.last_message_id: return None + event_context = cast( + EventContext, self.middleware_data.get(EVENT_CONTEXT_KEY), + ) return OldMessage( media_id=stack.last_media_id, media_uniq_id=stack.last_media_unique_id, text=UnknownText.UNKNOWN, has_reply_keyboard=stack.last_reply_keyboard, - chat=chat, + chat=event_context.chat, message_id=stack.last_message_id, + business_connection_id=event_context.business_connection_id, ) def _save_last_message(self, message: OldMessage): diff --git a/src/aiogram_dialog/manager/message_manager.py b/src/aiogram_dialog/manager/message_manager.py index ba2e5a7c..773c9cb2 100644 --- a/src/aiogram_dialog/manager/message_manager.py +++ b/src/aiogram_dialog/manager/message_manager.py @@ -65,6 +65,7 @@ def _combine(sent_message: NewMessage, message_result: Message) -> OldMessage: text=message_result.text, media_uniq_id=(media_id.file_unique_id if media_id else None), media_id=(media_id.file_id if media_id else None), + business_connection_id=message_result.business_connection_id, ) @@ -221,6 +222,7 @@ async def remove_inline_kbd( return await bot.edit_message_reply_markup( message_id=old_message.message_id, chat_id=old_message.chat.id, + business_connection_id=old_message.business_connection_id, ) except TelegramBadRequest as err: if "message is not modified" in err.message: @@ -244,6 +246,7 @@ async def remove_reply_kbd( chat=old_message.chat, text="...", reply_markup=ReplyKeyboardRemove(), + business_connection_id=old_message.business_connection_id, ), ) @@ -253,6 +256,9 @@ async def remove_message_safe( old_message: OldMessage, new_message: Optional[NewMessage], ) -> None: + if old_message.business_connection_id: + await self._remove_kbd(bot, old_message, new_message) + return try: await bot.delete_message( chat_id=old_message.chat.id, @@ -300,6 +306,7 @@ async def edit_caption( return await bot.edit_message_caption( message_id=old_message.message_id, chat_id=old_message.chat.id, + business_connection_id=new_message.business_connection_id, caption=new_message.text, reply_markup=new_message.reply_markup, parse_mode=new_message.parse_mode, @@ -312,6 +319,7 @@ async def edit_text( return await bot.edit_message_text( message_id=old_message.message_id, chat_id=old_message.chat.id, + business_connection_id=new_message.business_connection_id, text=new_message.text, reply_markup=new_message.reply_markup, parse_mode=new_message.parse_mode, @@ -349,11 +357,16 @@ async def send_message(self, bot: Bot, new_message: NewMessage) -> Message: return await self.send_text(bot, new_message) async def send_text(self, bot: Bot, new_message: NewMessage) -> Message: - logger.debug("send_text to chat %s, thread %s", new_message.chat.id, new_message.thread_id) + logger.debug( + "send_text to chat %s, thread %s, business_id %s", + new_message.chat.id, new_message.thread_id, + new_message.business_connection_id, + ) return await bot.send_message( new_message.chat.id, text=new_message.text, message_thread_id=new_message.thread_id, + business_connection_id=new_message.business_connection_id, disable_web_page_preview=new_message.disable_web_page_preview, reply_markup=new_message.reply_markup, parse_mode=new_message.parse_mode, @@ -374,6 +387,7 @@ async def send_media(self, bot: Bot, new_message: NewMessage) -> Message: new_message.chat.id, await self.get_media_source(new_message.media, bot), message_thread_id=new_message.thread_id, + business_connection_id=new_message.business_connection_id, caption=new_message.text, reply_markup=new_message.reply_markup, parse_mode=new_message.parse_mode, diff --git a/src/aiogram_dialog/setup.py b/src/aiogram_dialog/setup.py index e9761ac2..ad18481d 100644 --- a/src/aiogram_dialog/setup.py +++ b/src/aiogram_dialog/setup.py @@ -114,6 +114,7 @@ def _register_middleware( )) router.message.middleware(manager_middleware) + router.business_message.middleware(manager_middleware) router.callback_query.middleware(manager_middleware) update_handler.middleware(manager_middleware) router.my_chat_member.middleware(manager_middleware) @@ -121,6 +122,7 @@ def _register_middleware( router.errors.middleware(manager_middleware) router.message.outer_middleware(intent_middleware.process_message) + router.business_message.outer_middleware(intent_middleware.process_message) router.callback_query.outer_middleware( intent_middleware.process_callback_query, ) @@ -135,12 +137,14 @@ def _register_middleware( ) router.message.outer_middleware(context_unlocker_middleware) + router.business_message.outer_middleware(context_unlocker_middleware) router.callback_query.outer_middleware(context_unlocker_middleware) update_handler.outer_middleware(context_unlocker_middleware) router.my_chat_member.outer_middleware(context_unlocker_middleware) router.chat_join_request.outer_middleware(context_unlocker_middleware) router.message.middleware(context_saver_middleware) + router.business_message.middleware(context_saver_middleware) router.callback_query.middleware(context_saver_middleware) update_handler.middleware(context_saver_middleware) router.my_chat_member.middleware(context_saver_middleware) diff --git a/src/aiogram_dialog/test_tools/mock_message_manager.py b/src/aiogram_dialog/test_tools/mock_message_manager.py index 1c2e1288..42b66af8 100644 --- a/src/aiogram_dialog/test_tools/mock_message_manager.py +++ b/src/aiogram_dialog/test_tools/mock_message_manager.py @@ -138,4 +138,5 @@ async def show_message(self, bot: Bot, new_message: NewMessage, has_reply_keyboard=isinstance( new_message.reply_markup, ReplyKeyboardMarkup, ), + business_connection_id=None, ) diff --git a/src/aiogram_dialog/window.py b/src/aiogram_dialog/window.py index 98c7d504..e4ed4b01 100644 --- a/src/aiogram_dialog/window.py +++ b/src/aiogram_dialog/window.py @@ -1,6 +1,8 @@ from logging import getLogger -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, cast +from aiogram.dispatcher.middlewares.user_context import EVENT_CONTEXT_KEY, \ + EventContext from aiogram.fsm.state import State from aiogram.types import ( CallbackQuery, @@ -121,9 +123,13 @@ async def render( logger.error("Cannot get window data for state %s", self.state) raise try: + event_context = cast( + EventContext, manager.middleware_data.get(EVENT_CONTEXT_KEY), + ) return NewMessage( chat=chat, - thread_id=manager.middleware_data.get("event_thread_id"), + thread_id=event_context.thread_id, + business_connection_id=event_context.business_connection_id, text=await self.render_text(current_data, manager), reply_markup=await self.render_kbd(current_data, manager), parse_mode=self.parse_mode, From 52ac37023c0f03721a2e46a660a6763512e71933 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 23 Jun 2024 12:55:23 +0200 Subject: [PATCH 30/44] use own event context instead of aiogram --- example/multistack.py | 1 + src/aiogram_dialog/api/entities/__init__.py | 4 +- src/aiogram_dialog/api/entities/events.py | 18 +- .../context/intent_middleware.py | 154 ++++++++++++++---- src/aiogram_dialog/manager/manager.py | 5 +- src/aiogram_dialog/manager/updater.py | 1 + src/aiogram_dialog/window.py | 3 +- 7 files changed, 142 insertions(+), 44 deletions(-) diff --git a/example/multistack.py b/example/multistack.py index 4af52ff3..14e70ca3 100644 --- a/example/multistack.py +++ b/example/multistack.py @@ -104,6 +104,7 @@ async def main(): # register handler which resets stack and start dialogs on /start command dp.message.register(start, CommandStart()) + dp.business_message.register(start, CommandStart()) setup_dialogs(dp) await dp.start_polling(bot) diff --git a/src/aiogram_dialog/api/entities/__init__.py b/src/aiogram_dialog/api/entities/__init__.py index 01c89236..c7ff5d84 100644 --- a/src/aiogram_dialog/api/entities/__init__.py +++ b/src/aiogram_dialog/api/entities/__init__.py @@ -1,6 +1,6 @@ __all__ = [ "Context", "Data", - "ChatEvent", + "ChatEvent", "EVENT_CONTEXT_KEY", "EventContext", "LaunchMode", "MediaAttachment", "MediaId", "ShowMode", "StartMode", @@ -11,7 +11,7 @@ ] from .context import Context, Data -from .events import ChatEvent +from .events import ChatEvent, EVENT_CONTEXT_KEY, EventContext from .launch_mode import LaunchMode from .media import MediaAttachment, MediaId from .modes import ShowMode, StartMode diff --git a/src/aiogram_dialog/api/entities/events.py b/src/aiogram_dialog/api/entities/events.py index 83f89ca3..90a6785a 100644 --- a/src/aiogram_dialog/api/entities/events.py +++ b/src/aiogram_dialog/api/entities/events.py @@ -1,7 +1,9 @@ -from typing import Union +from dataclasses import dataclass +from typing import Union, Optional +from aiogram import Bot from aiogram.types import ( - CallbackQuery, ChatJoinRequest, ChatMemberUpdated, Message, + CallbackQuery, ChatJoinRequest, ChatMemberUpdated, Message, Chat, User, ) from .update_event import DialogUpdateEvent @@ -10,3 +12,15 @@ CallbackQuery, Message, DialogUpdateEvent, ChatMemberUpdated, ChatJoinRequest, ] + + +@dataclass +class EventContext: + bot: Bot + chat: Chat + user: User + thread_id: Optional[int] + business_connection_id: Optional[str] + + +EVENT_CONTEXT_KEY = "aiogd_event_context" \ No newline at end of file diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index d618ff6a..e66d7a05 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -3,15 +3,14 @@ from aiogram import Router from aiogram.dispatcher.middlewares.base import BaseMiddleware -from aiogram.fsm.storage.base import BaseEventIsolation -from aiogram.types import CallbackQuery, Chat, Message, User +from aiogram.fsm.storage.base import BaseEventIsolation, BaseStorage +from aiogram.types import CallbackQuery, Message, ChatMemberUpdated, \ + ChatJoinRequest from aiogram.types.error_event import ErrorEvent -from aiogram.dispatcher.middlewares.user_context import ( - EVENT_CONTEXT_KEY, EventContext -) from aiogram_dialog.api.entities import ( ChatEvent, Context, DEFAULT_STACK_ID, DialogUpdateEvent, Stack, + EVENT_CONTEXT_KEY, EventContext, ) from aiogram_dialog.api.exceptions import ( InvalidStackIdError, OutdatedIntent, UnknownIntent, UnknownState, @@ -29,6 +28,78 @@ logger = getLogger(__name__) +def get_thread_id(message: Message) -> Optional[str]: + if not message.is_topic_message: + return None + return message.message_thread_id + + +def event_context_from_callback(event: CallbackQuery) -> EventContext: + return EventContext( + bot=event.bot, + user=event.from_user, + chat=event.message.chat, + thread_id=( + get_thread_id(event.message) + if isinstance(event.message, Message) + else None + ), + business_connection_id=event.message.business_connection_id, + ) + + +def event_context_from_chat_member(event: ChatMemberUpdated) -> EventContext: + return EventContext( + bot=event.bot, + user=event.from_user, + chat=event.chat, + thread_id=None, + business_connection_id=None, + ) + + +def event_context_from_chat_join(event: ChatJoinRequest) -> EventContext: + return EventContext( + bot=event.bot, + user=event.from_user, + chat=event.chat, + thread_id=None, + business_connection_id=None, + ) + + +def event_context_from_message(event: Message) -> EventContext: + return EventContext( + bot=event.bot, + user=event.from_user, + chat=event.chat, + thread_id=get_thread_id(event), + business_connection_id=event.business_connection_id, + ) + + +def event_context_from_aiogd(event: DialogUpdateEvent) -> EventContext: + return EventContext( + bot=event.bot, + user=event.from_user, + chat=event.chat, + thread_id=event.thread_id, + business_connection_id=event.business_connection_id, + ) + +def event_context_from_error(event: ErrorEvent) -> EventContext: + if event.update.message: + return event_context_from_message(event.update.message) + elif event.update.business_message: + return event_context_from_message(event.update.business_message) + elif event.update.my_chat_member: + return event_context_from_chat_member(event.update.my_chat_member) + elif event.update.chat_join_request: + return event_context_from_chat_join(event.update.chat_join_request) + elif event.update.callback_query: + return event_context_from_callback(event.update.callback_query) + + class IntentMiddlewareFactory: def __init__( self, @@ -41,15 +112,16 @@ def __init__( self.access_validator = access_validator self.events_isolation = events_isolation - def storage_proxy(self, data: dict): - event_context = cast(EventContext, data.get(EVENT_CONTEXT_KEY)) + def storage_proxy( + self, event_context: EventContext, fsm_storage: BaseStorage, + ) -> StorageProxy: proxy = StorageProxy( - bot=data["bot"], - storage=data["fsm_storage"], + bot=event_context.bot, + storage=fsm_storage, events_isolation=self.events_isolation, state_groups=self.registry.states_groups(), - user_id=event_context.user_id, - chat_id=event_context.chat_id, + user_id=event_context.user.id, + chat_id=event_context.chat.id, thread_id=event_context.thread_id, business_connection_id=event_context.business_connection_id, ) @@ -97,12 +169,9 @@ async def _load_context_by_stack( stack_id: Optional[str], data: dict, ) -> None: - user = data["event_from_user"] - chat = data["event_chat"] - thread_id = data.get("event_thread_id") logger.debug( "Loading context for stack: `%s`, user: `%s`, chat: `%s`, thread: `%s`", - stack_id, user.id, chat.id, thread_id, + stack_id, proxy.user_id, proxy.chat_id, proxy.thread_id, ) stack = await self._load_stack(event, stack_id, proxy, data) if not stack: @@ -126,11 +195,9 @@ async def _load_context_by_intent( intent_id: Optional[str], data: dict, ) -> None: - user = data["event_from_user"] - chat = data["event_chat"] logger.debug( "Loading context for intent: `%s`, user: `%s`, chat: `%s`", - intent_id, user.id, chat.id, + intent_id, proxy.user_id, proxy.chat_id, ) context = await proxy.load_context(intent_id) stack = await self._load_stack(event, context.stack_id, proxy, data) @@ -147,11 +214,11 @@ async def _load_context_by_intent( data[CONTEXT_KEY] = context async def _load_default_context( - self, event: ChatEvent, data: dict, + self, event: ChatEvent, data: dict, event_context: EventContext, ) -> None: return await self._load_context_by_stack( event=event, - proxy=self.storage_proxy(data), + proxy=self.storage_proxy(event_context, data["fsm_storage"]), stack_id=DEFAULT_STACK_ID, data=data, ) @@ -179,6 +246,9 @@ async def process_message( event: Message, data: dict, ): + event_context = event_context_from_message(event) + data[EVENT_CONTEXT_KEY] = event_context + text, callback_data = split_reply_callback(event.text) if callback_data: query = ReplyCallbackQuery( @@ -201,30 +271,36 @@ async def process_message( if intent_id := self._intent_id_from_reply(event, data): await self._load_context_by_intent( event=event, - proxy=self.storage_proxy(data), + proxy=self.storage_proxy(event_context, data["fsm_storage"]), intent_id=intent_id, data=data, ) else: - await self._load_default_context(event, data) + await self._load_default_context(event, data, event_context) return await handler(event, data) async def process_my_chat_member( self, handler: Callable, - event: Message, + event: ChatMemberUpdated, data: dict, ) -> None: - await self._load_default_context(event, data) + event_context = event_context_from_chat_member(event) + data[EVENT_CONTEXT_KEY] = event_context + + await self._load_default_context(event, data, event_context) return await handler(event, data) async def process_chat_join_request( self, handler: Callable, - event: Message, + event: ChatJoinRequest, data: dict, ) -> None: - await self._load_default_context(event, data) + event_context = event_context_from_chat_join(event) + data[EVENT_CONTEXT_KEY] = event_context + + await self._load_default_context(event, data, event_context) return await handler(event, data) async def process_aiogd_update( @@ -233,17 +309,20 @@ async def process_aiogd_update( event: DialogUpdateEvent, data: dict, ): + event_context = event_context_from_aiogd(event) + data[EVENT_CONTEXT_KEY] = event_context + if event.intent_id: await self._load_context_by_intent( event=event, - proxy=self.storage_proxy(data), + proxy=self.storage_proxy(event_context, data["fsm_storage"]), intent_id=event.intent_id, data=data, ) else: await self._load_context_by_stack( event=event, - proxy=self.storage_proxy(data), + proxy=self.storage_proxy(event_context, data["fsm_storage"]), stack_id=event.stack_id, data=data, ) @@ -255,6 +334,9 @@ async def process_callback_query( event: CallbackQuery, data: dict, ): + event_context = event_context_from_callback(event) + data[EVENT_CONTEXT_KEY] = event_context + if "event_chat" not in data: return await handler(event, data) original_data = event.data @@ -263,15 +345,16 @@ async def process_callback_query( if intent_id: await self._load_context_by_intent( event=event, - proxy=self.storage_proxy(data), + proxy=self.storage_proxy(event_context, + data["fsm_storage"]), intent_id=intent_id, data=data, ) else: - await self._load_default_context(event, data) + await self._load_default_context(event, data, event_context) data[CALLBACK_DATA_KEY] = original_data else: - await self._load_default_context(event, data) + await self._load_default_context(event, data, event_context) return await handler(event, data) @@ -358,14 +441,15 @@ async def __call__( return await handler(event, data) try: - event_context = cast(EventContext, data.get(EVENT_CONTEXT_KEY)) + event_context = event_context_from_error(event) + data[EVENT_CONTEXT_KEY] = event_context proxy = StorageProxy( - bot=data["bot"], + bot=event_context.bot, storage=data["fsm_storage"], events_isolation=self.events_isolation, state_groups=self.registry.states_groups(), - user_id=event_context.user_id, - chat_id=event_context.chat_id, + user_id=event_context.user.id, + chat_id=event_context.chat.id, thread_id=event_context.thread_id, business_connection_id=event_context.business_connection_id, ) diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index 7e4acd9f..244864c9 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -2,8 +2,6 @@ from typing import Any, Dict, Optional, Union, cast from aiogram import Router -from aiogram.dispatcher.middlewares.user_context import EVENT_CONTEXT_KEY, \ - EventContext from aiogram.enums import ChatType from aiogram.fsm.state import State from aiogram.types import ( @@ -14,6 +12,7 @@ AccessSettings, ChatEvent, Context, Data, DEFAULT_STACK_ID, LaunchMode, MediaId, NewMessage, ShowMode, Stack, StartMode, + EVENT_CONTEXT_KEY, EventContext, ) from aiogram_dialog.api.entities import OldMessage, UnknownText from aiogram_dialog.api.exceptions import ( @@ -514,7 +513,7 @@ def bg( stack_id = DEFAULT_STACK_ID event_context = cast( - self.middleware_data.get(EVENT_CONTEXT_KEY), EventContext, + EventContext, self.middleware_data.get(EVENT_CONTEXT_KEY), ) if thread_id is UnsetId.UNSET: if same_chat: diff --git a/src/aiogram_dialog/manager/updater.py b/src/aiogram_dialog/manager/updater.py index df8926f7..c932ffd0 100644 --- a/src/aiogram_dialog/manager/updater.py +++ b/src/aiogram_dialog/manager/updater.py @@ -2,6 +2,7 @@ from contextvars import copy_context from aiogram import Bot, Dispatcher, Router +from aiogram.dispatcher.middlewares.user_context import EventContext from aiogram_dialog.api.entities import DialogUpdate diff --git a/src/aiogram_dialog/window.py b/src/aiogram_dialog/window.py index e4ed4b01..8a89430a 100644 --- a/src/aiogram_dialog/window.py +++ b/src/aiogram_dialog/window.py @@ -1,8 +1,6 @@ from logging import getLogger from typing import Any, Dict, List, Optional, cast -from aiogram.dispatcher.middlewares.user_context import EVENT_CONTEXT_KEY, \ - EventContext from aiogram.fsm.state import State from aiogram.types import ( CallbackQuery, @@ -13,6 +11,7 @@ from aiogram_dialog.api.entities import ( MarkupVariant, MediaAttachment, NewMessage, + EVENT_CONTEXT_KEY, EventContext, ) from aiogram_dialog.api.internal import Widget, WindowProtocol from .api.entities import Data From 5b05aa276bed93e6286579e712593e5b30e55fcf Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 23 Jun 2024 13:03:28 +0200 Subject: [PATCH 31/44] Fix bot in dialog events --- src/aiogram_dialog/manager/bg_manager.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/aiogram_dialog/manager/bg_manager.py b/src/aiogram_dialog/manager/bg_manager.py index 9054f6ca..7fca1aa1 100644 --- a/src/aiogram_dialog/manager/bg_manager.py +++ b/src/aiogram_dialog/manager/bg_manager.py @@ -121,9 +121,8 @@ def _base_event_params(self): } async def _notify(self, event: DialogUpdateEvent): - await self._updater.notify( - bot=self.bot, update=DialogUpdate(aiogd_update=event), - ) + update = DialogUpdate(aiogd_update=event.as_(self.bot)).as_(self.bot) + await self._updater.notify(bot=self.bot, update=update) async def _load(self): if self.load: From 7d4667ef710b1be265f7012dd77bfa43eae5eb7a Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Tue, 25 Jun 2024 23:58:12 +0200 Subject: [PATCH 32/44] fix bg manager and sort imports --- src/aiogram_dialog/api/entities/events.py | 18 +- src/aiogram_dialog/api/protocols/manager.py | 4 +- .../context/intent_middleware.py | 29 +++- src/aiogram_dialog/manager/bg_manager.py | 155 ++++++++++++------ src/aiogram_dialog/manager/manager.py | 84 +++++----- src/aiogram_dialog/manager/sub_manager.py | 2 +- src/aiogram_dialog/manager/updater.py | 1 - src/aiogram_dialog/tools/preview.py | 2 +- src/aiogram_dialog/window.py | 9 +- 9 files changed, 193 insertions(+), 111 deletions(-) diff --git a/src/aiogram_dialog/api/entities/events.py b/src/aiogram_dialog/api/entities/events.py index 90a6785a..bc1a2783 100644 --- a/src/aiogram_dialog/api/entities/events.py +++ b/src/aiogram_dialog/api/entities/events.py @@ -1,16 +1,24 @@ from dataclasses import dataclass -from typing import Union, Optional +from typing import Optional, Union from aiogram import Bot from aiogram.types import ( - CallbackQuery, ChatJoinRequest, ChatMemberUpdated, Message, Chat, User, + CallbackQuery, + Chat, + ChatJoinRequest, + ChatMemberUpdated, + Message, + User, ) from .update_event import DialogUpdateEvent ChatEvent = Union[ - CallbackQuery, Message, DialogUpdateEvent, - ChatMemberUpdated, ChatJoinRequest, + CallbackQuery, + ChatJoinRequest, + ChatMemberUpdated, + DialogUpdateEvent, + Message, ] @@ -23,4 +31,4 @@ class EventContext: business_connection_id: Optional[str] -EVENT_CONTEXT_KEY = "aiogd_event_context" \ No newline at end of file +EVENT_CONTEXT_KEY = "aiogd_event_context" diff --git a/src/aiogram_dialog/api/protocols/manager.py b/src/aiogram_dialog/api/protocols/manager.py index 686513bd..899c7e5e 100644 --- a/src/aiogram_dialog/api/protocols/manager.py +++ b/src/aiogram_dialog/api/protocols/manager.py @@ -58,7 +58,7 @@ def bg( chat_id: Optional[int] = None, stack_id: Optional[str] = None, thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, - business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, # load chat and user ) -> "BaseDialogManager": raise NotImplementedError @@ -73,7 +73,7 @@ def bg( chat_id: int, stack_id: Optional[str] = None, thread_id: Optional[int] = None, - business_connection_id: Optional[str] = None, + business_connection_id: Optional[str] = None, load: bool = False, # load chat and user ) -> "BaseDialogManager": raise NotImplementedError diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index e66d7a05..dcfbeb84 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -1,16 +1,25 @@ from logging import getLogger -from typing import Any, Awaitable, Callable, Dict, Optional, cast +from typing import Any, Awaitable, Callable, Dict, Optional from aiogram import Router from aiogram.dispatcher.middlewares.base import BaseMiddleware from aiogram.fsm.storage.base import BaseEventIsolation, BaseStorage -from aiogram.types import CallbackQuery, Message, ChatMemberUpdated, \ - ChatJoinRequest +from aiogram.types import ( + CallbackQuery, + ChatJoinRequest, + ChatMemberUpdated, + Message, +) from aiogram.types.error_event import ErrorEvent from aiogram_dialog.api.entities import ( - ChatEvent, Context, DEFAULT_STACK_ID, DialogUpdateEvent, Stack, - EVENT_CONTEXT_KEY, EventContext, + ChatEvent, + Context, + DEFAULT_STACK_ID, + DialogUpdateEvent, + EVENT_CONTEXT_KEY, + EventContext, + Stack, ) from aiogram_dialog.api.exceptions import ( InvalidStackIdError, OutdatedIntent, UnknownIntent, UnknownState, @@ -87,6 +96,7 @@ def event_context_from_aiogd(event: DialogUpdateEvent) -> EventContext: business_connection_id=event.business_connection_id, ) + def event_context_from_error(event: ErrorEvent) -> EventContext: if event.update.message: return event_context_from_message(event.update.message) @@ -170,7 +180,8 @@ async def _load_context_by_stack( data: dict, ) -> None: logger.debug( - "Loading context for stack: `%s`, user: `%s`, chat: `%s`, thread: `%s`", + "Loading context for stack: " + "`%s`, user: `%s`, chat: `%s`, thread: `%s`", stack_id, proxy.user_id, proxy.chat_id, proxy.thread_id, ) stack = await self._load_stack(event, stack_id, proxy, data) @@ -181,7 +192,7 @@ async def _load_context_by_stack( else: try: context = await proxy.load_context(stack.last_intent_id()) - except: + except: # noqa: B001,B901,E722 await proxy.unlock() raise data[STORAGE_KEY] = proxy @@ -205,7 +216,7 @@ async def _load_context_by_intent( return try: self._check_outdated(intent_id, stack) - except: + except: # noqa: B001,B901,E722 await proxy.unlock() raise @@ -363,6 +374,8 @@ async def process_callback_query( "callback_query", "my_chat_member", "aiogd_update", + "chat_join_request", + "business_message", } diff --git a/src/aiogram_dialog/manager/bg_manager.py b/src/aiogram_dialog/manager/bg_manager.py index 7fca1aa1..4e4adb34 100644 --- a/src/aiogram_dialog/manager/bg_manager.py +++ b/src/aiogram_dialog/manager/bg_manager.py @@ -14,6 +14,7 @@ DialogSwitchEvent, DialogUpdate, DialogUpdateEvent, + EventContext, ShowMode, StartMode, ) @@ -29,6 +30,38 @@ logger = getLogger(__name__) +def coalesce_business_connection_id( + *, + user: User, + chat: Chat, + business_connection_id: Union[str, None, UnsetId], + event_context: EventContext, +) -> Optional[str]: + if business_connection_id is not UnsetId.UNSET: + return business_connection_id + if user.id != event_context.user.id: + return None + if chat.id != event_context.chat.id: + return None + return event_context.business_connection_id + + +def coalesce_thread_id( + *, + user: User, + chat: Chat, + thread_id: Union[str, None, UnsetId], + event_context: EventContext, +) -> Optional[str]: + if thread_id is not UnsetId.UNSET: + return thread_id + if user.id != event_context.user.id: + return None + if chat.id != event_context.chat.id: + return None + return None + + class BgManager(BaseDialogManager): def __init__( self, @@ -38,47 +71,66 @@ def __init__( router: Router, intent_id: Optional[str], stack_id: Optional[str], - thread_id: Optional[int], - business_connection_id: Optional[str], + thread_id: Optional[int] = None, + business_connection_id: Optional[str] = None, load: bool = False, ): - self.user = user - self.chat = chat - self.bot = bot + self._event_context = EventContext( + chat=chat, + user=user, + bot=bot, + thread_id=thread_id, + business_connection_id=business_connection_id, + ) self._router = router self._updater = Updater(router) self.intent_id = intent_id self.stack_id = stack_id - self.thread_id = thread_id - self.business_connection_id = business_connection_id self.load = load + def _get_fake_user(self, user_id: Optional[int] = None) -> User: + """Get User if we have info about him or FakeUser instead.""" + if user_id in (None, self._event_context.user.id): + return self._event_context.user + return FakeUser(id=user_id, is_bot=False, first_name="") + + def _get_fake_chat(self, chat_id: Optional[int] = None) -> Chat: + """Get Chat if we have info about him or FakeChat instead.""" + if chat_id in (None, self._event_context.chat.id): + return self._event_context.chat + return FakeChat(id=chat_id, type="") + def bg( self, user_id: Optional[int] = None, chat_id: Optional[int] = None, stack_id: Optional[str] = None, thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, - business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, ) -> "BaseDialogManager": - if chat_id in (None, self.chat.id): - chat = self.chat - else: - chat = FakeChat(id=chat_id, type="") - - if user_id in (None, self.user.id): - user = self.user - else: - user = FakeUser(id=user_id, is_bot=False, first_name="") + chat = self._get_fake_chat(chat_id) + user = self._get_fake_user(user_id) - same_chat = ( - user.id == self.user.id and - chat.id == self.chat.id and - business_connection_id == self.business_connection_id + new_event_context = EventContext( + bot=self._event_context.bot, + chat=chat, + user=user, + thread_id=coalesce_thread_id( + chat=chat, + user=user, + thread_id=thread_id, + event_context=self._event_context, + ), + business_connection_id=coalesce_business_connection_id( + chat=chat, + user=user, + business_connection_id=business_connection_id, + event_context=self._event_context, + ), ) if stack_id is None: - if same_chat: + if self._event_context == new_event_context: stack_id = self.stack_id intent_id = self.intent_id else: @@ -87,56 +139,53 @@ def bg( else: intent_id = None - if thread_id is UnsetId.UNSET: - if same_chat: - thread_id = self.thread_id - else: - thread_id = None - if business_connection_id is UnsetId.UNSET: - if same_chat: - business_connection_id = self.business_connection_id - else: - business_connection_id = None - return BgManager( - user=user, - chat=chat, - bot=self.bot, + user=new_event_context.user, + chat=new_event_context.chat, + bot=new_event_context.bot, router=self._router, intent_id=intent_id, stack_id=stack_id, - thread_id=thread_id, - business_connection_id=business_connection_id, + thread_id=new_event_context.thread_id, + business_connection_id=new_event_context.business_connection_id, load=load, ) def _base_event_params(self): return { - "from_user": self.user, - "chat": self.chat, + "from_user": self._event_context.user, + "chat": self._event_context.chat, "intent_id": self.intent_id, "stack_id": self.stack_id, - "thread_id": self.thread_id, - "business_connection_id": self.business_connection_id, + "thread_id": self._event_context.thread_id, + "business_connection_id": + self._event_context.business_connection_id, } async def _notify(self, event: DialogUpdateEvent): - update = DialogUpdate(aiogd_update=event.as_(self.bot)).as_(self.bot) - await self._updater.notify(bot=self.bot, update=update) + bot = self._event_context.bot + update = DialogUpdate(aiogd_update=event.as_(bot)).as_(bot) + await self._updater.notify(bot=bot, update=update) async def _load(self): if self.load: - if not is_chat_loaded(self.chat): - logger.debug("load chat: %s", self.chat.id) - self.chat = await self.bot.get_chat(self.chat.id) - if not is_user_loaded(self.user): + bot = self._event_context.bot + if not is_chat_loaded(self._event_context.chat): + logger.debug( + "load chat: %s", self._event_context.chat.id, + ) + self._event_context.chat = await bot.get_chat( + self._event_context.chat.id, + ) + if not is_user_loaded(self._event_context.user): logger.debug( - "load user %s from chat %s", self.chat.id, self.user.id, + "load user %s from chat %s", + self._event_context.chat.id, self._event_context.user.id, ) - chat_member = await self.bot.get_chat_member( - self.chat.id, self.user.id, + chat_member = await bot.get_chat_member( + self._event_context.chat.id, self._event_context.user.id, ) - self.user = chat_member.user + self._event_context.user = chat_member.user async def done( self, @@ -215,7 +264,7 @@ def bg( chat_id: int, stack_id: Optional[str] = None, thread_id: Optional[int] = None, - business_connection_id: Optional[str] = None, + business_connection_id: Optional[str] = None, load: bool = False, ) -> "BaseDialogManager": chat = FakeChat(id=chat_id, type="") diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index 244864c9..5a8ee9e8 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -1,5 +1,5 @@ from logging import getLogger -from typing import Any, Dict, Optional, Union, cast +from typing import Any, cast, Dict, Optional, Union from aiogram import Router from aiogram.enums import ChatType @@ -10,9 +10,18 @@ from aiogram_dialog.api.entities import ( AccessSettings, - ChatEvent, Context, Data, DEFAULT_STACK_ID, LaunchMode, MediaId, - NewMessage, ShowMode, Stack, StartMode, - EVENT_CONTEXT_KEY, EventContext, + ChatEvent, + Context, + Data, + DEFAULT_STACK_ID, + EVENT_CONTEXT_KEY, + EventContext, + LaunchMode, + MediaId, + NewMessage, + ShowMode, + Stack, + StartMode, ) from aiogram_dialog.api.entities import OldMessage, UnknownText from aiogram_dialog.api.exceptions import ( @@ -29,8 +38,11 @@ ) from aiogram_dialog.context.storage import StorageProxy from aiogram_dialog.utils import get_media_id -from .bg_manager import BgManager - +from .bg_manager import ( + BgManager, + coalesce_business_connection_id, + coalesce_thread_id, +) logger = getLogger(__name__) @@ -418,7 +430,7 @@ def _get_last_message(self) -> Optional[OldMessage]: media_uniq_id=stack.last_media_unique_id, text=UnknownText.UNKNOWN, has_reply_keyboard=stack.last_reply_keyboard, - chat=event_context.chat, + chat=event_context.chat, message_id=stack.last_message_id, business_connection_id=event_context.business_connection_id, ) @@ -463,14 +475,6 @@ def find(self, widget_id) -> Optional[Any]: return None return widget.managed(self) - def is_same_chat(self, user: User, chat: Chat) -> bool: - if "event_chat" not in self._data: - return False - - current_chat = self._data["event_chat"] - current_user = self.event.from_user - return user.id == current_user.id and chat.id == current_chat.id - def _get_fake_user(self, user_id: Optional[int] = None) -> User: """Get User if we have info about him or FakeUser instead.""" current_user = self.event.from_user @@ -497,44 +501,50 @@ def bg( chat_id: Optional[int] = None, stack_id: Optional[str] = None, thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, - business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, ) -> BaseDialogManager: user = self._get_fake_user(user_id) chat = self._get_fake_chat(chat_id) intent_id = None - same_chat = self.is_same_chat(user, chat) + event_context = cast( + EventContext, self.middleware_data.get(EVENT_CONTEXT_KEY), + ) + new_event_context = EventContext( + bot=event_context.bot, + chat=chat, + user=user, + thread_id=coalesce_thread_id( + chat=chat, + user=user, + thread_id=thread_id, + event_context=event_context, + ), + business_connection_id=coalesce_business_connection_id( + chat=chat, + user=user, + business_connection_id=business_connection_id, + event_context=event_context, + ), + ) + if stack_id is None: - if same_chat: + if event_context == new_event_context: stack_id = self.current_stack().id if self.has_context(): intent_id = self.current_context().id else: stack_id = DEFAULT_STACK_ID - event_context = cast( - EventContext, self.middleware_data.get(EVENT_CONTEXT_KEY), - ) - if thread_id is UnsetId.UNSET: - if same_chat: - thread_id = event_context.thread_id - else: - thread_id = None - if business_connection_id is UnsetId.UNSET: - if same_chat: - business_connection_id = event_context.business_connection_id - else: - business_connection_id = None - return BgManager( - user=user, - chat=chat, - bot=self._data["bot"], + user=new_event_context.user, + chat=new_event_context.chat, + bot=new_event_context.bot, router=self._router, intent_id=intent_id, stack_id=stack_id, - thread_id=thread_id, - business_connection_id=business_connection_id, + thread_id=new_event_context.thread_id, + business_connection_id=new_event_context.business_connection_id, load=load, ) diff --git a/src/aiogram_dialog/manager/sub_manager.py b/src/aiogram_dialog/manager/sub_manager.py index 6d2f11f3..86069a34 100644 --- a/src/aiogram_dialog/manager/sub_manager.py +++ b/src/aiogram_dialog/manager/sub_manager.py @@ -145,7 +145,7 @@ def bg( chat_id: Optional[int] = None, stack_id: Optional[str] = None, thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, - business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, ) -> BaseDialogManager: return self.manager.bg( diff --git a/src/aiogram_dialog/manager/updater.py b/src/aiogram_dialog/manager/updater.py index c932ffd0..df8926f7 100644 --- a/src/aiogram_dialog/manager/updater.py +++ b/src/aiogram_dialog/manager/updater.py @@ -2,7 +2,6 @@ from contextvars import copy_context from aiogram import Bot, Dispatcher, Router -from aiogram.dispatcher.middlewares.user_context import EventContext from aiogram_dialog.api.entities import DialogUpdate diff --git a/src/aiogram_dialog/tools/preview.py b/src/aiogram_dialog/tools/preview.py index f58dc554..c509dce4 100644 --- a/src/aiogram_dialog/tools/preview.py +++ b/src/aiogram_dialog/tools/preview.py @@ -216,7 +216,7 @@ def bg( chat_id: Optional[int] = None, stack_id: Optional[str] = None, thread_id: Union[int, None, UnsetId] = UnsetId.UNSET, - business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, + business_connection_id: Union[str, None, UnsetId] = UnsetId.UNSET, load: bool = False, ) -> BaseDialogManager: return self diff --git a/src/aiogram_dialog/window.py b/src/aiogram_dialog/window.py index 8a89430a..f3ed80cd 100644 --- a/src/aiogram_dialog/window.py +++ b/src/aiogram_dialog/window.py @@ -1,5 +1,5 @@ from logging import getLogger -from typing import Any, Dict, List, Optional, cast +from typing import Any, cast, Dict, List, Optional from aiogram.fsm.state import State from aiogram.types import ( @@ -10,8 +10,11 @@ from aiogram.types.base import UNSET_DISABLE_WEB_PAGE_PREVIEW from aiogram_dialog.api.entities import ( - MarkupVariant, MediaAttachment, NewMessage, - EVENT_CONTEXT_KEY, EventContext, + EVENT_CONTEXT_KEY, + EventContext, + MarkupVariant, + MediaAttachment, + NewMessage, ) from aiogram_dialog.api.internal import Widget, WindowProtocol from .api.entities import Data From b8e182467753a7b560dc1f4b9075fabfa471c738 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 26 Jun 2024 00:25:42 +0200 Subject: [PATCH 33/44] Auto callback answer if stack forbidden, allow callback answer in tests --- src/aiogram_dialog/context/intent_middleware.py | 12 +++++++++--- src/aiogram_dialog/test_tools/bot_client.py | 10 ++++++++-- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index dcfbeb84..cbdbf466 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -2,6 +2,7 @@ from typing import Any, Awaitable, Callable, Dict, Optional from aiogram import Router +from aiogram.dispatcher.event.bases import UNHANDLED from aiogram.dispatcher.middlewares.base import BaseMiddleware from aiogram.fsm.storage.base import BaseEventIsolation, BaseStorage from aiogram.types import ( @@ -36,6 +37,8 @@ logger = getLogger(__name__) +FORBIDDEN_STACK_KEY = "aiogd_stack_forbidden" + def get_thread_id(message: Message) -> Optional[str]: if not message.is_topic_message: @@ -163,11 +166,11 @@ async def _load_stack( raise InvalidStackIdError("Both stack id and intent id are None") stack = await proxy.load_stack(stack_id) if not await self.access_validator.is_allowed(stack, event, data): - user = data["event_from_user"] logger.debug( "Stack %s is not allowed for user %s", - stack.id, user.id, + stack.id, proxy.user_id, ) + data[FORBIDDEN_STACK_KEY] = True await proxy.unlock() return return stack @@ -366,7 +369,10 @@ async def process_callback_query( data[CALLBACK_DATA_KEY] = original_data else: await self._load_default_context(event, data, event_context) - return await handler(event, data) + result = await handler(event, data) + if result is UNHANDLED and data.get(FORBIDDEN_STACK_KEY): + await event.answer() + return result SUPPORTED_ERROR_EVENTS = { diff --git a/src/aiogram_dialog/test_tools/bot_client.py b/src/aiogram_dialog/test_tools/bot_client.py index 7fa8d857..468098fa 100644 --- a/src/aiogram_dialog/test_tools/bot_client.py +++ b/src/aiogram_dialog/test_tools/bot_client.py @@ -1,8 +1,9 @@ import uuid from datetime import datetime -from typing import Optional, Union +from typing import Optional, Union, Any from aiogram import Bot, Dispatcher +from aiogram.methods import TelegramMethod, AnswerCallbackQuery from aiogram.types import ( CallbackQuery, Chat, ChatJoinRequest, ChatMemberAdministrator, ChatMemberBanned, ChatMemberLeft, @@ -21,7 +22,12 @@ def __init__(self): def id(self): return 1 - def __call__(self, *args, **kwargs) -> None: + async def __call__( + self, method: TelegramMethod[Any], + request_timeout: Optional[int] = None, + ) -> Any: + if isinstance(method, AnswerCallbackQuery): + return True raise RuntimeError("Fake bot should not be used to call telegram") def __hash__(self) -> int: From 8aeb99f99db79ebbf4500f871aa8389aa385df8e Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 26 Jun 2024 23:36:23 +0200 Subject: [PATCH 34/44] fix processing group stack in business chat --- src/aiogram_dialog/context/storage.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 75d57d83..02477191 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -129,8 +129,10 @@ def _context_key(self, intent_id: str) -> StorageKey: def _fixed_stack_id(self, stack_id: str) -> str: if stack_id != DEFAULT_STACK_ID: return stack_id + # private chat has chat_id=user_id and no business connection if self.user_id in (None, self.chat_id): - return stack_id + if self.business_connection_id is None: + return stack_id return f"<{self.user_id}>" def _stack_key(self, stack_id: str) -> StorageKey: From f675c8edaa9cbeb492befdebd5a6e02a8f9130d7 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Sun, 14 Jul 2024 23:23:01 +0200 Subject: [PATCH 35/44] 2.2.0a5 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3bf4b7a7..949986a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ where = ["src"] [project] name = "aiogram_dialog" -version = "2.2.0a4" +version = "2.2.0a5" readme = "README.md" authors = [ { name = "Andrey Tikhonov", email = "17@itishka.org" }, From f7151a83c055ddb6b31269716dd591661d93a279 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 17 Jul 2024 18:55:47 +0200 Subject: [PATCH 36/44] Save access settings per context --- src/aiogram_dialog/api/entities/__init__.py | 3 ++- src/aiogram_dialog/api/entities/access.py | 11 +++++++++++ src/aiogram_dialog/api/entities/context.py | 6 +++++- src/aiogram_dialog/api/entities/stack.py | 11 ++--------- src/aiogram_dialog/context/storage.py | 19 +++++++++---------- src/aiogram_dialog/manager/manager.py | 9 +++++++-- 6 files changed, 36 insertions(+), 23 deletions(-) create mode 100644 src/aiogram_dialog/api/entities/access.py diff --git a/src/aiogram_dialog/api/entities/__init__.py b/src/aiogram_dialog/api/entities/__init__.py index c7ff5d84..38b04326 100644 --- a/src/aiogram_dialog/api/entities/__init__.py +++ b/src/aiogram_dialog/api/entities/__init__.py @@ -10,13 +10,14 @@ "DialogStartEvent", "DialogSwitchEvent", "DialogUpdate", ] +from .access import AccessSettings from .context import Context, Data from .events import ChatEvent, EVENT_CONTEXT_KEY, EventContext from .launch_mode import LaunchMode from .media import MediaAttachment, MediaId from .modes import ShowMode, StartMode from .new_message import MarkupVariant, NewMessage, OldMessage, UnknownText -from .stack import AccessSettings, DEFAULT_STACK_ID, GROUP_STACK_ID, Stack +from .stack import DEFAULT_STACK_ID, GROUP_STACK_ID, Stack from .update_event import ( DIALOG_EVENT_NAME, DialogAction, DialogStartEvent, DialogSwitchEvent, DialogUpdate, DialogUpdateEvent, diff --git a/src/aiogram_dialog/api/entities/access.py b/src/aiogram_dialog/api/entities/access.py new file mode 100644 index 00000000..edab28b6 --- /dev/null +++ b/src/aiogram_dialog/api/entities/access.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass +from typing import Any, List, Optional + +from aiogram.enums import ChatMemberStatus + + +@dataclass +class AccessSettings: + user_ids: List[int] + member_status: Optional[ChatMemberStatus] = None + custom: Any = None diff --git a/src/aiogram_dialog/api/entities/context.py b/src/aiogram_dialog/api/entities/context.py index 60c4555a..262da6f1 100644 --- a/src/aiogram_dialog/api/entities/context.py +++ b/src/aiogram_dialog/api/entities/context.py @@ -1,7 +1,8 @@ from dataclasses import dataclass, field -from typing import Dict, List, Union +from typing import Dict, List, Optional, Union from aiogram.fsm.state import State +from .access import AccessSettings Data = Union[Dict, List, int, str, float, None] DataDict = Dict[str, Data] @@ -15,6 +16,9 @@ class Context: start_data: Data = field(compare=False) dialog_data: DataDict = field(compare=False, default_factory=dict) widget_data: DataDict = field(compare=False, default_factory=dict) + access_settings: Optional[AccessSettings] = field( + compare=False, default=None, + ) @property def id(self) -> str: diff --git a/src/aiogram_dialog/api/entities/stack.py b/src/aiogram_dialog/api/entities/stack.py index cd5d46f5..30fbd8cb 100644 --- a/src/aiogram_dialog/api/entities/stack.py +++ b/src/aiogram_dialog/api/entities/stack.py @@ -2,12 +2,12 @@ import string import time from dataclasses import dataclass, field -from typing import Any, List, Optional +from typing import List, Optional -from aiogram.enums import ChatMemberStatus from aiogram.fsm.state import State from aiogram_dialog.api.exceptions import DialogStackOverflow +from .access import AccessSettings from .context import Context, Data DEFAULT_STACK_ID = "" @@ -35,13 +35,6 @@ def new_id(): return id_to_str(new_int_id()) -@dataclass -class AccessSettings: - user_ids: List[int] - member_status: Optional[ChatMemberStatus] = None - custom: Any = None - - @dataclass(unsafe_hash=True) class Stack: _id: str = field(compare=True, default_factory=new_id) diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 02477191..91902d6c 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -53,11 +53,14 @@ async def load_context(self, intent_id: str) -> Context: raise UnknownIntent( f"Context not found for intent id: {intent_id}", ) + data["access_settings"] = self._parse_access_settings( + data.pop("access_settings", None), + ) data["state"] = self._state(data["state"]) return Context(**data) def _default_access_settings(self, stack_id: str) -> AccessSettings: - if stack_id == DEFAULT_STACK_ID: + if stack_id == DEFAULT_STACK_ID and self.user_id: return AccessSettings(user_ids=[self.user_id]) else: return AccessSettings(user_ids=[]) @@ -67,13 +70,9 @@ async def load_stack(self, stack_id: str = DEFAULT_STACK_ID) -> Stack: key = self._stack_key(fixed_stack_id) await self.lock(key) data = await self.storage.get_data(key) + access_settings = self._default_access_settings(stack_id) if not data: - access_settings = self._default_access_settings(stack_id) - return Stack(_id=fixed_stack_id, access_settings=access_settings) - - access_settings = self._parse_access_settings( - data.pop("access_settings", None), - ) + return Stack(_id=stack_id, access_settings=access_settings) return Stack(access_settings=access_settings, **data) async def save_context(self, context: Optional[Context]) -> None: @@ -81,6 +80,9 @@ async def save_context(self, context: Optional[Context]) -> None: return data = copy(vars(context)) data["state"] = data["state"].state + data["access_settings"] = self._dump_access_settings( + context.access_settings, + ) await self.storage.set_data( key=self._context_key(context.id), data=data, @@ -108,9 +110,6 @@ async def save_stack(self, stack: Optional[Stack]) -> None: ) else: data = copy(vars(stack)) - data["access_settings"] = self._dump_access_settings( - stack.access_settings, - ) await self.storage.set_data( key=self._stack_key(stack.id), data=data, diff --git a/src/aiogram_dialog/manager/manager.py b/src/aiogram_dialog/manager/manager.py index 5a8ee9e8..113cec19 100644 --- a/src/aiogram_dialog/manager/manager.py +++ b/src/aiogram_dialog/manager/manager.py @@ -1,3 +1,4 @@ +from copy import deepcopy from logging import getLogger from typing import Any, cast, Dict, Optional, Union @@ -255,8 +256,6 @@ async def _start_normal( access_settings: Optional[AccessSettings], ) -> None: stack = self.current_stack() - if access_settings is not None: - stack.access_settings = access_settings old_dialog: Optional[DialogProtocol] = None if not stack.empty(): old_dialog = self.dialog() @@ -270,7 +269,13 @@ async def _start_normal( await self._process_launch_mode(old_dialog, new_dialog) if self.has_context(): await self.storage().save_context(self.current_context()) + if access_settings is None: + access_settings = self.current_context().access_settings + if access_settings is None: + access_settings = stack.access_settings + context = stack.push(state, data) + context.access_settings = deepcopy(access_settings) self._data[CONTEXT_KEY] = context await self.dialog().process_start(self, data, state) new_context = self._current_context_unsafe() From 3f9bd5d8638055e93c0ce100d707c6b60b4ac2cf Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 17 Jul 2024 20:59:55 +0200 Subject: [PATCH 37/44] validate access --- .../context/intent_middleware.py | 42 ++++++++++++++----- src/aiogram_dialog/setup.py | 1 + 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index cbdbf466..b6ddc710 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -165,14 +165,6 @@ async def _load_stack( if stack_id is None: raise InvalidStackIdError("Both stack id and intent id are None") stack = await proxy.load_stack(stack_id) - if not await self.access_validator.is_allowed(stack, event, data): - logger.debug( - "Stack %s is not allowed for user %s", - stack.id, proxy.user_id, - ) - data[FORBIDDEN_STACK_KEY] = True - await proxy.unlock() - return return stack async def _load_context_by_stack( @@ -198,6 +190,15 @@ async def _load_context_by_stack( except: # noqa: B001,B901,E722 await proxy.unlock() raise + + if not await self.access_validator.is_allowed(stack, event, data): + logger.debug( + "Stack %s is not allowed for user %s", + stack.id, proxy.user_id, + ) + data[FORBIDDEN_STACK_KEY] = True + await proxy.unlock() + return data[STORAGE_KEY] = proxy data[STACK_KEY] = stack data[CONTEXT_KEY] = context @@ -223,6 +224,14 @@ async def _load_context_by_intent( await proxy.unlock() raise + if not await self.access_validator.is_allowed(stack, event, data): + logger.debug( + "Stack %s is not allowed for user %s", + stack.id, proxy.user_id, + ) + data[FORBIDDEN_STACK_KEY] = True + await proxy.unlock() + return data[STORAGE_KEY] = proxy data[STACK_KEY] = stack data[CONTEXT_KEY] = context @@ -408,11 +417,13 @@ class IntentErrorMiddleware(BaseMiddleware): def __init__( self, registry: DialogRegistryProtocol, + access_validator: StackAccessValidator, events_isolation: BaseEventIsolation, ): super().__init__() self.registry = registry self.events_isolation = events_isolation + self.access_validator = access_validator def _is_error_supported( self, event: ErrorEvent, data: Dict[str, Any], @@ -484,8 +495,19 @@ async def __call__( storage=proxy, stack=stack, ) - data[STACK_KEY] = stack - data[CONTEXT_KEY] = context + + if await self.access_validator.is_allowed( + stack, event.update.event, data, + ): + data[STACK_KEY] = stack + data[CONTEXT_KEY] = context + else: + logger.debug( + "Stack %s is not allowed for user %s", + stack.id, proxy.user_id, + ) + data[FORBIDDEN_STACK_KEY] = True + await proxy.unlock() return await handler(event, data) finally: proxy: StorageProxy = data.pop(STORAGE_KEY, None) diff --git a/src/aiogram_dialog/setup.py b/src/aiogram_dialog/setup.py index ad18481d..c7825799 100644 --- a/src/aiogram_dialog/setup.py +++ b/src/aiogram_dialog/setup.py @@ -111,6 +111,7 @@ def _register_middleware( router.errors.middleware(IntentErrorMiddleware( registry=registry, events_isolation=events_isolation, + access_validator=stack_access_validator, )) router.message.middleware(manager_middleware) From 054da0462676993ac42e22ffe94482f2ba01a9ba Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 17 Jul 2024 21:03:02 +0200 Subject: [PATCH 38/44] compat with previous alpha --- src/aiogram_dialog/context/storage.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index 91902d6c..ab628de9 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -70,6 +70,7 @@ async def load_stack(self, stack_id: str = DEFAULT_STACK_ID) -> Stack: key = self._stack_key(fixed_stack_id) await self.lock(key) data = await self.storage.get_data(key) + data.pop("access_settings", None) # compat with 2.2a5 access_settings = self._default_access_settings(stack_id) if not data: return Stack(_id=stack_id, access_settings=access_settings) From 84b83ff7e9dd0c61482c2c81d74d4e1451b9da20 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Fri, 19 Jul 2024 00:12:48 +0200 Subject: [PATCH 39/44] Fix access settings processing --- .../api/protocols/stack_access.py | 10 +++-- .../context/access_validator.py | 20 +++++++--- .../context/intent_middleware.py | 10 +++-- src/aiogram_dialog/context/storage.py | 2 +- tests/test_group.py | 39 ++++++++++++++++++- 5 files changed, 67 insertions(+), 14 deletions(-) diff --git a/src/aiogram_dialog/api/protocols/stack_access.py b/src/aiogram_dialog/api/protocols/stack_access.py index 30091553..eb7eaee3 100644 --- a/src/aiogram_dialog/api/protocols/stack_access.py +++ b/src/aiogram_dialog/api/protocols/stack_access.py @@ -1,13 +1,17 @@ from abc import abstractmethod -from typing import Protocol +from typing import Protocol, Optional from aiogram_dialog import ChatEvent -from aiogram_dialog.api.entities import Stack +from aiogram_dialog.api.entities import Stack, Context class StackAccessValidator(Protocol): @abstractmethod async def is_allowed( - self, stack: Stack, event: ChatEvent, data: dict, + self, + stack: Stack, + context: Optional[Context], + event: ChatEvent, + data: dict, ) -> bool: raise NotImplementedError diff --git a/src/aiogram_dialog/context/access_validator.py b/src/aiogram_dialog/context/access_validator.py index f6a33e51..d2a3c902 100644 --- a/src/aiogram_dialog/context/access_validator.py +++ b/src/aiogram_dialog/context/access_validator.py @@ -1,10 +1,11 @@ from logging import getLogger +from typing import Optional from aiogram.enums import ChatType from aiogram_dialog import ChatEvent from aiogram_dialog.api.entities import ( - Stack, + Stack, Context, ) from aiogram_dialog.api.protocols import StackAccessValidator @@ -13,15 +14,24 @@ class DefaultAccessValidator(StackAccessValidator): async def is_allowed( - self, stack: Stack, event: ChatEvent, data: dict, + self, + stack: Stack, + context: Optional[Context], + event: ChatEvent, + data: dict, ) -> bool: - if not stack.access_settings: + if context: + access_settings = context.access_settings + else: + access_settings = stack.access_settings + + if not access_settings: return True chat = data["event_chat"] if chat.type is ChatType.PRIVATE: return True - if stack.access_settings.user_ids: + if access_settings.user_ids: user = data["event_from_user"] - if user.id not in stack.access_settings.user_ids: + if user.id not in access_settings.user_ids: return False return True diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index b6ddc710..bf6b5f56 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -191,7 +191,9 @@ async def _load_context_by_stack( await proxy.unlock() raise - if not await self.access_validator.is_allowed(stack, event, data): + if not await self.access_validator.is_allowed( + stack, context, event, data, + ): logger.debug( "Stack %s is not allowed for user %s", stack.id, proxy.user_id, @@ -224,7 +226,9 @@ async def _load_context_by_intent( await proxy.unlock() raise - if not await self.access_validator.is_allowed(stack, event, data): + if not await self.access_validator.is_allowed( + stack, context, event, data, + ): logger.debug( "Stack %s is not allowed for user %s", stack.id, proxy.user_id, @@ -497,7 +501,7 @@ async def __call__( ) if await self.access_validator.is_allowed( - stack, event.update.event, data, + stack, context, event.update.event, data, ): data[STACK_KEY] = stack data[CONTEXT_KEY] = context diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index ab628de9..b1122acb 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -73,7 +73,7 @@ async def load_stack(self, stack_id: str = DEFAULT_STACK_ID) -> Stack: data.pop("access_settings", None) # compat with 2.2a5 access_settings = self._default_access_settings(stack_id) if not data: - return Stack(_id=stack_id, access_settings=access_settings) + return Stack(_id=fixed_stack_id, access_settings=access_settings) return Stack(access_settings=access_settings, **data) async def save_context(self, context: Optional[Context]) -> None: diff --git a/tests/test_group.py b/tests/test_group.py index db08463a..3c150050 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -3,13 +3,13 @@ import pytest from aiogram import Dispatcher -from aiogram.filters import CommandStart +from aiogram.filters import CommandStart, Command from aiogram.fsm.state import State, StatesGroup from aiogram_dialog import ( Dialog, DialogManager, setup_dialogs, StartMode, Window, ) -from aiogram_dialog.api.entities import GROUP_STACK_ID +from aiogram_dialog.api.entities import GROUP_STACK_ID, AccessSettings from aiogram_dialog.test_tools import BotClient, MockMessageManager from aiogram_dialog.test_tools.keyboard import InlineButtonTextLocator from aiogram_dialog.widgets.kbd import Button @@ -36,6 +36,12 @@ async def start_shared(event: Any, dialog_manager: DialogManager): await dialog_manager.start(MainSG.start, mode=StartMode.RESET_STACK) +async def add_shared(event: Any, dialog_manager: DialogManager): + await dialog_manager.start(MainSG.start, access_settings=AccessSettings( + user_ids=[1, 2], + )) + + @pytest.fixture() def message_manager(): return MockMessageManager() @@ -74,6 +80,35 @@ async def test_second_user(dp, client, second_client, message_manager): assert not message_manager.sent_messages +@pytest.mark.asyncio +async def test_change_seettings(dp, client, second_client, message_manager): + dp.message.register(start, CommandStart()) + dp.message.register(add_shared, Command("add")) + + await client.send("/start") + message_manager.reset_history() + + await client.send("/add") + window_message = message_manager.one_message() + message_manager.reset_history() + + await second_client.click( + window_message, InlineButtonTextLocator("Button"), + ) + window_message = message_manager.one_message() + message_manager.reset_history() + assert window_message.text == "stub" + + await client.send("/start") + window_message = message_manager.one_message() + message_manager.reset_history() + + await second_client.click( + window_message, InlineButtonTextLocator("Button"), + ) + assert not message_manager.sent_messages + + @pytest.mark.asyncio async def test_same_user(dp, client, message_manager): dp.message.register(start, CommandStart()) From 1b7bf760ba422a1e0b568d8e3a3e41b08d793bd1 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Mon, 19 Aug 2024 20:59:46 +0200 Subject: [PATCH 40/44] check if need resend dialog --- src/aiogram_dialog/__init__.py | 4 +- src/aiogram_dialog/api/internal/window.py | 6 ++- src/aiogram_dialog/api/protocols/__init__.py | 4 +- src/aiogram_dialog/api/protocols/dialog.py | 4 ++ src/aiogram_dialog/dialog.py | 51 +++++++++++++++----- src/aiogram_dialog/tools/preview.py | 23 ++++----- src/aiogram_dialog/widgets/utils.py | 6 +-- src/aiogram_dialog/window.py | 14 ++++-- tests/test_group.py | 35 ++++++++++++-- 9 files changed, 109 insertions(+), 38 deletions(-) diff --git a/src/aiogram_dialog/__init__.py b/src/aiogram_dialog/__init__.py index 54fb11cb..82d2a1c1 100644 --- a/src/aiogram_dialog/__init__.py +++ b/src/aiogram_dialog/__init__.py @@ -7,6 +7,7 @@ "StartMode", "BaseDialogManager", "BgManagerFactory", + "CancelEventProcessing", "DialogManager", "DialogProtocol", "UnsetId", @@ -22,7 +23,8 @@ ChatEvent, Data, DEFAULT_STACK_ID, LaunchMode, ShowMode, StartMode, ) from .api.protocols import ( - BaseDialogManager, BgManagerFactory, DialogManager, DialogProtocol, + BaseDialogManager, BgManagerFactory, CancelEventProcessing, + DialogManager, DialogProtocol, UnsetId, ) from .dialog import Dialog diff --git a/src/aiogram_dialog/api/internal/window.py b/src/aiogram_dialog/api/internal/window.py index 503678e7..2dce844c 100644 --- a/src/aiogram_dialog/api/internal/window.py +++ b/src/aiogram_dialog/api/internal/window.py @@ -19,7 +19,8 @@ async def process_message( message: Message, dialog: "DialogProtocol", manager: DialogManager, - ) -> None: + ) -> bool: + """Return True if message in handled.""" raise NotImplementedError @abstractmethod @@ -28,7 +29,8 @@ async def process_callback( callback: CallbackQuery, dialog: "DialogProtocol", manager: DialogManager, - ) -> None: + ) -> bool: + """Return True if callback in handled.""" raise NotImplementedError @abstractmethod diff --git a/src/aiogram_dialog/api/protocols/__init__.py b/src/aiogram_dialog/api/protocols/__init__.py index b83cef41..73f1ecb1 100644 --- a/src/aiogram_dialog/api/protocols/__init__.py +++ b/src/aiogram_dialog/api/protocols/__init__.py @@ -1,5 +1,5 @@ __all__ = [ - "DialogProtocol", + "CancelEventProcessing", "DialogProtocol", "BaseDialogManager", "BgManagerFactory", "DialogManager", "UnsetId", "MediaIdStorageProtocol", "MessageManagerProtocol", "MessageNotModified", @@ -7,7 +7,7 @@ "StackAccessValidator", ] -from .dialog import DialogProtocol +from .dialog import CancelEventProcessing, DialogProtocol from .manager import ( BaseDialogManager, BgManagerFactory, DialogManager, UnsetId, ) diff --git a/src/aiogram_dialog/api/protocols/dialog.py b/src/aiogram_dialog/api/protocols/dialog.py index 40fac5cc..941566dd 100644 --- a/src/aiogram_dialog/api/protocols/dialog.py +++ b/src/aiogram_dialog/api/protocols/dialog.py @@ -9,6 +9,10 @@ from .manager import DialogManager +class CancelEventProcessing(Exception): + pass + + @runtime_checkable class DialogProtocol(Protocol): @property diff --git a/src/aiogram_dialog/dialog.py b/src/aiogram_dialog/dialog.py index bb6932b4..3afe0777 100644 --- a/src/aiogram_dialog/dialog.py +++ b/src/aiogram_dialog/dialog.py @@ -13,15 +13,16 @@ from aiogram import Router from aiogram.fsm.state import State, StatesGroup -from aiogram.types import CallbackQuery, Message +from aiogram.types import CallbackQuery, Message, Chat +from aiogram.enums import ChatType -from aiogram_dialog.api.entities import Data, LaunchMode, NewMessage +from aiogram_dialog.api.entities import Data, LaunchMode, NewMessage, Context from aiogram_dialog.api.exceptions import ( UnregisteredWindowError, ) from aiogram_dialog.api.internal import Widget, WindowProtocol from aiogram_dialog.api.protocols import ( - DialogManager, DialogProtocol, + CancelEventProcessing, DialogManager, DialogProtocol, ) from .context.intent_filter import IntentFilter from .utils import remove_intent_id @@ -128,10 +129,14 @@ async def _message_handler( ): old_context = dialog_manager.current_context() window = await self._current_window(dialog_manager) - await window.process_message(message, self, dialog_manager) - if dialog_manager.has_context(): - if dialog_manager.current_context() == old_context: # same dialog - await dialog_manager.show() + try: + processed = await window.process_message( + message, self, dialog_manager, + ) + except CancelEventProcessing: + processed = False + if self._need_refresh(processed, old_context, dialog_manager): + await dialog_manager.show() async def _callback_handler( self, @@ -142,12 +147,36 @@ async def _callback_handler( intent_id, callback_data = remove_intent_id(callback.data) cleaned_callback = callback.model_copy(update={"data": callback_data}) window = await self._current_window(dialog_manager) - await window.process_callback(cleaned_callback, self, dialog_manager) - if dialog_manager.has_context(): - if dialog_manager.current_context() == old_context: # same dialog - await dialog_manager.show() + try: + processed = await window.process_callback( + cleaned_callback, self, dialog_manager, + ) + except CancelEventProcessing: + processed = False + if self._need_refresh(processed, old_context, dialog_manager): + await dialog_manager.show() await dialog_manager.answer_callback() + def _need_refresh( + self, processed: bool, + old_context: Context, + dialog_manager: DialogManager, + ): + if not dialog_manager.has_context(): + # nothing to show + return False + if dialog_manager.current_context() != old_context: + # dialog switched, so it is already refreshed + return False + if processed: + # something happened + return True + event_chat: Chat = dialog_manager.middleware_data["event_chat"] + if event_chat.type == ChatType.PRIVATE: + # for private chats we can ensure dialog is visible + return True + return False + def _setup_filter(self): intent_filter = IntentFilter( aiogd_intent_state_group=self.states_group(), diff --git a/src/aiogram_dialog/tools/preview.py b/src/aiogram_dialog/tools/preview.py index c509dce4..296d411d 100644 --- a/src/aiogram_dialog/tools/preview.py +++ b/src/aiogram_dialog/tools/preview.py @@ -26,7 +26,7 @@ NewMessage, ShowMode, Stack, - StartMode, + StartMode, EVENT_CONTEXT_KEY, EventContext, ) from aiogram_dialog.api.exceptions import NoContextError from aiogram_dialog.api.protocols import UnsetId @@ -72,13 +72,21 @@ def __init__(self): intent_id=None, stack_id=None, thread_id=None, + business_connection_id=None, ) self._context: Optional[Context] = None self._dialog: Optional[DialogProtocol] = None self._data = { "dialog_manager": self, - "event_chat": Chat(id=1, type="private"), - "event_from_user": User(id=1, is_bot=False, first_name="Fake"), + "event_chat": self._event.chat, + "event_from_user": self._event.from_user, + EVENT_CONTEXT_KEY: EventContext( + bot=None, + thread_id=None, + chat=self._event.chat, + user= self._event.from_user, + business_connection_id=None, + ), } async def next(self, show_mode: Optional[ShowMode] = None) -> None: @@ -94,7 +102,7 @@ async def back(self, show_mode: Optional[ShowMode] = None) -> None: await self.switch_to(new_state, show_mode) @property - def data(self) -> Dict: + def middleware_data(self) -> Dict: return self._data @property @@ -120,13 +128,6 @@ def set_state(self, state: State): def is_preview(self) -> bool: return True - @property - def middleware_data(self) -> Dict: - return { - "event_chat": self.event.chat, - "dialog_manager": self, - } - @property def dialog_data(self) -> Dict: return self.current_context().dialog_data diff --git a/src/aiogram_dialog/widgets/utils.py b/src/aiogram_dialog/widgets/utils.py index cd0c36bb..3deb8807 100644 --- a/src/aiogram_dialog/widgets/utils.py +++ b/src/aiogram_dialog/widgets/utils.py @@ -45,12 +45,12 @@ def ensure_input( BaseInput, Sequence[BaseInput], ], -) -> BaseInput: +) -> Union[BaseInput, None]: if isinstance(widget, BaseInput): return widget elif isinstance(widget, Sequence): if len(widget) == 0: - return MessageInput(None) + return None elif len(widget) == 1: return widget[0] else: @@ -71,7 +71,7 @@ def ensure_media(widget: Union[Media, Sequence[Media]]) -> Media: def ensure_widgets( widgets: Sequence[WidgetSrc], -) -> Tuple[Text, Keyboard, BaseInput, Media]: +) -> Tuple[Text, Keyboard, Union[BaseInput, None], Media]: texts = [] keyboards = [] inputs = [] diff --git a/src/aiogram_dialog/window.py b/src/aiogram_dialog/window.py index f3ed80cd..af02a0d2 100644 --- a/src/aiogram_dialog/window.py +++ b/src/aiogram_dialog/window.py @@ -96,16 +96,22 @@ async def load_data( async def process_message( self, message: Message, dialog: DialogProtocol, manager: DialogManager, - ) -> None: + ) -> bool: if self.on_message: - await self.on_message.process_message(message, dialog, manager) + return await self.on_message.process_message( + message, dialog, manager, + ) + return False async def process_callback( self, callback: CallbackQuery, dialog: DialogProtocol, manager: DialogManager, - ) -> None: + ) -> bool: if self.keyboard: - await self.keyboard.process_callback(callback, dialog, manager) + return await self.keyboard.process_callback( + callback, dialog, manager, + ) + return False async def process_result( self, start_data: Data, result: Any, manager: DialogManager, diff --git a/tests/test_group.py b/tests/test_group.py index 3c150050..bada638e 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -81,7 +81,36 @@ async def test_second_user(dp, client, second_client, message_manager): @pytest.mark.asyncio -async def test_change_seettings(dp, client, second_client, message_manager): +async def test_change_settings(dp, client, second_client, message_manager): + dp.message.register(start, CommandStart()) + dp.message.register(add_shared, Command("add")) + + await client.send("/start") + message_manager.reset_history() + + await client.send("/add") + window_message = message_manager.one_message() + message_manager.reset_history() + + await second_client.click( + window_message, InlineButtonTextLocator("Button"), + ) + window_message = message_manager.one_message() + message_manager.reset_history() + assert window_message.text == "stub" + + await client.send("/start") + window_message = message_manager.one_message() + message_manager.reset_history() + + await second_client.click( + window_message, InlineButtonTextLocator("Button"), + ) + assert not message_manager.sent_messages + + +@pytest.mark.asyncio +async def test_change_settings_bg(dp, client, second_client, message_manager): dp.message.register(start, CommandStart()) dp.message.register(add_shared, Command("add")) @@ -117,9 +146,7 @@ async def test_same_user(dp, client, message_manager): assert first_message.text == "stub" message_manager.reset_history() await client.send("test") - first_message = message_manager.one_message() - assert first_message.text == "stub" - message_manager.reset_history() + assert not message_manager.sent_messages # no resend await client.click( first_message, InlineButtonTextLocator("Button"), ) From cd124e021ce29163374755efb301b3adb77e465d Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Mon, 19 Aug 2024 22:22:33 +0200 Subject: [PATCH 41/44] Fix linter errors, fix unlocking on error --- src/aiogram_dialog/api/entities/context.py | 1 + src/aiogram_dialog/api/protocols/stack_access.py | 4 ++-- src/aiogram_dialog/context/access_validator.py | 2 +- src/aiogram_dialog/context/intent_middleware.py | 14 ++++++++++---- src/aiogram_dialog/dialog.py | 6 +++--- src/aiogram_dialog/test_tools/bot_client.py | 4 ++-- src/aiogram_dialog/tools/preview.py | 6 ++++-- tests/test_group.py | 4 ++-- 8 files changed, 25 insertions(+), 16 deletions(-) diff --git a/src/aiogram_dialog/api/entities/context.py b/src/aiogram_dialog/api/entities/context.py index 262da6f1..b4882088 100644 --- a/src/aiogram_dialog/api/entities/context.py +++ b/src/aiogram_dialog/api/entities/context.py @@ -2,6 +2,7 @@ from typing import Dict, List, Optional, Union from aiogram.fsm.state import State + from .access import AccessSettings Data = Union[Dict, List, int, str, float, None] diff --git a/src/aiogram_dialog/api/protocols/stack_access.py b/src/aiogram_dialog/api/protocols/stack_access.py index eb7eaee3..0f3561ce 100644 --- a/src/aiogram_dialog/api/protocols/stack_access.py +++ b/src/aiogram_dialog/api/protocols/stack_access.py @@ -1,8 +1,8 @@ from abc import abstractmethod -from typing import Protocol, Optional +from typing import Optional, Protocol from aiogram_dialog import ChatEvent -from aiogram_dialog.api.entities import Stack, Context +from aiogram_dialog.api.entities import Context, Stack class StackAccessValidator(Protocol): diff --git a/src/aiogram_dialog/context/access_validator.py b/src/aiogram_dialog/context/access_validator.py index d2a3c902..1120be57 100644 --- a/src/aiogram_dialog/context/access_validator.py +++ b/src/aiogram_dialog/context/access_validator.py @@ -5,7 +5,7 @@ from aiogram_dialog import ChatEvent from aiogram_dialog.api.entities import ( - Stack, Context, + Context, Stack, ) from aiogram_dialog.api.protocols import StackAccessValidator diff --git a/src/aiogram_dialog/context/intent_middleware.py b/src/aiogram_dialog/context/intent_middleware.py index bf6b5f56..b20d0d6e 100644 --- a/src/aiogram_dialog/context/intent_middleware.py +++ b/src/aiogram_dialog/context/intent_middleware.py @@ -462,6 +462,14 @@ async def _load_last_context( await self._fix_broken_stack(storage, stack) return None + async def _load_stack( + self, proxy: StorageProxy, error: Exception, + ) -> Stack: + if isinstance(error, OutdatedIntent): + return await proxy.load_stack(stack_id=error.stack_id) + else: + return await proxy.load_stack() + async def __call__( self, handler: Callable[ @@ -488,10 +496,7 @@ async def __call__( business_connection_id=event_context.business_connection_id, ) data[STORAGE_KEY] = proxy - if isinstance(error, OutdatedIntent): - stack = await proxy.load_stack(stack_id=error.stack_id) - else: - stack = await proxy.load_stack() + stack = await self._load_stack(proxy, event) if stack.empty() or isinstance(error, UnknownState): context = None else: @@ -511,6 +516,7 @@ async def __call__( stack.id, proxy.user_id, ) data[FORBIDDEN_STACK_KEY] = True + del data[STORAGE_KEY] await proxy.unlock() return await handler(event, data) finally: diff --git a/src/aiogram_dialog/dialog.py b/src/aiogram_dialog/dialog.py index 3afe0777..164d045c 100644 --- a/src/aiogram_dialog/dialog.py +++ b/src/aiogram_dialog/dialog.py @@ -12,11 +12,11 @@ ) from aiogram import Router -from aiogram.fsm.state import State, StatesGroup -from aiogram.types import CallbackQuery, Message, Chat from aiogram.enums import ChatType +from aiogram.fsm.state import State, StatesGroup +from aiogram.types import CallbackQuery, Chat, Message -from aiogram_dialog.api.entities import Data, LaunchMode, NewMessage, Context +from aiogram_dialog.api.entities import Context, Data, LaunchMode, NewMessage from aiogram_dialog.api.exceptions import ( UnregisteredWindowError, ) diff --git a/src/aiogram_dialog/test_tools/bot_client.py b/src/aiogram_dialog/test_tools/bot_client.py index 468098fa..24735dbf 100644 --- a/src/aiogram_dialog/test_tools/bot_client.py +++ b/src/aiogram_dialog/test_tools/bot_client.py @@ -1,9 +1,9 @@ import uuid from datetime import datetime -from typing import Optional, Union, Any +from typing import Any, Optional, Union from aiogram import Bot, Dispatcher -from aiogram.methods import TelegramMethod, AnswerCallbackQuery +from aiogram.methods import AnswerCallbackQuery, TelegramMethod from aiogram.types import ( CallbackQuery, Chat, ChatJoinRequest, ChatMemberAdministrator, ChatMemberBanned, ChatMemberLeft, diff --git a/src/aiogram_dialog/tools/preview.py b/src/aiogram_dialog/tools/preview.py index 296d411d..63b24188 100644 --- a/src/aiogram_dialog/tools/preview.py +++ b/src/aiogram_dialog/tools/preview.py @@ -22,11 +22,13 @@ Data, DialogAction, DialogUpdateEvent, + EVENT_CONTEXT_KEY, + EventContext, MediaAttachment, NewMessage, ShowMode, Stack, - StartMode, EVENT_CONTEXT_KEY, EventContext, + StartMode, ) from aiogram_dialog.api.exceptions import NoContextError from aiogram_dialog.api.protocols import UnsetId @@ -84,7 +86,7 @@ def __init__(self): bot=None, thread_id=None, chat=self._event.chat, - user= self._event.from_user, + user=self._event.from_user, business_connection_id=None, ), } diff --git a/tests/test_group.py b/tests/test_group.py index bada638e..70e8008e 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -3,13 +3,13 @@ import pytest from aiogram import Dispatcher -from aiogram.filters import CommandStart, Command +from aiogram.filters import Command, CommandStart from aiogram.fsm.state import State, StatesGroup from aiogram_dialog import ( Dialog, DialogManager, setup_dialogs, StartMode, Window, ) -from aiogram_dialog.api.entities import GROUP_STACK_ID, AccessSettings +from aiogram_dialog.api.entities import AccessSettings, GROUP_STACK_ID from aiogram_dialog.test_tools import BotClient, MockMessageManager from aiogram_dialog.test_tools.keyboard import InlineButtonTextLocator from aiogram_dialog.widgets.kbd import Button From 0c2fb85f6a5331981c49482d99b80fd854d7bdcb Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Mon, 19 Aug 2024 22:24:36 +0200 Subject: [PATCH 42/44] fix unused --- src/aiogram_dialog/test_tools/bot_client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/aiogram_dialog/test_tools/bot_client.py b/src/aiogram_dialog/test_tools/bot_client.py index 24735dbf..961b5a92 100644 --- a/src/aiogram_dialog/test_tools/bot_client.py +++ b/src/aiogram_dialog/test_tools/bot_client.py @@ -26,6 +26,7 @@ async def __call__( self, method: TelegramMethod[Any], request_timeout: Optional[int] = None, ) -> Any: + del request_timeout # unused if isinstance(method, AnswerCallbackQuery): return True raise RuntimeError("Fake bot should not be used to call telegram") From 921ff65164ac910fe549a5bcd2018ef4ae75fe76 Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 21 Aug 2024 22:30:21 +0200 Subject: [PATCH 43/44] Docs for group dialogs --- docs/group_business.rst | 58 +++++++++++++++++++ docs/how_are_messages_updated/index.rst | 2 +- docs/index.rst | 1 + src/aiogram_dialog/__init__.py | 5 +- src/aiogram_dialog/api/entities/access.py | 5 +- .../api/protocols/stack_access.py | 1 + 6 files changed, 66 insertions(+), 6 deletions(-) create mode 100644 docs/group_business.rst diff --git a/docs/group_business.rst b/docs/group_business.rst new file mode 100644 index 00000000..862195dd --- /dev/null +++ b/docs/group_business.rst @@ -0,0 +1,58 @@ +****************************************** +Groups and business chats +****************************************** + +.. warning:: + Telegram has very strong limitations on amount of operations in groups, + so it is not recommended to use interactive menus there + +Support of groups, supergroups and business chats is based on usage of additional dialog stacks. + +Starting shared dialogs +================================= + +When user sends message or other event not attached directly to some dialog, default stack is used. If you start dialogs in that stack, they can be accessed only by that user. So, the default stack in **personal**. + +To send a *shared* dialog from *personal*, you need to use other stacks. It can be ``aiogram_dialog.GROUP_STACK_ID``, other predefined string or starting via ``StartMode.NEW_STACK``. + +.. code-block:: python + + bg = dialog_manager.bg(stack_id=GROUP_STACK_ID) + bg.start( + MyStateGroup.MY_STATE, + mode=StartMode.RESET_STACK, + ) + +If there are different topics in chat, stacks between them are isolated. To start dialog in different topic pass ``thread_id`` as ``.bg()`` argument + +Limiting access +====================== + +To set limitations on who can interact with that dialog, you can pass ``AccessSettings`` when starting new dialog. If not access settings are set, they will be copied from last opened dialog in stack. + +.. code-block:: python + + dialog_manager.start( + MyStateGroup.MY_STATE, + mode=StartMode.RESET_STACK, + access_settings=AccessSettings(user_ids=[123456]), + ) + +In this example, pre-defined group stack will be used and new dialogs will be available only for user with id ``123456``. If later user clicks on a specific dialog, stack of that dialog is used, so you won't need to call ``.bg()`` + +Currently, only check by ``user.id`` is supported, but you bring your own logic implementing ``StackAccessValidator`` protocol and passing it so ``setup_dialogs`` function. + +Handling forbidden interactions +================================= + +If user is not allowed to interact with dialog his event is not routed to dialogs and you can handle it in aiogram. To filter this situation you can rely on ``aiogd_stack_forbidden`` key of middleware data. + +Classes +=========== + + +.. autoclass:: aiogram_dialog.AccessSettings + + +.. autoclass:: aiogram_dialog.api.protocols.StackAccessValidator + :members: is_allowed diff --git a/docs/how_are_messages_updated/index.rst b/docs/how_are_messages_updated/index.rst index 7c4a8818..af2cf6ff 100644 --- a/docs/how_are_messages_updated/index.rst +++ b/docs/how_are_messages_updated/index.rst @@ -1,5 +1,5 @@ How are messages updated -===================== +=========================== ShowMode ******************** diff --git a/docs/index.rst b/docs/index.rst index a812232d..12874a06 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -10,6 +10,7 @@ aiogram-dialog widgets/index transitions/index how_are_messages_updated/index + group_business helper_tools/index migration faq diff --git a/src/aiogram_dialog/__init__.py b/src/aiogram_dialog/__init__.py index 82d2a1c1..d6d649e5 100644 --- a/src/aiogram_dialog/__init__.py +++ b/src/aiogram_dialog/__init__.py @@ -1,7 +1,9 @@ __all__ = [ + "AccessSettings", "DEFAULT_STACK_ID", "Dialog", "Data", + "GROUP_STACK_ID", "ChatEvent", "LaunchMode", "StartMode", @@ -20,7 +22,8 @@ import importlib.metadata as _metadata from .api.entities import ( - ChatEvent, Data, DEFAULT_STACK_ID, LaunchMode, ShowMode, StartMode, + AccessSettings, ChatEvent, Data, DEFAULT_STACK_ID, GROUP_STACK_ID, + LaunchMode, ShowMode, StartMode, ) from .api.protocols import ( BaseDialogManager, BgManagerFactory, CancelEventProcessing, diff --git a/src/aiogram_dialog/api/entities/access.py b/src/aiogram_dialog/api/entities/access.py index edab28b6..bf49c19f 100644 --- a/src/aiogram_dialog/api/entities/access.py +++ b/src/aiogram_dialog/api/entities/access.py @@ -1,11 +1,8 @@ from dataclasses import dataclass -from typing import Any, List, Optional - -from aiogram.enums import ChatMemberStatus +from typing import Any, List @dataclass class AccessSettings: user_ids: List[int] - member_status: Optional[ChatMemberStatus] = None custom: Any = None diff --git a/src/aiogram_dialog/api/protocols/stack_access.py b/src/aiogram_dialog/api/protocols/stack_access.py index 0f3561ce..79158e00 100644 --- a/src/aiogram_dialog/api/protocols/stack_access.py +++ b/src/aiogram_dialog/api/protocols/stack_access.py @@ -14,4 +14,5 @@ async def is_allowed( event: ChatEvent, data: dict, ) -> bool: + """Check if current user is allowed to interactor with dialog.""" raise NotImplementedError From e82658614e10f33b4158b0942c0405ceb3704ecb Mon Sep 17 00:00:00 2001 From: Andrey Tikhonov <17@itishka.org> Date: Wed, 21 Aug 2024 22:33:21 +0200 Subject: [PATCH 44/44] Fix after access settings cleanup --- src/aiogram_dialog/context/storage.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/aiogram_dialog/context/storage.py b/src/aiogram_dialog/context/storage.py index b1122acb..17cbd0a0 100644 --- a/src/aiogram_dialog/context/storage.py +++ b/src/aiogram_dialog/context/storage.py @@ -3,7 +3,6 @@ from typing import Dict, Optional, Type from aiogram import Bot -from aiogram.enums import ChatMemberStatus from aiogram.fsm.state import State, StatesGroup from aiogram.fsm.storage.base import ( BaseEventIsolation, BaseStorage, StorageKey, @@ -161,13 +160,8 @@ def _parse_access_settings( ) -> Optional[AccessSettings]: if not raw: return None - if raw_member_status := raw.get("member_status"): - member_status = ChatMemberStatus(raw_member_status) - else: - member_status = None return AccessSettings( user_ids=raw.get("user_ids") or [], - member_status=member_status, custom=raw.get("custom"), ) @@ -178,6 +172,5 @@ def _dump_access_settings( return None return { "user_ids": access_settings.user_ids, - "member_status": access_settings.member_status, "custom": access_settings.custom, }