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