toms-fast 0.2.1__py3-none-any.whl

tomskit/server/middleware.py

@@ -0,0 +1,371 @@
+"""
+FastAPI middleware for request ID tracking and resource cleanup.
+
+This module provides middleware to:
+1. Automatically handle X-Request-ID headers
+2. Clean up resources (database sessions, Redis connections, etc.) after request completion
+"""
+
+import asyncio
+import logging
+import uuid
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+from fastapi import Request
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from tomskit.logger import set_app_trace_id
+
+logger = logging.getLogger(__name__)
+
+
+class RequestIDMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to handle X-Request-ID header for request tracking.
+    
+    This middleware:
+    1. Reads X-Request-ID from request headers (if present)
+    2. Generates a new UUID if not present
+    3. Sets the request ID to the logging context via set_app_trace_id()
+    4. Adds X-Request-ID to response headers
+    
+    This enables distributed tracing and request correlation across services.
+    
+    Example:
+        from tomskit.server import FastApp, RequestIDMiddleware
+        
+        app = FastApp()
+        app.add_middleware(RequestIDMiddleware)
+    """
+    
+    def __init__(
+        self,
+        app: ASGIApp,
+        header_name: str = "X-Request-ID",
+        include_in_response: bool = True,
+        generate_on_missing: bool = True,
+    ):
+        """
+        Initialize RequestIDMiddleware.
+        
+        Args:
+            app: The ASGI application.
+            header_name: The header name to read/write request ID. Default: "X-Request-ID".
+            include_in_response: Whether to include request ID in response headers. Default: True.
+            generate_on_missing: Whether to generate a new ID if not present in request. Default: True.
+        """
+        super().__init__(app)
+        self.header_name = header_name
+        self.include_in_response = include_in_response
+        self.generate_on_missing = generate_on_missing
+    
+    async def dispatch(self, request: Request, call_next):
+        """
+        Process the request and set request ID in context.
+        
+        Args:
+            request: The incoming request.
+            call_next: The next middleware or route handler.
+            
+        Returns:
+            Response with X-Request-ID header (if enabled).
+        """
+        # 1. Try to get request ID from request headers
+        request_id = request.headers.get(self.header_name)
+        
+        # 2. Generate new ID if not present and generation is enabled
+        if not request_id and self.generate_on_missing:
+            request_id = str(uuid.uuid4())
+        
+        # 3. Set request ID to logging context (if we have one)
+        if request_id:
+            set_app_trace_id(request_id)
+        
+        # 4. Process the request
+        response = await call_next(request)
+        
+        # 5. Add request ID to response headers (if enabled and we have one)
+        if self.include_in_response and request_id:
+            response.headers[self.header_name] = request_id
+        
+        return response
+
+
+# ============================================================================
+# Resource Cleanup Middleware (Strategy Pattern)
+# ============================================================================
+
+class CleanupStrategy(ABC):
+    """
+    Abstract base class for resource cleanup strategies.
+    
+    Defines a unified interface for resource cleanup, supporting multiple resource types
+    (database, Redis, etc.). Each strategy is responsible for cleaning up a specific type of resource.
+    """
+    
+    @abstractmethod
+    async def cleanup(self, state: dict[str, Any]) -> None:
+        """
+        Execute resource cleanup.
+        
+        Args:
+            state: ASGI scope state dictionary containing resource identifiers created during the request
+        
+        Raises:
+            Exception: Any exceptions during cleanup should be caught and logged, not raised
+        """
+        pass
+    
+    @property
+    @abstractmethod
+    def resource_name(self) -> str:
+        """
+        Return resource name for logging purposes.
+        
+        Returns:
+            Resource name string, e.g., "database_session" or "redis_client"
+        """
+        pass
+
+
+class ResourceCleanupMiddleware:
+    """
+    Resource cleanup middleware.
+    
+    Uses strategy pattern to manage cleanup of multiple resources. Automatically cleans up
+    resources created during the request (such as database sessions, Redis connections, etc.)
+    after HTTP response completion to prevent resource leaks.
+    
+    How it works:
+    1. Intercepts ASGI send calls using a message queue
+    2. Listens for response completion event (http.response.body with more_body=False)
+    3. Executes all registered cleanup strategies after response completion
+    4. Ensures cleanup is executed even if application raises exceptions
+    
+    Features:
+    - Supports cleanup of multiple resource types (via strategy pattern)
+    - Automatic exception handling ensures cleanup logic doesn't affect responses
+    - Supports streaming responses (Server-Sent Events, large file downloads, etc.)
+    - Complete error handling and logging
+    
+    Example:
+        from tomskit.server import FastApp, ResourceCleanupMiddleware
+        from tomskit.sqlalchemy.database import DatabaseCleanupStrategy
+        from tomskit.redis.redis_pool import RedisCleanupStrategy
+        
+        app = FastApp()
+        
+        # 使用数据库清理策略
+        app.add_middleware(
+            ResourceCleanupMiddleware,
+            strategies=[DatabaseCleanupStrategy()]
+        )
+        
+        # 或使用多个策略
+        app.add_middleware(
+            ResourceCleanupMiddleware,
+            strategies=[
+                DatabaseCleanupStrategy(),
+                RedisCleanupStrategy(),
+            ]
+        )
+        
+        # If no strategies provided, no cleanup will be performed
+        app.add_middleware(ResourceCleanupMiddleware)
+    """
+    
+    def __init__(
+        self,
+        app: ASGIApp,
+        strategies: Optional[list[CleanupStrategy]] = None,
+        cleanup_timeout: float = 5.0,
+    ):
+        """
+        Initialize resource cleanup middleware.
+        
+        Args:
+            app: ASGI application instance
+            strategies: List of cleanup strategies. If None or empty, no cleanup will be performed
+            cleanup_timeout: Timeout for cleanup operations in seconds, default 5.0
+        """
+        self.app = app
+        self.cleanup_timeout = cleanup_timeout
+        
+        # If no strategies provided, use empty list (no cleanup will be performed)
+        if strategies is None:
+            self.strategies: list[CleanupStrategy] = []
+        else:
+            self.strategies = strategies
+    
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """
+        Handle ASGI request.
+        
+        Args:
+            scope: ASGI scope dictionary
+            receive: ASGI receive callable
+            send: ASGI send callable
+        """
+        # Only handle HTTP requests
+        if scope['type'] != 'http':
+            await self.app(scope, receive, send)
+            return
+        
+        # If no cleanup strategies, skip the queue mechanism for efficiency
+        if not self.strategies:
+            await self.app(scope, receive, send)
+            return
+        
+        # Get state (shared with request.state)
+        # In Starlette, scope['state'] is a State object, not a dict
+        # We should use it directly if it exists, otherwise create a dict
+        if 'state' not in scope:
+            scope['state'] = {}
+        state = scope['state']
+        
+        # Create message queue to intercept send calls
+        send_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        
+        # Track if cleanup has been executed to avoid duplicate execution
+        cleanup_executed = asyncio.Event()
+        
+        # Wrap send function to put messages into queue
+        async def send_wrapper(message: dict[str, Any]) -> None:  # type: ignore
+            await send_queue.put(message)
+        
+        # Sender coroutine: get messages from queue and send
+        async def sender() -> None:
+            """
+            Sender coroutine responsible for getting messages from queue and sending to client.
+            Triggers resource cleanup when response completion is detected.
+            """
+            response_completed = False
+            try:
+                while True:
+                    message = await send_queue.get()
+                    await send(message)
+                    
+                    # Detect response completion: http.response.body with more_body=False
+                    if (
+                        message['type'] == 'http.response.body' 
+                        and not message.get('more_body', False)
+                    ):
+                        response_completed = True
+                        break
+                        
+            except asyncio.CancelledError:
+                # If cancelled, still attempt to cleanup resources
+                response_completed = True
+            except Exception as e:
+                logger.error(f"Sender exception: {e}", exc_info=True)
+                response_completed = True
+            finally:
+                # Execute resource cleanup regardless of normal completion
+                if response_completed and not cleanup_executed.is_set():
+                    cleanup_executed.set()
+                    await self._execute_cleanup(state)
+        
+        # Application handler coroutine
+        async def app_handler() -> None:
+            """
+            Application handler coroutine that executes actual request processing.
+            """
+            try:
+                await self.app(scope, receive, send_wrapper)  # type: ignore[arg-type]
+            except Exception as e:
+                logger.error(f"Application handler exception: {e}", exc_info=True)
+                # Even if application raises exception, ensure response completion signal is sent
+                # so that sender can execute cleanup
+                if send_queue.empty():
+                    # If queue is empty, response has not been sent yet
+                    # Send an error response to ensure cleanup execution
+                    try:
+                        await send_wrapper({
+                            'type': 'http.response.start',
+                            'status': 500,
+                            'headers': [],
+                        })
+                        await send_wrapper({
+                            'type': 'http.response.body',
+                            'body': b'',
+                            'more_body': False,
+                        })
+                    except Exception:
+                        pass
+                raise
+        
+        # Run application handler and sender concurrently
+        # Use return_exceptions=True to prevent exceptions from stopping gather
+        results = await asyncio.gather(
+            app_handler(),
+            sender(),
+            return_exceptions=True,
+        )
+        
+        # If cleanup hasn't been executed yet (e.g., app_handler failed before sending anything),
+        # execute it now
+        if not cleanup_executed.is_set():
+            cleanup_executed.set()
+            await self._execute_cleanup(state)
+        
+        # Re-raise exceptions from app_handler if any
+        app_result = results[0]
+        if isinstance(app_result, Exception):
+            raise app_result
+    
+    async def _execute_cleanup(self, state: dict[str, Any]) -> None:
+        """
+        Execute all registered cleanup strategies.
+        
+        Args:
+            state: ASGI scope state dictionary
+        """
+        if not self.strategies:
+            return
+        
+        # Execute all cleanup strategies concurrently with timeout
+        cleanup_tasks = [
+            self._cleanup_with_timeout(strategy, state)
+            for strategy in self.strategies
+        ]
+        
+        results = await asyncio.gather(*cleanup_tasks, return_exceptions=True)
+        
+        # Check results and log
+        for strategy, result in zip(self.strategies, results):
+            if isinstance(result, Exception):
+                logger.error(
+                    f"Cleanup strategy '{strategy.resource_name}' failed: {result}",
+                    exc_info=result if isinstance(result, BaseException) else None,
+                )
+    
+    async def _cleanup_with_timeout(
+        self,
+        strategy: CleanupStrategy,
+        state: dict[str, Any],
+    ) -> None:
+        """
+        Execute a single cleanup strategy with timeout protection.
+        
+        Args:
+            strategy: Cleanup strategy instance
+            state: ASGI scope state dictionary
+        """
+        try:
+            await asyncio.wait_for(
+                strategy.cleanup(state),
+                timeout=self.cleanup_timeout,
+            )
+        except asyncio.TimeoutError:
+            logger.warning(
+                f"Cleanup strategy '{strategy.resource_name}' timed out "
+                f"(timeout: {self.cleanup_timeout}s)"
+            )
+        except Exception as e:
+            # Exceptions from cleanup strategies should not be raised
+            logger.error(
+                f"Cleanup strategy '{strategy.resource_name}' exception: {e}",
+                exc_info=True,
+            )

tomskit/server/parser.py

@@ -0,0 +1,312 @@
+from typing import (
+    Any, Dict, List, Literal, Mapping, Optional, Type, get_args, get_origin
+)
+
+from fastapi import HTTPException, Request, UploadFile
+from starlette.datastructures import FormData
+from pydantic import (
+    BaseModel, ConfigDict, Field, ValidationError, create_model, TypeAdapter
+)
+
+
+class ParserModel(BaseModel):
+    """
+    Base model that allows extra fields and dictionary‐style access.
+    """
+    model_config = ConfigDict(
+        extra="allow",
+        populate_by_name=True,
+        arbitrary_types_allowed=True # ← 允许任意类作为字段类型
+    )
+
+    def __getitem__(self, key: str) -> Any:
+        if key == "model_config":
+            key = "model_config_field"
+        return getattr(self, key, None)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        if key == "model_config":
+            key = "model_config_field"
+        setattr(self, key, value)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        if key == "model_config":
+            key = "model_config_field"
+        return getattr(self, key, default)
+
+    def __contains__(self, key: str) -> bool:
+        if key == "model_config":
+            key = "model_config_field"
+        return hasattr(self, key)
+    
+    def to_mapping(self) -> Mapping[str, Any]:
+        return self.model_dump(by_alias=True)
+
+class RequestParser:
+    """
+    A Flask-RESTful-like request parser built on FastAPI + Pydantic v2.
+    Supports:
+      - JSON, query, form, header, cookie, path parameters
+      - store/append/count actions
+      - required/default/nullable/choices semantics
+      - automatic TypeAdapter validation for built-ins, Annotated, custom types
+      - UploadFile and List[UploadFile]
+    """
+    # 类级缓存：parser_id -> 模型类
+    _model_cache: Dict[str, Type[ParserModel]] = {}
+
+    # 允许的 action
+    ALLOWED_ACTIONS = {"store", "append", "count"}
+
+    # 允许的 location
+    ALLOWED_LOCATIONS = {"args", "json", "form", "header", "cookie", "path"}
+
+    def __init__(self, parser_id: Optional[str] = None) -> None:
+        self._parser_id = parser_id
+        self._arg_defs: List[Dict[str, Any]] = []
+        self._model_cls: Optional[Type[ParserModel]] = None
+
+    def add_argument(
+        self,
+        name: str,
+        arg_type: Any,
+        required: bool = False,
+        default: Any = None,
+        choices: Optional[List[Any]] = None,
+        nullable: bool = False,
+        location: Literal["args", "json", "form", "header", "cookie", "path"] = "json",
+        action: Literal["store", "append", "count"] = "store",
+        help: Optional[str] = None,
+    ) -> None:
+        """
+        Register a new argument.
+        Raises ValueError on invalid configuration.
+        """
+        if not name or not isinstance(name, str):
+            raise ValueError("RequestParser: `name` must be a non-empty string")
+        if action not in self.ALLOWED_ACTIONS:
+            raise ValueError(f"RequestParser: `action` must be one of {self.ALLOWED_ACTIONS}")
+        if location not in self.ALLOWED_LOCATIONS:
+            raise ValueError(f"RequestParser: `location` must be one of {self.ALLOWED_LOCATIONS}")
+        if choices is not None and not isinstance(choices, list):
+            raise ValueError("RequestParser: `choices` must be a list if provided")
+        if required and default is not None:
+            raise ValueError("RequestParser: `default` must be None when `required=True`")
+        if not required and default is None and not nullable:
+            raise ValueError("RequestParser: `default=None` with `required=False` requires `nullable=True`")
+        if choices and default is not None and default not in choices:
+            raise ValueError("RequestParser: `default` must be one of `choices`")
+
+        # Validate that arg_type is supported by Pydantic
+        try:
+            adapter = TypeAdapter(arg_type)
+        except Exception as e:
+            raise ValueError(f"RequestParser: Invalid arg_type for '{name}': {e}")
+
+        # Validate default value with the same adapter
+        if default is not None:
+            try:
+                adapter.validate_python(default)
+            except ValidationError as e:
+                raise ValueError(f"RequestParser: Invalid default for '{name}': {e}")
+
+        # Validate each choice
+        for choice in (choices or []):
+            try:
+                adapter.validate_python(choice)
+            except ValidationError as e:
+                raise ValueError(f"RequestParser: Invalid choice {choice!r} for '{name}': {e}")
+
+        self._arg_defs.append({
+            "name": name,
+            "type": arg_type,
+            "required": required,
+            "default": default,
+            "choices": choices or [],
+            "nullable": nullable,
+            "location": location,
+            "action": action,
+            "help": help,
+        })
+
+    def _build_model(self) -> Type[ParserModel]:
+        """
+        Dynamically create a Pydantic model class based on registered arguments.
+        """
+
+        if self._parser_id is not None:
+            cached = RequestParser._model_cache.get(self._parser_id)
+            if cached is not None:
+                return cached
+
+        fields: Dict[str, Any] = {}
+
+        for d in self._arg_defs:
+            name = d["name"]
+            base_type = d["type"]
+            action = d["action"]
+            
+            # Determine annotation and default
+            ann: Any = None
+            if action == "append":
+                if get_origin(base_type) is list:
+                    item_type = get_args(base_type)[0]
+                else:
+                    item_type = base_type
+                ann = List[item_type]  # type: ignore
+                default = d["default"] if d["default"] is not None else []
+            elif action == "count":
+                ann = int
+                default = 0
+            else:  # store
+                ann = base_type
+                default = ... if d["required"] and d["default"] is None else d["default"]
+
+            # nullable → Optional
+            if d["nullable"]:
+                ann = Optional[ann]  # type: ignore
+
+            # choices → Literal
+            if d["choices"]:
+                # ann = Literal[tuple(d["choices"])]  # type: ignore
+                ann = Literal.__getitem__(tuple[Any, ...](d["choices"]))
+
+            internal_name = name
+            field_kwargs = {}
+            if name == "model_config":
+                internal_name = "model_config_field"
+                field_kwargs["alias"] = "model_config"
+
+            metadata: Dict[str, Any] = {}
+            if d["help"]:
+                metadata["description"] = d["help"]
+
+            pydantic_field_kwargs = {**metadata, **field_kwargs}
+
+            fields[internal_name] = (ann, Field(default, **pydantic_field_kwargs))
+
+        try:
+            cached_model = create_model("ParsedModel", __base__=ParserModel, **fields)
+        except Exception as e:
+            raise RuntimeError(f"RequestParser: Failed to create parser model: {e}")
+        if self._parser_id is not None:
+            RequestParser._model_cache[self._parser_id] = cached_model
+        return cached_model  # type: ignore
+
+    async def parse_args(self, request: Request) -> ParserModel:
+        """
+        Extract raw values from Request by location, merge per action,
+        then validate & coerce via the generated Pydantic model.
+        Raises HTTPException(422) on validation errors.
+        """
+        if self._model_cls is None:
+            self._model_cls = self._build_model()
+
+        # 1) 根据 Content-Type 决定是否解析 form
+        content_type = request.headers.get("content-type", "")
+        if content_type.startswith("multipart/form-data") or \
+           content_type.startswith("application/x-www-form-urlencoded"):
+            try:
+                form = await request.form()
+            except Exception:
+                form = FormData()
+        else:
+            form = FormData()
+
+        # 2) JSON 解析（仅在不是表单或 multipart 时才尝试）
+        try:
+            json_body = await request.json()
+            json_body = json_body if isinstance(json_body, dict) else {}
+        except Exception:
+            json_body = {}
+
+        # 3) 其余数据源
+        query = dict(request.query_params)
+        headers = dict(request.headers)
+        cookies = request.cookies
+        path_params = request.path_params
+
+        raw: Dict[str, Any] = {}
+
+        # 2) extract & merge
+        for d in self._arg_defs:
+            key = d["name"]
+            loc = d["location"]
+            action = d["action"]
+            value: Any = None
+            value = [] if action == "append" else 0 if action == "count" else None
+            present = False
+
+            if loc == "args" and key in query:
+                present = True
+                value = self._merge(value, query[key], action)
+
+            elif loc == "json" and key in json_body:
+                present = True
+                value = self._merge(value, json_body[key], action)
+
+            elif loc == "form":
+                # 多文件
+                if get_origin(d["type"]) is list and get_args(d["type"])[0] is UploadFile:
+                    fl = form.getlist(key)  # type: ignore[attr-defined]
+                    if fl:
+                        present = True
+                        for f in fl:
+                            value = self._merge(value, f, action)
+                # 单文件
+                elif d["type"] is UploadFile:
+                    fl = form.getlist(key)  # type: ignore[attr-defined]
+                    if fl:
+                        present = True
+                        # 如果想严格单文件，可检查 len(fl)>1
+                        value = self._merge(value, fl[0], action)
+                # 普通表单字段
+                elif key in form:
+                    present = True
+                    value = self._merge(value, form[key], action)
+
+            elif loc == "header" and key in headers:
+                present = True
+                value = self._merge(value, headers[key], action)
+
+            elif loc == "cookie" and key in cookies:
+                present = True
+                value = self._merge(value, cookies[key], action)
+
+            elif loc == "path" and key in path_params:
+                present = True
+                value = self._merge(value, path_params[key], action)
+
+            # 仅当请求里确实提供过这个字段时，才把它塞进 raw
+            if present:
+                raw[key] = value
+
+        # 3. 交给 Pydantic 去补默认 & 验证
+        try:
+            # Pydantic v2 用 model_validate
+            return self._model_cls.model_validate(raw)  # type: ignore
+        except ValidationError as e:
+            # 格式化错误，使运维快速定位字段与原因
+            error_messages = []
+            for err in e.errors():
+                loc = ".".join(str(x) for x in err.get("loc", []))
+                msg = err.get("msg", "")
+                err_type = err.get("type", "")
+                inp = err.get("input", None)
+                error_messages.append(
+                    f"RequestParser: Field '{loc}': {msg} (type={err_type}, input={inp!r})"
+                )
+            # 最终把格式化后的信息放到 detail 里
+            raise HTTPException(status_code=422, detail={"errors": error_messages})
+
+    @staticmethod
+    def _merge(prev: Any, new: Any, action: str) -> Any:
+        """Merge raw values according to action."""
+        if action == "store":
+            return new
+        if action == "append":
+            lst = prev if isinstance(prev, list) else []
+            lst.append(new)
+            return lst
+        # count
+        return (prev or 0) + 1