telegram-helper-bot/helper_bot/middlewares/metrics_middleware.py

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


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."""
        
        # Record basic event metrics
        if isinstance(event, Message):
            await self._record_message_metrics(event)
        elif isinstance(event, CallbackQuery):
            await self._record_callback_metrics(event)
        
        # Execute handler with timing
        start_time = time.time()
        try:
            result = await handler(event, data)
            duration = time.time() - start_time
            
            # Record successful execution
            handler_name = self._get_handler_name(handler)
            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 = self._get_handler_name(handler)
            metrics.record_method_duration(
                handler_name,
                duration,
                "handler",
                "error"
            )
            metrics.record_error(
                type(e).__name__,
                "handler",
                handler_name
            )
            raise
    
    def _get_handler_name(self, handler: Callable) -> str:
        """Extract handler name efficiently."""
        if hasattr(handler, '__name__'):
            return handler.__name__
        elif hasattr(handler, '__qualname__'):
            return handler.__qualname__
        return "unknown"
    
    async def _record_message_metrics(self, message: Message):
        """Record message metrics efficiently."""
        # 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"
        
        # Record message processing
        metrics.record_message(message_type, chat_type, "message_handler")
        
        # Record command if applicable
        if message.text and message.text.startswith('/'):
            command = message.text.split()[0][1:]  # Remove '/' and get command name
            user_type = "user" if message.from_user else "unknown"
            metrics.record_command(command, "message_handler", user_type)
    
    async def _record_callback_metrics(self, callback: CallbackQuery):
        """Record callback metrics efficiently."""
        metrics.record_message("callback_query", "callback", "callback_handler")
        
        if callback.data:
            parts = callback.data.split(':', 1)
            if parts:
                command = parts[0]
                user_type = "user" if callback.from_user else "unknown"
                metrics.record_command(command, "callback_handler", user_type)


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