amo-server/workers/middleware.py

99 lines
2.8 KiB
Python

"""
FastStream middleware for observability and error handling.
This module provides middleware for adding metrics, tracing, and
enhanced error handling to FastStream message processing.
"""
import logging
from typing import Any, Awaitable, Callable
from faststream import BaseMiddleware
from faststream.redis import RedisPublishCommand
from faststream.prometheus import PrometheusMiddleware
from faststream.observability.middleware import TelemetryMiddleware
logger = logging.getLogger(__name__)
class ErrorHandlingMiddleware(BaseMiddleware):
"""Middleware for enhanced error handling and logging."""
async def consume_scope(
self,
call_next: Callable[..., Awaitable[Any]],
msg: Any,
) -> Any:
"""
Handle message consumption with error logging.
Args:
call_next: Next middleware/handler in the chain
msg: Incoming message
Returns:
Result from the handler
"""
try:
logger.debug(f"Processing message: {type(msg).__name__}")
result = await call_next(msg)
logger.debug("Message processed successfully")
return result
except Exception as e:
logger.error(f"Error processing message: {str(e)}", exc_info=True)
# TODO: Implement dead letter queue publishing for permanent failures
# For now, re-raise to let FastStream handle retries
raise
class RedisPublishMiddleware(BaseMiddleware[RedisPublishCommand]):
"""Middleware for Redis publishing operations."""
async def publish_scope(
self,
call_next: Callable[[RedisPublishCommand], Awaitable[Any]],
cmd: RedisPublishCommand,
) -> Any:
"""
Handle Redis publish operations with logging.
Args:
call_next: Next middleware/handler in the chain
cmd: Redis publish command
Returns:
Result from the publish operation
"""
try:
logger.debug(f"Publishing to Redis: {cmd}")
result = await call_next(cmd)
logger.debug("Redis publish successful")
return result
except Exception as e:
logger.error(f"Redis publish failed: {str(e)}", exc_info=True)
raise
def setup_middleware(broker):
"""
Set up all middleware for the FastStream broker.
Args:
broker: FastStream Redis broker instance
"""
# Add Prometheus metrics middleware
broker.add_middleware(PrometheusMiddleware)
# Add OpenTelemetry tracing middleware
broker.add_middleware(TelemetryMiddleware)
# Add custom error handling middleware
broker.add_middleware(ErrorHandlingMiddleware)
# Add custom Redis publish middleware
broker.add_middleware(RedisPublishMiddleware)
logger.info("FastStream middleware configured successfully")