""" Metrics middleware for aiogram 3.x. Automatically collects metrics for message processing, command execution, and errors. """ from typing import Any, Awaitable, Callable, Dict from aiogram import BaseMiddleware from aiogram.types import TelegramObject, Message, CallbackQuery from aiogram.enums import ChatType import time from ..utils.metrics import metrics, track_middleware class MetricsMiddleware(BaseMiddleware): """Middleware for automatic metrics collection in aiogram handlers.""" async def __call__( self, handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]], event: TelegramObject, data: Dict[str, Any] ) -> Any: """Process event and collect metrics.""" async with track_middleware("metrics_middleware"): # Record message processing if isinstance(event, Message): await self._record_message_metrics(event, data) elif isinstance(event, CallbackQuery): await self._record_callback_metrics(event, data) # Execute handler and collect timing start_time = time.time() try: result = await handler(event, data) duration = time.time() - start_time # Record successful execution handler_name = handler.__name__ if hasattr(handler, '__name__') else "unknown" metrics.record_method_duration( handler_name, duration, "handler", "success" ) return result except Exception as e: duration = time.time() - start_time # Record error and timing handler_name = handler.__name__ if hasattr(handler, '__name__') else "unknown" metrics.record_method_duration( handler_name, duration, "handler", "error" ) metrics.record_error( type(e).__name__, "handler", handler_name ) raise async def _record_message_metrics(self, message: Message, data: Dict[str, Any]): """Record metrics for message processing.""" # Determine message type message_type = "text" if message.photo: message_type = "photo" elif message.video: message_type = "video" elif message.audio: message_type = "audio" elif message.document: message_type = "document" elif message.voice: message_type = "voice" elif message.sticker: message_type = "sticker" elif message.animation: message_type = "animation" # Determine chat type chat_type = "private" if message.chat.type == ChatType.GROUP: chat_type = "group" elif message.chat.type == ChatType.SUPERGROUP: chat_type = "supergroup" elif message.chat.type == ChatType.CHANNEL: chat_type = "channel" # Determine handler type handler_type = "unknown" if message.text and message.text.startswith('/'): handler_type = "command" # Record command specifically command = message.text.split()[0][1:] # Remove '/' and get command name metrics.record_command( command, "message_handler", "user" if message.from_user else "unknown" ) # Record message processing metrics.record_message(message_type, chat_type, handler_type) async def _record_callback_metrics(self, callback: CallbackQuery, data: Dict[str, Any]): """Record metrics for callback query processing.""" # Record callback processing metrics.record_message( "callback_query", "callback", "callback_handler" ) # Record callback command if available if callback.data: # Extract command from callback data (assuming format like "command:param") parts = callback.data.split(':', 1) if parts: command = parts[0] metrics.record_command( command, "callback_handler", "user" if callback.from_user else "unknown" ) class DatabaseMetricsMiddleware(BaseMiddleware): """Middleware for database operation metrics.""" async def __call__( self, handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]], event: TelegramObject, data: Dict[str, Any] ) -> Any: """Process event and collect database metrics.""" # Check if this handler involves database operations handler_name = handler.__name__ if hasattr(handler, '__name__') else "unknown" # You can add specific database operation detection logic here # For now, we'll just pass through and let individual decorators handle it return await handler(event, data) class ErrorMetricsMiddleware(BaseMiddleware): """Middleware for error tracking and metrics.""" async def __call__( self, handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]], event: TelegramObject, data: Dict[str, Any] ) -> Any: """Process event and collect error metrics.""" try: return await handler(event, data) except Exception as e: # Record error metrics handler_name = handler.__name__ if hasattr(handler, '__name__') else "unknown" metrics.record_error( type(e).__name__, "handler", handler_name ) raise