mirror of
https://gitlab.com/MoonTestUse1/AdministrationItDepartmens.git
synced 2025-08-14 00:25:46 +02:00
599 lines
23 KiB
Python
599 lines
23 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import signal
|
|
import warnings
|
|
from asyncio import CancelledError, Event, Future, Lock
|
|
from contextlib import suppress
|
|
from typing import Any, AsyncGenerator, Dict, List, Optional, Set, Union
|
|
|
|
from .. import loggers
|
|
from ..client.bot import Bot
|
|
from ..exceptions import TelegramAPIError
|
|
from ..fsm.middleware import FSMContextMiddleware
|
|
from ..fsm.storage.base import BaseEventIsolation, BaseStorage
|
|
from ..fsm.storage.memory import DisabledEventIsolation, MemoryStorage
|
|
from ..fsm.strategy import FSMStrategy
|
|
from ..methods import GetUpdates, TelegramMethod
|
|
from ..methods.base import TelegramType
|
|
from ..types import Update, User
|
|
from ..types.base import UNSET, UNSET_TYPE
|
|
from ..types.update import UpdateTypeLookupError
|
|
from ..utils.backoff import Backoff, BackoffConfig
|
|
from .event.bases import UNHANDLED, SkipHandler
|
|
from .event.telegram import TelegramEventObserver
|
|
from .middlewares.error import ErrorsMiddleware
|
|
from .middlewares.user_context import UserContextMiddleware
|
|
from .router import Router
|
|
|
|
DEFAULT_BACKOFF_CONFIG = BackoffConfig(min_delay=1.0, max_delay=5.0, factor=1.3, jitter=0.1)
|
|
|
|
|
|
class Dispatcher(Router):
|
|
"""
|
|
Root router
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*, # * - Preventing to pass instance of Bot to the FSM storage
|
|
storage: Optional[BaseStorage] = None,
|
|
fsm_strategy: FSMStrategy = FSMStrategy.USER_IN_CHAT,
|
|
events_isolation: Optional[BaseEventIsolation] = None,
|
|
disable_fsm: bool = False,
|
|
name: Optional[str] = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""
|
|
Root router
|
|
|
|
:param storage: Storage for FSM
|
|
:param fsm_strategy: FSM strategy
|
|
:param events_isolation: Events isolation
|
|
:param disable_fsm: Disable FSM, note that if you disable FSM
|
|
then you should not use storage and events isolation
|
|
:param kwargs: Other arguments, will be passed as keyword arguments to handlers
|
|
"""
|
|
super(Dispatcher, self).__init__(name=name)
|
|
|
|
if storage and not isinstance(storage, BaseStorage):
|
|
raise TypeError(
|
|
f"FSM storage should be instance of 'BaseStorage' not {type(storage).__name__}"
|
|
)
|
|
|
|
# Telegram API provides originally only one event type - Update
|
|
# For making easily interactions with events here is registered handler which helps
|
|
# to separate Update to different event types like Message, CallbackQuery etc.
|
|
self.update = self.observers["update"] = TelegramEventObserver(
|
|
router=self, event_name="update"
|
|
)
|
|
self.update.register(self._listen_update)
|
|
|
|
# Error handlers should work is out of all other functions
|
|
# and should be registered before all others middlewares
|
|
self.update.outer_middleware(ErrorsMiddleware(self))
|
|
|
|
# User context middleware makes small optimization for all other builtin
|
|
# middlewares via caching the user and chat instances in the event context
|
|
self.update.outer_middleware(UserContextMiddleware())
|
|
|
|
# FSM middleware should always be registered after User context middleware
|
|
# because here is used context from previous step
|
|
self.fsm = FSMContextMiddleware(
|
|
storage=storage or MemoryStorage(),
|
|
strategy=fsm_strategy,
|
|
events_isolation=events_isolation or DisabledEventIsolation(),
|
|
)
|
|
if not disable_fsm:
|
|
# Note that when FSM middleware is disabled, the event isolation is also disabled
|
|
# Because the isolation mechanism is a part of the FSM
|
|
self.update.outer_middleware(self.fsm)
|
|
self.shutdown.register(self.fsm.close)
|
|
|
|
self.workflow_data: Dict[str, Any] = kwargs
|
|
self._running_lock = Lock()
|
|
self._stop_signal: Optional[Event] = None
|
|
self._stopped_signal: Optional[Event] = None
|
|
self._handle_update_tasks: Set[asyncio.Task[Any]] = set()
|
|
|
|
def __getitem__(self, item: str) -> Any:
|
|
return self.workflow_data[item]
|
|
|
|
def __setitem__(self, key: str, value: Any) -> None:
|
|
self.workflow_data[key] = value
|
|
|
|
def __delitem__(self, key: str) -> None:
|
|
del self.workflow_data[key]
|
|
|
|
def get(self, key: str, /, default: Optional[Any] = None) -> Optional[Any]:
|
|
return self.workflow_data.get(key, default)
|
|
|
|
@property
|
|
def storage(self) -> BaseStorage:
|
|
return self.fsm.storage
|
|
|
|
@property
|
|
def parent_router(self) -> Optional[Router]:
|
|
"""
|
|
Dispatcher has no parent router and can't be included to any other routers or dispatchers
|
|
|
|
:return:
|
|
"""
|
|
return None # noqa: RET501
|
|
|
|
@parent_router.setter
|
|
def parent_router(self, value: Router) -> None:
|
|
"""
|
|
Dispatcher is root Router then configuring parent router is not allowed
|
|
|
|
:param value:
|
|
:return:
|
|
"""
|
|
raise RuntimeError("Dispatcher can not be attached to another Router.")
|
|
|
|
async def feed_update(self, bot: Bot, update: Update, **kwargs: Any) -> Any:
|
|
"""
|
|
Main entry point for incoming updates
|
|
Response of this method can be used as Webhook response
|
|
|
|
:param bot:
|
|
:param update:
|
|
"""
|
|
loop = asyncio.get_running_loop()
|
|
handled = False
|
|
start_time = loop.time()
|
|
|
|
if update.bot != bot:
|
|
# Re-mounting update to the current bot instance for making possible to
|
|
# use it in shortcuts.
|
|
# Here is update is re-created because we need to propagate context to
|
|
# all nested objects and attributes of the Update, but it
|
|
# is impossible without roundtrip to JSON :(
|
|
# The preferred way is that pass already mounted Bot instance to this update
|
|
# before call feed_update method
|
|
update = Update.model_validate(update.model_dump(), context={"bot": bot})
|
|
|
|
try:
|
|
response = await self.update.wrap_outer_middleware(
|
|
self.update.trigger,
|
|
update,
|
|
{
|
|
**self.workflow_data,
|
|
**kwargs,
|
|
"bot": bot,
|
|
},
|
|
)
|
|
handled = response is not UNHANDLED
|
|
return response
|
|
finally:
|
|
finish_time = loop.time()
|
|
duration = (finish_time - start_time) * 1000
|
|
loggers.event.info(
|
|
"Update id=%s is %s. Duration %d ms by bot id=%d",
|
|
update.update_id,
|
|
"handled" if handled else "not handled",
|
|
duration,
|
|
bot.id,
|
|
)
|
|
|
|
async def feed_raw_update(self, bot: Bot, update: Dict[str, Any], **kwargs: Any) -> Any:
|
|
"""
|
|
Main entry point for incoming updates with automatic Dict->Update serializer
|
|
|
|
:param bot:
|
|
:param update:
|
|
:param kwargs:
|
|
"""
|
|
parsed_update = Update.model_validate(update, context={"bot": bot})
|
|
return await self._feed_webhook_update(bot=bot, update=parsed_update, **kwargs)
|
|
|
|
@classmethod
|
|
async def _listen_updates(
|
|
cls,
|
|
bot: Bot,
|
|
polling_timeout: int = 30,
|
|
backoff_config: BackoffConfig = DEFAULT_BACKOFF_CONFIG,
|
|
allowed_updates: Optional[List[str]] = None,
|
|
) -> AsyncGenerator[Update, None]:
|
|
"""
|
|
Endless updates reader with correctly handling any server-side or connection errors.
|
|
|
|
So you may not worry that the polling will stop working.
|
|
"""
|
|
backoff = Backoff(config=backoff_config)
|
|
get_updates = GetUpdates(timeout=polling_timeout, allowed_updates=allowed_updates)
|
|
kwargs = {}
|
|
if bot.session.timeout:
|
|
# Request timeout can be lower than session timeout and that's OK.
|
|
# To prevent false-positive TimeoutError we should wait longer than polling timeout
|
|
kwargs["request_timeout"] = int(bot.session.timeout + polling_timeout)
|
|
failed = False
|
|
while True:
|
|
try:
|
|
updates = await bot(get_updates, **kwargs)
|
|
except Exception as e:
|
|
failed = True
|
|
# In cases when Telegram Bot API was inaccessible don't need to stop polling
|
|
# process because some developers can't make auto-restarting of the script
|
|
loggers.dispatcher.error("Failed to fetch updates - %s: %s", type(e).__name__, e)
|
|
# And also backoff timeout is best practice to retry any network activity
|
|
loggers.dispatcher.warning(
|
|
"Sleep for %f seconds and try again... (tryings = %d, bot id = %d)",
|
|
backoff.next_delay,
|
|
backoff.counter,
|
|
bot.id,
|
|
)
|
|
await backoff.asleep()
|
|
continue
|
|
|
|
# In case when network connection was fixed let's reset the backoff
|
|
# to initial value and then process updates
|
|
if failed:
|
|
loggers.dispatcher.info(
|
|
"Connection established (tryings = %d, bot id = %d)",
|
|
backoff.counter,
|
|
bot.id,
|
|
)
|
|
backoff.reset()
|
|
failed = False
|
|
|
|
for update in updates:
|
|
yield update
|
|
# The getUpdates method returns the earliest 100 unconfirmed updates.
|
|
# To confirm an update, use the offset parameter when calling getUpdates
|
|
# All updates with update_id less than or equal to offset will be marked
|
|
# as confirmed on the server and will no longer be returned.
|
|
get_updates.offset = update.update_id + 1
|
|
|
|
async def _listen_update(self, update: Update, **kwargs: Any) -> Any:
|
|
"""
|
|
Main updates listener
|
|
|
|
Workflow:
|
|
- Detect content type and propagate to observers in current router
|
|
- If no one filter is pass - propagate update to child routers as Update
|
|
|
|
:param update:
|
|
:param kwargs:
|
|
:return:
|
|
"""
|
|
try:
|
|
update_type = update.event_type
|
|
event = update.event
|
|
except UpdateTypeLookupError as e:
|
|
warnings.warn(
|
|
"Detected unknown update type.\n"
|
|
"Seems like Telegram Bot API was updated and you have "
|
|
"installed not latest version of aiogram framework"
|
|
f"\nUpdate: {update.model_dump_json(exclude_unset=True)}",
|
|
RuntimeWarning,
|
|
)
|
|
raise SkipHandler() from e
|
|
|
|
kwargs.update(event_update=update)
|
|
|
|
return await self.propagate_event(update_type=update_type, event=event, **kwargs)
|
|
|
|
@classmethod
|
|
async def silent_call_request(cls, bot: Bot, result: TelegramMethod[Any]) -> None:
|
|
"""
|
|
Simulate answer into WebHook
|
|
|
|
:param bot:
|
|
:param result:
|
|
:return:
|
|
"""
|
|
try:
|
|
await bot(result)
|
|
except TelegramAPIError as e:
|
|
# In due to WebHook mechanism doesn't allow getting response for
|
|
# requests called in answer to WebHook request.
|
|
# Need to skip unsuccessful responses.
|
|
# For debugging here is added logging.
|
|
loggers.event.error("Failed to make answer: %s: %s", e.__class__.__name__, e)
|
|
|
|
async def _process_update(
|
|
self, bot: Bot, update: Update, call_answer: bool = True, **kwargs: Any
|
|
) -> bool:
|
|
"""
|
|
Propagate update to event listeners
|
|
|
|
:param bot: instance of Bot
|
|
:param update: instance of Update
|
|
:param call_answer: need to execute response as Telegram method (like answer into webhook)
|
|
:param kwargs: contextual data for middlewares, filters and handlers
|
|
:return: status
|
|
"""
|
|
try:
|
|
response = await self.feed_update(bot, update, **kwargs)
|
|
if call_answer and isinstance(response, TelegramMethod):
|
|
await self.silent_call_request(bot=bot, result=response)
|
|
return response is not UNHANDLED
|
|
|
|
except Exception as e:
|
|
loggers.event.exception(
|
|
"Cause exception while process update id=%d by bot id=%d\n%s: %s",
|
|
update.update_id,
|
|
bot.id,
|
|
e.__class__.__name__,
|
|
e,
|
|
)
|
|
return True # because update was processed but unsuccessful
|
|
|
|
async def _polling(
|
|
self,
|
|
bot: Bot,
|
|
polling_timeout: int = 30,
|
|
handle_as_tasks: bool = True,
|
|
backoff_config: BackoffConfig = DEFAULT_BACKOFF_CONFIG,
|
|
allowed_updates: Optional[List[str]] = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""
|
|
Internal polling process
|
|
|
|
:param bot:
|
|
:param kwargs:
|
|
:return:
|
|
"""
|
|
user: User = await bot.me()
|
|
loggers.dispatcher.info(
|
|
"Run polling for bot @%s id=%d - %r", user.username, bot.id, user.full_name
|
|
)
|
|
try:
|
|
async for update in self._listen_updates(
|
|
bot,
|
|
polling_timeout=polling_timeout,
|
|
backoff_config=backoff_config,
|
|
allowed_updates=allowed_updates,
|
|
):
|
|
handle_update = self._process_update(bot=bot, update=update, **kwargs)
|
|
if handle_as_tasks:
|
|
handle_update_task = asyncio.create_task(handle_update)
|
|
self._handle_update_tasks.add(handle_update_task)
|
|
handle_update_task.add_done_callback(self._handle_update_tasks.discard)
|
|
else:
|
|
await handle_update
|
|
finally:
|
|
loggers.dispatcher.info(
|
|
"Polling stopped for bot @%s id=%d - %r", user.username, bot.id, user.full_name
|
|
)
|
|
|
|
async def _feed_webhook_update(self, bot: Bot, update: Update, **kwargs: Any) -> Any:
|
|
"""
|
|
The same with `Dispatcher.process_update()` but returns real response instead of bool
|
|
"""
|
|
try:
|
|
return await self.feed_update(bot, update, **kwargs)
|
|
except Exception as e:
|
|
loggers.event.exception(
|
|
"Cause exception while process update id=%d by bot id=%d\n%s: %s",
|
|
update.update_id,
|
|
bot.id,
|
|
e.__class__.__name__,
|
|
e,
|
|
)
|
|
raise
|
|
|
|
async def feed_webhook_update(
|
|
self, bot: Bot, update: Union[Update, Dict[str, Any]], _timeout: float = 55, **kwargs: Any
|
|
) -> Optional[TelegramMethod[TelegramType]]:
|
|
if not isinstance(update, Update): # Allow to use raw updates
|
|
update = Update.model_validate(update, context={"bot": bot})
|
|
|
|
ctx = contextvars.copy_context()
|
|
loop = asyncio.get_running_loop()
|
|
waiter = loop.create_future()
|
|
|
|
def release_waiter(*_: Any) -> None:
|
|
if not waiter.done():
|
|
waiter.set_result(None)
|
|
|
|
timeout_handle = loop.call_later(_timeout, release_waiter)
|
|
|
|
process_updates: Future[Any] = asyncio.ensure_future(
|
|
self._feed_webhook_update(bot=bot, update=update, **kwargs)
|
|
)
|
|
process_updates.add_done_callback(release_waiter, context=ctx)
|
|
|
|
def process_response(task: Future[Any]) -> None:
|
|
warnings.warn(
|
|
"Detected slow response into webhook.\n"
|
|
"Telegram is waiting for response only first 60 seconds and then re-send update.\n"
|
|
"For preventing this situation response into webhook returned immediately "
|
|
"and handler is moved to background and still processing update.",
|
|
RuntimeWarning,
|
|
)
|
|
try:
|
|
result = task.result()
|
|
except Exception as e:
|
|
raise e
|
|
if isinstance(result, TelegramMethod):
|
|
asyncio.ensure_future(self.silent_call_request(bot=bot, result=result))
|
|
|
|
try:
|
|
try:
|
|
await waiter
|
|
except CancelledError: # pragma: no cover
|
|
process_updates.remove_done_callback(release_waiter)
|
|
process_updates.cancel()
|
|
raise
|
|
|
|
if process_updates.done():
|
|
# TODO: handle exceptions
|
|
response: Any = process_updates.result()
|
|
if isinstance(response, TelegramMethod):
|
|
return response
|
|
|
|
else:
|
|
process_updates.remove_done_callback(release_waiter)
|
|
process_updates.add_done_callback(process_response, context=ctx)
|
|
|
|
finally:
|
|
timeout_handle.cancel()
|
|
|
|
return None
|
|
|
|
async def stop_polling(self) -> None:
|
|
"""
|
|
Execute this method if you want to stop polling programmatically
|
|
|
|
:return:
|
|
"""
|
|
if not self._running_lock.locked():
|
|
raise RuntimeError("Polling is not started")
|
|
if not self._stop_signal or not self._stopped_signal:
|
|
return
|
|
self._stop_signal.set()
|
|
await self._stopped_signal.wait()
|
|
|
|
def _signal_stop_polling(self, sig: signal.Signals) -> None:
|
|
if not self._running_lock.locked():
|
|
return
|
|
|
|
loggers.dispatcher.warning("Received %s signal", sig.name)
|
|
if not self._stop_signal:
|
|
return
|
|
self._stop_signal.set()
|
|
|
|
async def start_polling(
|
|
self,
|
|
*bots: Bot,
|
|
polling_timeout: int = 10,
|
|
handle_as_tasks: bool = True,
|
|
backoff_config: BackoffConfig = DEFAULT_BACKOFF_CONFIG,
|
|
allowed_updates: Optional[Union[List[str], UNSET_TYPE]] = UNSET,
|
|
handle_signals: bool = True,
|
|
close_bot_session: bool = True,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""
|
|
Polling runner
|
|
|
|
:param bots: Bot instances (one or more)
|
|
:param polling_timeout: Long-polling wait time
|
|
:param handle_as_tasks: Run task for each event and no wait result
|
|
:param backoff_config: backoff-retry config
|
|
:param allowed_updates: List of the update types you want your bot to receive
|
|
By default, all used update types are enabled (resolved from handlers)
|
|
:param handle_signals: handle signals (SIGINT/SIGTERM)
|
|
:param close_bot_session: close bot sessions on shutdown
|
|
:param kwargs: contextual data
|
|
:return:
|
|
"""
|
|
if not bots:
|
|
raise ValueError("At least one bot instance is required to start polling")
|
|
if "bot" in kwargs:
|
|
raise ValueError(
|
|
"Keyword argument 'bot' is not acceptable, "
|
|
"the bot instance should be passed as positional argument"
|
|
)
|
|
|
|
async with self._running_lock: # Prevent to run this method twice at a once
|
|
if self._stop_signal is None:
|
|
self._stop_signal = Event()
|
|
if self._stopped_signal is None:
|
|
self._stopped_signal = Event()
|
|
|
|
if allowed_updates is UNSET:
|
|
allowed_updates = self.resolve_used_update_types()
|
|
|
|
self._stop_signal.clear()
|
|
self._stopped_signal.clear()
|
|
|
|
if handle_signals:
|
|
loop = asyncio.get_running_loop()
|
|
with suppress(NotImplementedError): # pragma: no cover
|
|
# Signals handling is not supported on Windows
|
|
# It also can't be covered on Windows
|
|
loop.add_signal_handler(
|
|
signal.SIGTERM, self._signal_stop_polling, signal.SIGTERM
|
|
)
|
|
loop.add_signal_handler(
|
|
signal.SIGINT, self._signal_stop_polling, signal.SIGINT
|
|
)
|
|
|
|
workflow_data = {
|
|
"dispatcher": self,
|
|
"bots": bots,
|
|
**self.workflow_data,
|
|
**kwargs,
|
|
}
|
|
if "bot" in workflow_data:
|
|
workflow_data.pop("bot")
|
|
|
|
await self.emit_startup(bot=bots[-1], **workflow_data)
|
|
loggers.dispatcher.info("Start polling")
|
|
try:
|
|
tasks: List[asyncio.Task[Any]] = [
|
|
asyncio.create_task(
|
|
self._polling(
|
|
bot=bot,
|
|
handle_as_tasks=handle_as_tasks,
|
|
polling_timeout=polling_timeout,
|
|
backoff_config=backoff_config,
|
|
allowed_updates=allowed_updates,
|
|
**workflow_data,
|
|
)
|
|
)
|
|
for bot in bots
|
|
]
|
|
tasks.append(asyncio.create_task(self._stop_signal.wait()))
|
|
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
|
|
|
|
for task in pending:
|
|
# (mostly) Graceful shutdown unfinished tasks
|
|
task.cancel()
|
|
with suppress(CancelledError):
|
|
await task
|
|
# Wait finished tasks to propagate unhandled exceptions
|
|
await asyncio.gather(*done)
|
|
|
|
finally:
|
|
loggers.dispatcher.info("Polling stopped")
|
|
try:
|
|
await self.emit_shutdown(bot=bots[-1], **workflow_data)
|
|
finally:
|
|
if close_bot_session:
|
|
await asyncio.gather(*(bot.session.close() for bot in bots))
|
|
self._stopped_signal.set()
|
|
|
|
def run_polling(
|
|
self,
|
|
*bots: Bot,
|
|
polling_timeout: int = 10,
|
|
handle_as_tasks: bool = True,
|
|
backoff_config: BackoffConfig = DEFAULT_BACKOFF_CONFIG,
|
|
allowed_updates: Optional[Union[List[str], UNSET_TYPE]] = UNSET,
|
|
handle_signals: bool = True,
|
|
close_bot_session: bool = True,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""
|
|
Run many bots with polling
|
|
|
|
:param bots: Bot instances (one or more)
|
|
:param polling_timeout: Long-polling wait time
|
|
:param handle_as_tasks: Run task for each event and no wait result
|
|
:param backoff_config: backoff-retry config
|
|
:param allowed_updates: List of the update types you want your bot to receive
|
|
:param handle_signals: handle signals (SIGINT/SIGTERM)
|
|
:param close_bot_session: close bot sessions on shutdown
|
|
:param kwargs: contextual data
|
|
:return:
|
|
"""
|
|
with suppress(KeyboardInterrupt):
|
|
return asyncio.run(
|
|
self.start_polling(
|
|
*bots,
|
|
**kwargs,
|
|
polling_timeout=polling_timeout,
|
|
handle_as_tasks=handle_as_tasks,
|
|
backoff_config=backoff_config,
|
|
allowed_updates=allowed_updates,
|
|
handle_signals=handle_signals,
|
|
close_bot_session=close_bot_session,
|
|
)
|
|
)
|