fastapi-augment 0.1.0__py3-none-any.whl

fastapi_augment/log/config.py

@@ -0,0 +1,201 @@
+"""
+@Author     : zarkhan
+@CreateDate : 2026/9/6
+@Description: 日志配置与管理——日志格式、级别、轮转、控制台/文件输出
+"""
+import logging
+import sys
+from pathlib import Path
+
+from .factory import install_request_id_factory
+from .filters import UvicornNameRewriteFilter
+from .handlers import (
+    MonthlyRotatingFileHandler,
+    MultiProcessTimedRotatingFileHandler,
+    YearlyRotatingFileHandler,
+)
+
+# -------------------------------------
+# 日志格式
+# -------------------------------------
+NORMAL_FORMAT = (
+    '%(asctime)s.%(msecs)03d | %(levelname)-8s | %(process)d:%(thread)d | '
+    '%(name)s | %(lineno)d | %(request_id)s | %(message)s'
+)
+_DATE_FORMAT = '%Y-%m-%d %H:%M:%S'
+
+# 轮转周期映射：语义字符串 → (when, interval)，month/year 使用自定义Handler
+_ROTATION_MAP: dict[str, tuple[str, int] | None] = {
+    'second': ('S', 1),
+    'minute': ('M', 1),
+    'hour': ('H', 1),
+    'day': ('midnight', 1),
+    'week': ('W0', 1),
+    'month': None,
+    'year': None,
+}
+
+# 需要接管的日志名称（清空默认处理器，走根日志）
+_TAKEOVER_LOGGERS = ('uvicorn', 'uvicorn.access', 'uvicorn.error', 'fastapi')
+
+
+def _takeover_uvicorn() -> None:
+    """强制接管uvicorn/fastapi日志：清空handler、propagate回根、附加名称重写过滤器"""
+    for logger_name in _TAKEOVER_LOGGERS:
+        logger = logging.getLogger(logger_name)
+        logger.handlers.clear()
+        logger.propagate = True
+        logger.setLevel(logging.INFO)
+
+        if logger_name in ('uvicorn.error', 'uvicorn.access'):
+            logger.addFilter(UvicornNameRewriteFilter())
+
+
+def _init_root_logger() -> None:
+    """初始化根日志：安装request_id工厂、设置格式、接管uvicorn
+
+    根日志级别设为WARNING，第三方库的DEBUG/INFO默认不输出。
+    uvicorn/fastapi已单独设为INFO，不受根日志影响。
+    """
+    install_request_id_factory()
+    logging.basicConfig(
+        level=logging.WARNING,
+        format=NORMAL_FORMAT,
+        datefmt=_DATE_FORMAT,
+        force=True,
+    )
+    _takeover_uvicorn()
+
+
+# 包导入时自动执行一次初始化
+_init_root_logger()
+
+
+# -------------------------------------
+# 对外API
+# -------------------------------------
+def set_log_level(level: str | int, logger_name: str | None = None) -> None:
+    """修改项目日志级别（不影响第三方库）
+
+    Args:
+        level: 支持传入 'debug' / 'info' / 'warn' / 'error' 或logging.DEBUG等
+        logger_name: 项目logger名称，不传则仅修改根日志级别
+    """
+    if isinstance(level, str):
+        level = level.upper()
+        level = getattr(logging, level, logging.INFO)
+
+    if logger_name:
+        logging.getLogger(logger_name).setLevel(level)
+    else:
+        logging.getLogger().setLevel(level)
+
+
+def set_log_format(log_format: str) -> None:
+    """修改全局日志格式，所有输出立即生效
+
+    Args:
+        log_format: 新的日志格式字符串
+    """
+    root_logger = logging.getLogger()
+    formatter = logging.Formatter(log_format)
+
+    for handler in root_logger.handlers:
+        handler.setFormatter(formatter)
+
+
+def setup_logger(
+    log_dir: str | Path | None = None,
+    filename: str = 'app.log',
+    rotation: str = 'day',
+    backup_count: int = 30,
+    encoding: str = 'utf-8',
+    enable_console: bool = True,
+) -> None:
+    """一键配置应用日志（支持控制台开关与文件日志）
+
+    支持以整点为节点的日志轮转，可选粒度：
+    - ``'second'`` — 每整秒
+    - ``'minute'`` — 每整分钟（:00秒）
+    - ``'hour'``   — 每整点小时（:00分）
+    - ``'day'``    — 每天00:00（默认）
+    - ``'week'``   — 每周一00:00
+    - ``'month'``  — 每月1日00:00
+    - ``'year'``   — 每年1月1日00:00
+
+    Args:
+        log_dir: 日志保存目录，若不提供则不开启文件日志
+        filename: 日志文件名
+        rotation: 轮转粒度，见上方说明，默认 ``'day'``
+        backup_count: 保留的历史日志文件数量
+        encoding: 文件编码
+        enable_console: 是否在控制台输出日志
+    """
+    root_logger = logging.getLogger()
+
+    # 1. 获取现有formatter
+    current_formatter = None
+    for h in root_logger.handlers:
+        if h.formatter:
+            current_formatter = h.formatter
+            break
+
+    if not current_formatter:
+        current_formatter = logging.Formatter(NORMAL_FORMAT, datefmt=_DATE_FORMAT)
+
+    # 2. 控制台输出
+    if enable_console:
+        has_console = any(
+            isinstance(h, logging.StreamHandler) and not isinstance(h, logging.FileHandler)
+            for h in root_logger.handlers
+        )
+        if not has_console:
+            console_handler = logging.StreamHandler(sys.stdout)
+            console_handler.setFormatter(current_formatter)
+            root_logger.addHandler(console_handler)
+    else:
+        handlers_to_remove = [
+            h for h in root_logger.handlers
+            if isinstance(h, logging.StreamHandler) and not isinstance(h, logging.FileHandler)
+        ]
+        for h in handlers_to_remove:
+            root_logger.removeHandler(h)
+
+    # 3. 文件输出
+    if log_dir:
+        log_dir = Path(log_dir)
+        log_dir.mkdir(parents=True, exist_ok=True)
+        file_path = log_dir / filename
+
+        has_file_handler = any(
+            isinstance(h, logging.FileHandler)
+            and getattr(h, 'baseFilename', '') == str(file_path.absolute())
+            for h in root_logger.handlers
+        )
+
+        if not has_file_handler:
+            rotation_key = rotation.lower()
+
+            if rotation_key not in _ROTATION_MAP:
+                raise ValueError(
+                    f'不支持的轮转粒度 {rotation!r}，'
+                    f'可选值：{list(_ROTATION_MAP.keys())}'
+                )
+
+            if rotation_key == 'month':
+                file_handler = MonthlyRotatingFileHandler(
+                    str(file_path), backup_count=backup_count, encoding=encoding
+                )
+            elif rotation_key == 'year':
+                file_handler = YearlyRotatingFileHandler(
+                    str(file_path), backup_count=backup_count, encoding=encoding
+                )
+            else:
+                when, interval = _ROTATION_MAP[rotation_key]
+                file_handler = MultiProcessTimedRotatingFileHandler(
+                    filename=str(file_path), when=when, interval=interval,
+                    backupCount=backup_count, encoding=encoding
+                )
+
+            file_handler.setFormatter(current_formatter)
+            root_logger.addHandler(file_handler)

fastapi_augment/log/factory.py

@@ -0,0 +1,32 @@
+"""
+@Author     : zarkhan
+@CreateDate : 2026/9/6
+@Description: 自定义日志记录工厂，注入request_id到每条日志
+"""
+import logging
+from typing import Any
+
+from ..middlewares.request_id import request_id_ctx_var
+
+# 保存原始工厂
+_old_factory = logging.getLogRecordFactory()
+
+
+def _record_factory(*args: Any, **kwargs: Any) -> logging.LogRecord:
+    """自定义日志记录工厂，在每条日志记录上注入当前请求的request_id
+
+    Args:
+        *args: 传递给原始工厂的位置参数
+        **kwargs: 传递给原始工厂的关键字参数
+
+    Returns:
+        注入了request_id属性的日志记录对象
+    """
+    record = _old_factory(*args, **kwargs)
+    record.request_id = request_id_ctx_var.get() or '-'
+    return record
+
+
+def install_request_id_factory() -> None:
+    """安装自定义日志记录工厂，使所有日志自动携带request_id"""
+    logging.setLogRecordFactory(_record_factory)

fastapi_augment/log/filters.py

@@ -0,0 +1,23 @@
+"""
+@Author     : zarkhan
+@CreateDate : 2026/9/6
+@Description: 日志过滤器，统一uvicorn日志名称
+"""
+import logging
+
+
+class UvicornNameRewriteFilter(logging.Filter):
+    """日志过滤器，统一uvicorn.error / uvicorn.access的名称为uvicorn"""
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        """将uvicorn.error / uvicorn.access的日志名称统一重写为uvicorn
+
+        Args:
+            record: 当前日志记录对象
+
+        Returns:
+            始终返回True，确保日志记录正常输出
+        """
+        if record.name in ('uvicorn.error', 'uvicorn.access'):
+            record.name = 'uvicorn'
+        return True

fastapi_augment/log/handlers.py

@@ -0,0 +1,81 @@
+"""
+@Author     : zarkhan
+@CreateDate : 2026/9/6
+@Description: 多进程安全的日志轮转处理器，支持秒/分/时/天/周及自定义月/年轮转
+"""
+from datetime import datetime
+from logging.handlers import TimedRotatingFileHandler
+
+
+class MultiProcessTimedRotatingFileHandler(TimedRotatingFileHandler):
+    """支持多进程安全的按时间轮转处理器"""
+
+    def doRollover(self) -> None:
+        """执行轮转，捕获PermissionError以兼容多进程场景"""
+        try:
+            super().doRollover()
+        except PermissionError:
+            # 没抢到锁，说明其他进程正在轮转，重新打开新的文件流即可
+            if self.stream:
+                self.stream.close()
+                self.stream = None
+            self.stream = self._open()
+
+
+class MonthlyRotatingFileHandler(MultiProcessTimedRotatingFileHandler):
+    """按月轮转的日志处理器，以每月1日00:00:00为节点"""
+
+    def __init__(self, filename: str, backup_count: int = 12, encoding: str = 'utf-8'):
+        """初始化按月轮转处理器
+
+        Args:
+            filename: 日志文件路径
+            backup_count: 保留的日志文件数量
+            encoding: 文件编码
+        """
+        super().__init__(filename, when='midnight', backupCount=backup_count, encoding=encoding)
+
+    def computeRollover(self, current_time: float) -> float:
+        """计算下次轮转时间（下月1日00:00:00）
+
+        Args:
+            current_time: 当前UNIX时间戳
+
+        Returns:
+            下次轮转的UNIX时间戳
+        """
+        dt = datetime.fromtimestamp(current_time)
+
+        if dt.month == 12:
+            next_month = datetime(dt.year + 1, 1, 1)
+        else:
+            next_month = datetime(dt.year, dt.month + 1, 1)
+
+        return next_month.timestamp()
+
+
+class YearlyRotatingFileHandler(MultiProcessTimedRotatingFileHandler):
+    """按年轮转的日志处理器，以每年1月1日00:00:00为节点"""
+
+    def __init__(self, filename: str, backup_count: int = 5, encoding: str = 'utf-8'):
+        """初始化按年轮转处理器
+
+        Args:
+            filename: 日志文件路径
+            backup_count: 保留的日志文件数量
+            encoding: 文件编码
+        """
+        super().__init__(filename, when='midnight', backupCount=backup_count, encoding=encoding)
+
+    def computeRollover(self, current_time: float) -> float:
+        """计算下次轮转时间（下一年1月1日00:00:00）
+
+        Args:
+            current_time: 当前UNIX时间戳
+
+        Returns:
+            下次轮转的UNIX时间戳
+        """
+        dt = datetime.fromtimestamp(current_time)
+        next_year = datetime(dt.year + 1, 1, 1)
+        return next_year.timestamp()

fastapi_augment/middlewares/__init__.py

@@ -0,0 +1,20 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/9/1
+@Description    : 中间件模块
+"""
+from .base import BaseASGIMiddleware
+from .request_id import (
+    get_request_id,
+    set_request_id,
+    reset_request_id,
+    RequestIdMiddleware
+)
+
+__all__ = [
+    'BaseASGIMiddleware',
+    'get_request_id',
+    'set_request_id',
+    'reset_request_id',
+    'RequestIdMiddleware'
+]

fastapi_augment/middlewares/base.py

@@ -0,0 +1,79 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/9/4
+@Description    : 通用ASGI中间件基类
+"""
+from contextvars import Token
+
+from starlette.types import ASGIApp, Scope, Receive, Send, Message
+
+
+class BaseASGIMiddleware:
+    """
+    原生ASGI中间件抽象基类
+    剥离重复ASGI样板代码，子类只实现钩子即可
+    兼容 http / websocket；http支持包装send修改响应头；finally保证清理
+    """
+
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def on_request(self, scope: Scope) -> Token | None:
+        """请求进入钩子：http/websocket都会调用
+
+        返回token，如果使用ContextVar，返回set得到的token；没有返回None
+
+        Args:
+            scope: ASGI scope对象
+
+        Returns:
+            ContextVar token或None
+        """
+        pass
+
+    async def wrap_send(self, message: Message) -> Message:  # noqa: no-self-use
+        """http.response.start消息钩子，可以修改headers等。
+
+        返回修改后的message对象。
+        仅HTTP模式生效；websocket不会进入此逻辑。
+        默认实现直接返回原message，子类按需覆盖。
+        （保留self参数，子类覆盖时需要访问实例属性）
+
+        Args:
+            message: ASGI message对象
+
+        Returns:
+            修改后的message对象
+        """
+        return message
+
+    async def on_finish(self, token: Token | None) -> None:
+        """请求结束finally钩子，用于reset ContextVar等清理工作
+
+        Args:
+            token: ContextVar token或None
+        """
+        pass
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        scope_type = scope['type']
+
+        if scope_type not in ('http', 'websocket'):
+            await self.app(scope, receive, send)
+            return
+
+        token: Token | None = await self.on_request(scope)
+
+        try:
+            if scope_type == 'http':
+                async def send_wrapper(message: Message) -> None:
+                    if message['type'] == 'http.response.start':
+                        message = await self.wrap_send(message)
+                    await send(message)
+
+                await self.app(scope, receive, send_wrapper)
+            else:
+                # websocket，不包装send，只执行上下文
+                await self.app(scope, receive, send)
+        finally:
+            await self.on_finish(token)

fastapi_augment/middlewares/request_id.py

@@ -0,0 +1,82 @@
+"""
+@Author     : zarkhan
+@CreateDate : 2026/7/4
+@Description: Request‑ID 追踪 ASGI 中间件与上下文工具
+"""
+from contextvars import ContextVar, Token
+from uuid import uuid4
+
+from starlette.types import Scope, Message
+
+from .base import BaseASGIMiddleware
+
+request_id_ctx_var: ContextVar[str | None] = ContextVar('request_id', default=None)
+
+
+def get_request_id() -> str | None:
+    """获取当前请求的Request ID
+
+    Returns:
+        Request ID or None
+    """
+    return request_id_ctx_var.get()
+
+
+def set_request_id(request_id: str) -> Token:
+    """设置当前请求的Request ID，返回token用于请求结束重置上下文
+
+    Args:
+        request_id: Request ID
+
+    Returns:
+        用于请求结束重置上下文的token
+    """
+    return request_id_ctx_var.set(request_id)
+
+
+def reset_request_id(token: Token) -> None:
+    """重置当前请求的Request ID
+
+    Args:
+        token: 用于重置上下文的token
+    """
+    request_id_ctx_var.reset(token)
+
+
+class RequestIdMiddleware(BaseASGIMiddleware):
+    """Request‑ID追踪中间件，继承通用ASGI基类"""
+
+    async def on_request(self, scope: Scope) -> Token | None:
+        # 1. 从请求头提取X‑Request‑Id
+        request_id = None
+
+        for name, value in scope.get('headers', []):
+            if name == b'x-request-id':
+                request_id = value.decode('latin-1')
+                break
+
+        if not request_id:
+            request_id = str(uuid4())
+
+        token = set_request_id(request_id)
+        return token
+
+    async def wrap_send(self, message: Message) -> Message:
+        """注入响应头 X‑Request‑Id"""
+        req_id = get_request_id()
+
+        if not req_id:
+            return message
+
+        headers = list(message.get('headers', []))
+
+        if not any(k == b'x-request-id' for k, _ in headers):
+            headers.append((b'x-request-id', req_id.encode('latin-1')))
+
+        message['headers'] = headers
+        return message
+
+    async def on_finish(self, token: Token | None) -> None:
+        """请求结束重置上下文，防止泄漏串请求"""
+        if token is not None:
+            reset_request_id(token)

fastapi_augment/openapi.py

@@ -0,0 +1,110 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/9/4
+@Description    : OpenAPI文档自定义配置，适配FastAPI factory工厂调用
+"""
+from logging import getLogger
+from typing import Any
+
+from fastapi import FastAPI
+from fastapi.openapi.utils import get_openapi
+
+_logger = getLogger(__name__)
+
+
+class OpenAPICustomConfig:
+    """OpenAPI自定义配置参数，方便工厂传入控制行为"""
+
+    def __init__(
+            self,
+            remove_422: bool = True,
+            remove_validation_error_schema: bool = True,
+            enable_bearer_auth: bool = False,
+            bearer_auth_name: str = 'BearerAuth'
+    ):
+        self.remove_422 = remove_422
+        self.remove_validation_error_schema = remove_validation_error_schema
+        self.enable_bearer_auth = enable_bearer_auth
+        self.bearer_auth_name = bearer_auth_name
+
+
+def configure_openapi_schema(
+        app: FastAPI,
+        config: OpenAPICustomConfig | None = None
+) -> None:
+    """配置OpenAPI Schema，在FastAPI factory(create_app)中调用
+
+    Args:
+        app: FastAPI实例对象
+        config: 自定义openapi配置，不传使用默认配置
+    """
+    cfg = config or OpenAPICustomConfig()
+
+    def custom_openapi() -> dict[str, Any] | None:
+        # OpenAPI被禁用场景，直接返回None
+        if app.openapi_url is None:
+            return None
+
+        # 已经生成过schema，直接复用缓存
+        if app.openapi_schema:
+            return app.openapi_schema
+
+        try:
+            # 兼容低版本fastapi，过滤不存在的参数
+            kwargs = dict(
+                title=app.title,
+                version=app.version,
+                openapi_version=app.openapi_version,
+                summary=app.summary,
+                description=app.description,
+                routes=app.routes,
+                tags=app.openapi_tags,
+                servers=app.servers,
+                terms_of_service=app.terms_of_service,
+                contact=app.contact,
+                license_info=app.license_info,
+            )
+            # separate_input_output_schemas 0.95+才存在
+            if hasattr(app, 'separate_input_output_schemas'):
+                kwargs['separate_input_output_schemas'] = app.separate_input_output_schemas
+
+            openapi_schema = get_openapi(**kwargs)
+        except Exception as exc:
+            # 生成openapi异常，不阻断服务启动
+            _logger.warning('Generate openapi schema failed: %s', exc)
+            return None
+
+        components = openapi_schema.setdefault('components', {})
+        schemas = components.setdefault('schemas', {})
+
+        # 移除校验错误模型
+        if cfg.remove_validation_error_schema:
+            schemas.pop('ValidationError', None)
+            schemas.pop('HTTPValidationError', None)
+
+        # 移除全部接口422响应
+        if cfg.remove_422:
+            paths = openapi_schema.get('paths', {})
+            for path_item in paths.values():
+                # http method: get/post/put/delete/patch/options/head/trace
+                for http_method in ('get', 'post', 'put', 'delete', 'patch', 'options', 'head', 'trace'):
+                    method_obj = path_item.get(http_method)
+                    if not isinstance(method_obj, dict):
+                        continue
+                    responses = method_obj.get('responses', {})
+                    responses.pop('422', None)
+
+        # 开启Bearer鉴权文档
+        if cfg.enable_bearer_auth:
+            security_schemes = components.setdefault('securitySchemes', {})
+            security_schemes[cfg.bearer_auth_name] = {
+                'type': 'http',
+                'scheme': 'bearer'
+            }
+            openapi_schema.setdefault('security', [{cfg.bearer_auth_name: []}])
+
+        app.openapi_schema = openapi_schema
+        return app.openapi_schema
+
+    # 替换openapi生成函数；去除 type: ignore，类型是FastAPI内部动态属性
+    app.openapi = custom_openapi

fastapi_augment/schemas/__init__.py

@@ -0,0 +1,29 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/9/1
+@Description    : 模型 schemas 定义
+"""
+from .base import SchemaBase, ORMSchemaBase
+from .pagination import PageData
+from .request import PageParams, TimeRangeParams, KeywordParams
+from .response import (
+    APIResponse,
+    response_success,
+    response_fail
+)
+
+__all__ = [
+    # base
+    'SchemaBase',
+    'ORMSchemaBase',
+    # pagination
+    'PageData',
+    # request params
+    'PageParams',
+    'TimeRangeParams',
+    'KeywordParams',
+    # response
+    'APIResponse',
+    'response_success',
+    'response_fail'
+]

fastapi_augment/schemas/base.py

@@ -0,0 +1,32 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/9/1
+@Description    : Schema全局基类
+"""
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class SchemaBase(BaseModel):
+    """
+    API对外输出基类：仅用于接口返回JSON，**不做ORM读取**
+    关闭 from_attributes，避免误用；保留驼峰、时间序列化
+    """
+    model_config = ConfigDict(
+        populate_by_name=True,
+        extra='ignore',
+        alias_generator=to_camel
+    )
+
+
+class ORMSchemaBase(SchemaBase):
+    """全局所有Pydantic Schema基类
+    统一配置、统一行为
+    """
+    model_config = ConfigDict(
+        **SchemaBase.model_config,  # 复制父类配置
+        from_attributes=True,
+        arbitrary_types_allowed=True,
+    )