awesome-errors 0.3.0__py3-none-any.whl

awesome_errors/websocket/error_handler.py

@@ -0,0 +1,231 @@
+"""WebSocket error handler for awesome-errors"""
+
+import asyncio
+import json
+import logging
+from typing import Dict, Any, Type, Callable, Optional
+from fastapi import FastAPI, WebSocket
+
+from awesome_errors.core.error_codes import ErrorCode
+from awesome_errors.core.exceptions import AppError
+from awesome_errors.websocket.exceptions import WebSocketError, WebSocketInternalError
+from starlette.websockets import WebSocketState
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketErrorHandler:
+    """
+    Centralized error handler for WebSocket connections
+
+    Provides automatic error conversion, logging, and connection management
+    similar to Anaya's approach but for WebSocket.
+    """
+
+    def __init__(self):
+        self.error_mappings: Dict[Type[Exception], Dict[str, Any]] = {}
+        self._setup_default_mappings()
+
+    def _setup_default_mappings(self):
+        """Setup default error mappings"""
+        from .exceptions import (
+            WebSocketValidationError,
+            WebSocketInternalError,
+            JSONRPCErrorCode,
+        )
+
+        # Map common exceptions to WebSocket errors
+        self.error_mappings.update(
+            {
+                # Pydantic validation
+                ValueError: {
+                    "converter": lambda e: WebSocketValidationError(
+                        message=str(e), request_id=None
+                    )
+                },
+                # JSON parsing
+                json.JSONDecodeError: {
+                    "converter": lambda e: WebSocketError(
+                        code=ErrorCode.INVALID_FORMAT,
+                        message="Invalid JSON",
+                        ws_error_code=JSONRPCErrorCode.PARSE_ERROR,
+                        request_id=None,
+                    )
+                },
+                # Generic exceptions
+                Exception: {
+                    "converter": lambda e: WebSocketInternalError(
+                        message="Internal server error",
+                        original_error=str(e),
+                        request_id=None,
+                    )
+                },
+            }
+        )
+
+    def register_error_mapping(
+        self,
+        error_type: Type[Exception],
+        converter: Callable[[Exception], WebSocketError],
+    ):
+        """Register custom error mapping"""
+        self.error_mappings[error_type] = {"converter": converter}
+
+    async def handle_websocket_error(
+        self, websocket: WebSocket, error: Exception, request_id: Optional[str] = None
+    ) -> bool:
+        """
+        Handle WebSocket error
+
+        Returns:
+            bool: True if connection should be closed
+        """
+        try:
+            # Convert to WebSocketError
+            ws_error = self._convert_to_websocket_error(error, request_id)
+
+            # Log the error
+            self._log_error(ws_error)
+
+            # Send error response
+            await self._send_error_response(websocket, ws_error)
+
+            # Handle connection closing if needed
+            if ws_error.close_connection:
+                await self._close_connection(websocket, ws_error)
+                return True
+
+            return False
+
+        except Exception as handler_error:
+            logger.error(f"Error in error handler: {handler_error}")
+            # Try to close connection gracefully
+            try:
+                await websocket.close(code=1011, reason="Error handler failed")
+            except Exception:
+                pass
+            return True
+
+    def _convert_to_websocket_error(
+        self, error: Exception, request_id: Optional[str] = None
+    ) -> WebSocketError:
+        """Convert any exception to WebSocketError"""
+
+        # If already a WebSocketError, just update request_id
+        if isinstance(error, WebSocketError):
+            if request_id and not error.request_id:
+                error.request_id = request_id
+            return error
+
+        # If it's an AppError, convert it
+        if isinstance(error, AppError):
+            return WebSocketError.from_app_error(error, request_id)
+
+        # Check registered mappings
+        for error_type, mapping in self.error_mappings.items():
+            if isinstance(error, error_type):
+                ws_error = mapping["converter"](error)
+                if request_id and not ws_error.request_id:
+                    ws_error.request_id = request_id
+                return ws_error
+
+        # Fallback to internal error
+        return WebSocketInternalError(
+            message="Unexpected error occurred",
+            original_error=str(error),
+            request_id=request_id,
+        )
+
+    def _log_error(self, error: WebSocketError):
+        """Log error based on severity"""
+        if error.ws_error_code < -32600:  # Standard JSON-RPC errors
+            logger.error(
+                f"WebSocket error: {error.message}",
+                extra={
+                    "error_code": error.code.value,
+                    "ws_error_code": error.ws_error_code,
+                    "request_id": error.request_id,
+                },
+            )
+        else:  # Custom errors
+            logger.warning(
+                f"WebSocket error: {error.message}",
+                extra={
+                    "error_code": error.code.value,
+                    "ws_error_code": error.ws_error_code,
+                    "request_id": error.request_id,
+                },
+            )
+
+    async def _send_error_response(self, websocket: WebSocket, error: WebSocketError):
+        """Send error response via WebSocket"""
+        try:
+            if getattr(websocket, "client_state", None) not in (
+                WebSocketState.CONNECTED,
+            ):
+                return
+
+            error_response = json.dumps(error.to_jsonrpc_error())
+            await websocket.send_text(error_response)
+        except Exception as send_error:
+            logger.error(f"Failed to send error response: {send_error}")
+
+    async def _close_connection(self, websocket: WebSocket, error: WebSocketError):
+        """Close WebSocket connection with proper code and reason"""
+        try:
+            # Wait a bit to ensure error message is delivered
+            await asyncio.sleep(0.1)
+
+            # Map error codes to WebSocket close codes
+            close_code = self._get_close_code(error)
+            close_reason = error.message[:123]  # Max 123 bytes for close reason
+
+            await websocket.close(code=close_code, reason=close_reason)
+            if getattr(websocket, "client_state", None) not in (
+                WebSocketState.CONNECTED,
+            ):
+                await websocket.close(code=close_code, reason=close_reason)
+
+        except Exception as close_error:
+            logger.error(f"Failed to close WebSocket: {close_error}")
+
+    def _get_close_code(self, error: WebSocketError) -> int:
+        """Get appropriate WebSocket close code"""
+        from .exceptions import JSONRPCErrorCode
+
+        # Map error codes to WebSocket close codes
+        close_code_mapping = {
+            JSONRPCErrorCode.AUTH_REQUIRED: 1008,  # Policy Violation
+            JSONRPCErrorCode.TOKEN_EXPIRED: 1008,  # Policy Violation
+            JSONRPCErrorCode.RATE_LIMITED: 1008,  # Policy Violation
+            JSONRPCErrorCode.INTERNAL_ERROR: 1011,  # Internal Error
+            JSONRPCErrorCode.INVALID_REQUEST: 1003,  # Unsupported Data
+            JSONRPCErrorCode.PARSE_ERROR: 1007,  # Invalid frame payload data
+        }
+
+        return close_code_mapping.get(error.ws_error_code, 1000)  # Normal closure
+
+
+def setup_websocket_error_handling(app: FastAPI) -> WebSocketErrorHandler:
+    """
+    Setup WebSocket error handling for FastAPI app
+
+    Usage:
+        ws_error_handler = setup_websocket_error_handling(app)
+
+        # In WebSocket endpoint:
+        try:
+            # ... handle message
+        except Exception as e:
+            should_close = await ws_error_handler.handle_websocket_error(
+                websocket, e, request_id
+            )
+            if should_close:
+                break
+    """
+    handler = WebSocketErrorHandler()
+
+    # Store handler in app state for access
+    app.state.ws_error_handler = handler
+
+    return handler

awesome_errors/websocket/exceptions.py

@@ -0,0 +1,286 @@
+"""
+WebSocket exceptions extending awesome-errors for JSON-RPC 2.0 compatibility
+
+This module provides WebSocket-specific error classes that integrate with
+the existing awesome-errors system while adding JSON-RPC 2.0 support.
+"""
+
+from typing import Any, Dict, Optional
+
+from awesome_errors.core.error_codes import ErrorCode
+from awesome_errors.core.exceptions import AppError
+
+
+class JSONRPCErrorCode:
+    """Standard JSON-RPC 2.0 error codes"""
+
+    # Standard codes defined by JSON-RPC 2.0 specification
+    PARSE_ERROR = -32700  # Invalid JSON was received
+    INVALID_REQUEST = -32600  # The JSON sent is not a valid Request object
+    METHOD_NOT_FOUND = -32601  # The method does not exist / is not available
+    INVALID_PARAMS = -32602  # Invalid method parameter(s)
+    INTERNAL_ERROR = -32603  # Internal JSON-RPC error
+
+    # Server error codes (reserved range: -32000 to -32099)
+    # Custom codes for WebSocket-specific errors
+    AUTH_REQUIRED = -32000  # Authentication required
+    TOKEN_EXPIRED = -32001  # Authentication token expired
+    RATE_LIMITED = -32002  # Rate limit exceeded
+    SUBSCRIPTION_DENIED = -32003  # Subscription permission denied
+    CONNECTION_LIMIT = -32004  # Too many connections
+
+
+class WebSocketError(AppError):
+    """
+    WebSocket error extending AppError with JSON-RPC 2.0 compatibility
+
+    This class bridges awesome-errors with WebSocket/JSON-RPC requirements,
+    providing automatic error code mapping and connection management.
+    """
+
+    def __init__(
+        self,
+        code: ErrorCode,
+        message: str,
+        ws_error_code: Optional[int] = None,
+        request_id: Optional[str] = None,
+        details: Optional[Dict[str, Any]] = None,
+        close_connection: bool = False,
+    ):
+        """
+        Initialize WebSocket error
+
+        Args:
+            code: ErrorCode from awesome-errors
+            message: Human-readable error message
+            ws_error_code: JSON-RPC error code (auto-mapped if not provided)
+            request_id: JSON-RPC request ID for correlation
+            details: Additional error details
+            close_connection: Whether to close WebSocket connection after error
+        """
+        super().__init__(code, message, details)
+
+        # WebSocket specific fields
+        self.ws_error_code = ws_error_code or self._map_to_jsonrpc_code(code)
+        self.request_id = request_id
+        self.close_connection = close_connection
+
+    def _map_to_jsonrpc_code(self, error_code: ErrorCode) -> int:
+        """
+        Map ErrorCode to JSON-RPC error code
+
+        Provides sensible defaults for mapping awesome-errors codes
+        to JSON-RPC 2.0 standard codes.
+        """
+        mapping = {
+            # Authentication errors
+            ErrorCode.AUTH_REQUIRED: JSONRPCErrorCode.AUTH_REQUIRED,
+            ErrorCode.AUTH_PERMISSION_DENIED: JSONRPCErrorCode.AUTH_REQUIRED,
+            ErrorCode.AUTH_TOKEN_EXPIRED: JSONRPCErrorCode.TOKEN_EXPIRED,
+            ErrorCode.AUTH_INVALID_TOKEN: JSONRPCErrorCode.AUTH_REQUIRED,
+            # Validation errors
+            ErrorCode.VALIDATION_FAILED: JSONRPCErrorCode.INVALID_PARAMS,
+            ErrorCode.INVALID_INPUT: JSONRPCErrorCode.INVALID_PARAMS,
+            ErrorCode.MISSING_REQUIRED_FIELD: JSONRPCErrorCode.INVALID_PARAMS,
+            # Resource errors
+            ErrorCode.RESOURCE_NOT_FOUND: JSONRPCErrorCode.METHOD_NOT_FOUND,
+            ErrorCode.RESOURCE_ALREADY_EXISTS: JSONRPCErrorCode.INVALID_REQUEST,
+            # Rate limiting
+            ErrorCode.RATE_LIMIT_EXCEEDED: JSONRPCErrorCode.RATE_LIMITED,
+            # Database errors (all map to internal error)
+            ErrorCode.DB_CONNECTION_ERROR: JSONRPCErrorCode.INTERNAL_ERROR,
+            ErrorCode.DB_QUERY_ERROR: JSONRPCErrorCode.INTERNAL_ERROR,
+            ErrorCode.DB_CONSTRAINT_VIOLATION: JSONRPCErrorCode.INVALID_PARAMS,
+            # Business logic
+            ErrorCode.BUSINESS_RULE_VIOLATION: JSONRPCErrorCode.INVALID_REQUEST,
+            # Generic errors
+            ErrorCode.INTERNAL_ERROR: JSONRPCErrorCode.INTERNAL_ERROR,
+            ErrorCode.SERVICE_UNAVAILABLE: JSONRPCErrorCode.INTERNAL_ERROR,
+        }
+
+        return mapping.get(error_code, JSONRPCErrorCode.INTERNAL_ERROR)
+
+    def to_jsonrpc_error(self) -> Dict[str, Any]:
+        """
+        Convert to JSON-RPC 2.0 error format
+
+        Returns a dictionary conforming to JSON-RPC 2.0 error object specification.
+        """
+        error_data: Dict[str, Any] | None = None
+
+        # Build error data if we have details or metadata
+        if self.details or self.request_id:
+            error_data = {}
+
+            if self.details:
+                error_data["details"] = self.details
+
+            # Include original error code for clients that understand it
+            error_data["error_code"] = self.code.value
+
+            # Include timestamp for debugging
+            error_data["timestamp"] = self.timestamp.isoformat() + "Z"
+
+            # Include request ID if available
+            if self.request_id:
+                error_data["request_id"] = self.request_id
+
+        return {
+            "jsonrpc": "2.0",
+            "error": {
+                "code": self.ws_error_code,
+                "message": self.message,
+                "data": error_data,
+            },
+            "id": self.request_id,
+        }
+
+    @classmethod
+    def from_app_error(
+        cls,
+        error: AppError,
+        request_id: Optional[str] = None,
+        close_connection: bool = False,
+    ) -> "WebSocketError":
+        """
+        Create WebSocketError from existing AppError
+
+        This allows seamless conversion of any AppError to WebSocket format.
+        """
+        return cls(
+            code=error.code,
+            message=error.message,
+            details=error.details,
+            request_id=request_id,
+            close_connection=close_connection,
+        )
+
+
+class WebSocketAuthError(WebSocketError):
+    """Authentication error - always closes connection"""
+
+    def __init__(
+        self,
+        message: str = "Authentication required",
+        request_id: Optional[str] = None,
+        details: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(
+            code=ErrorCode.AUTH_REQUIRED,
+            message=message,
+            ws_error_code=JSONRPCErrorCode.AUTH_REQUIRED,
+            request_id=request_id,
+            details=details,
+            close_connection=True,  # Always close on auth failure
+        )
+
+
+class WebSocketTokenExpiredError(WebSocketError):
+    """Token expired error - requires re-authentication"""
+
+    def __init__(
+        self, request_id: Optional[str] = None, details: Optional[Dict[str, Any]] = None
+    ):
+        super().__init__(
+            code=ErrorCode.AUTH_TOKEN_EXPIRED,
+            message="Token expired",
+            ws_error_code=JSONRPCErrorCode.TOKEN_EXPIRED,
+            request_id=request_id,
+            details=details,
+            close_connection=True,  # Force reconnection with new token
+        )
+
+
+class WebSocketRateLimitError(WebSocketError):
+    """Rate limit exceeded error"""
+
+    def __init__(
+        self,
+        retry_after: int,
+        request_id: Optional[str] = None,
+        limit: Optional[int] = None,
+        window: Optional[int] = None,
+    ):
+        details: Dict[str, Any] = {
+            "retry_after": retry_after,
+            "limit": limit,
+            "window": window,
+        }
+
+        super().__init__(
+            code=ErrorCode.RATE_LIMIT_EXCEEDED,
+            message=f"Rate limit exceeded. Retry after {retry_after} seconds",
+            ws_error_code=JSONRPCErrorCode.RATE_LIMITED,
+            details={k: v for k, v in details.items() if v is not None},
+            request_id=request_id,
+            close_connection=False,  # Don't close, client can retry
+        )
+
+
+class WebSocketValidationError(WebSocketError):
+    """Validation error with field details"""
+
+    def __init__(
+        self,
+        message: str,
+        validation_errors: Optional[list[Any]] = None,
+        request_id: Optional[str] = None,
+    ):
+        details: Dict[str, Any] = {}
+        if validation_errors:
+            details["validation_errors"] = validation_errors
+
+        super().__init__(
+            code=ErrorCode.VALIDATION_FAILED,
+            message=message,
+            ws_error_code=JSONRPCErrorCode.INVALID_PARAMS,
+            details=details if details else None,
+            request_id=request_id,
+            close_connection=False,
+        )
+
+
+class WebSocketMethodNotFoundError(WebSocketError):
+    """Method not found error"""
+
+    def __init__(
+        self,
+        method: str,
+        request_id: Optional[str] = None,
+        available_methods: Optional[list[str]] = None,
+    ):
+        details: Dict[str, Any] = {"method": method}
+        if available_methods:
+            details["available_methods"] = available_methods
+
+        super().__init__(
+            code=ErrorCode.RESOURCE_NOT_FOUND,
+            message=f"Method not found: {method}",
+            ws_error_code=JSONRPCErrorCode.METHOD_NOT_FOUND,
+            details=details,
+            request_id=request_id,
+            close_connection=False,
+        )
+
+
+class WebSocketInternalError(WebSocketError):
+    """Internal server error"""
+
+    def __init__(
+        self,
+        message: str = "Internal server error",
+        request_id: Optional[str] = None,
+        original_error: Optional[str] = None,
+    ):
+        details: Dict[str, Any] = {}
+        if original_error:
+            details["original_error"] = str(original_error)
+
+        super().__init__(
+            code=ErrorCode.INTERNAL_ERROR,
+            message=message,
+            ws_error_code=JSONRPCErrorCode.INTERNAL_ERROR,
+            details=details if details else None,
+            request_id=request_id,
+            close_connection=False,
+        )