fastapi-augment 0.1.0__py3-none-any.whl

fastapi_augment/db/sqlalchemy/session.py

@@ -0,0 +1,160 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/8/31
+@Description    : Async session factory with read/write splitting, built on top of EngineManager.
+"""
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import Any
+
+from sqlalchemy.ext.asyncio import (
+    AsyncSession,
+    async_sessionmaker,
+)
+
+from .engine import EngineManager
+
+
+class SessionFactory:
+    """
+    Async session factory that transparently routes reads and writes
+    to the correct database engine.
+
+    Usage::
+
+        manager = EngineManager(topology).start()
+        factory = SessionFactory(manager)
+
+        # Write session (always hits primary)
+        async with factory.write_session() as session:
+            session.add(User(name='alice'))
+            await session.commit()
+
+        # Read session (round-robin across replicas)
+        async with factory.read_session() as session:
+            result = await session.execute(select(User))
+
+        # FastAPI dependencies
+        @app.get('/users')
+        async def list_users(session: AsyncSession = Depends(factory.depends_read)):
+            ...
+    """
+
+    def __init__(
+        self,
+        engine_manager: EngineManager,
+        *,
+        expire_on_commit: bool = False,
+        session_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        self._manager = engine_manager
+        self._session_kwargs = session_kwargs or {}
+        self._expire_on_commit = expire_on_commit
+
+        self._write_factory = async_sessionmaker(
+            bind=engine_manager.write_engine,
+            class_=AsyncSession,
+            expire_on_commit=expire_on_commit,
+            **self._session_kwargs,
+        )
+
+        # 缓存读引擎对应的 session factory，避免每次 read_session 重复创建
+        self._read_factories: dict[str, async_sessionmaker[AsyncSession]] = {}
+
+    # ── Properties ───────────────────────────────────────────────────────
+
+    @property
+    def engine_manager(self) -> EngineManager:
+        """The underlying engine manager.
+
+        Returns:
+            The engine manager.
+        """
+
+        return self._manager
+
+    # ── Session Factories (async context managers) ───────────────────────
+
+    @asynccontextmanager
+    async def write_session(self) -> AsyncGenerator[AsyncSession, None]:
+        """Yield a session bound to the **primary** (write) engine
+
+        Returns:
+            An async session bound to the primary engine.
+        """
+        async with self._write_factory() as session:
+            yield session
+
+    @asynccontextmanager
+    async def read_session(self) -> AsyncGenerator[AsyncSession, None]:
+        """Yield a session bound to a **read** engine (round-robin across replicas).
+
+        Returns:
+            An async session bound to a read engine.
+        """
+        read_engine = self._manager.next_read_engine()
+        engine_key = id(read_engine)
+
+        factory = self._read_factories.get(engine_key)
+        if factory is None:
+            factory = async_sessionmaker(
+                bind=read_engine,
+                class_=AsyncSession,
+                expire_on_commit=self._expire_on_commit,
+                **self._session_kwargs,
+            )
+            self._read_factories[engine_key] = factory
+
+        async with factory() as session:
+            yield session
+
+    # ── FastAPI Dependencies ─────────────────────────────────────────────
+
+    async def depends_write(self) -> AsyncGenerator[AsyncSession, None]:
+        """FastAPI ``Depends()`` — inject a write session.
+
+        Returns:
+            An async session bound to the primary engine.
+        """
+        async with self.write_session() as session:
+            yield session
+
+    async def depends_read(self) -> AsyncGenerator[AsyncSession, None]:
+        """FastAPI ``Depends()`` — inject a read session.
+
+        Returns:
+            An async session bound to a read engine.
+        """
+        async with self.read_session() as session:
+            yield session
+
+    # ── Transactional helpers ────────────────────────────────────────────
+
+    @asynccontextmanager
+    async def transaction(self) -> AsyncGenerator[AsyncSession, None]:
+        """Write session with automatic commit / rollback.
+
+        Commits on clean exit, rolls back on any exception::
+
+            async with factory.transaction() as session:
+                session.add(obj)
+                # auto-commit here
+
+        Returns:
+            An async session bound to the primary engine.
+        """
+        async with self.write_session() as session:
+            try:
+                yield session
+                await session.commit()
+            except Exception:
+                await session.rollback()
+                raise
+
+    # ── Lifecycle ────────────────────────────────────────────────────────
+
+    async def dispose(self) -> None:
+        """Dispose the underlying engine manager and all connection pools."""
+        await self._manager.dispose()

fastapi_augment/factory.py

@@ -0,0 +1,238 @@
+"""
+@Author         : hangu
+@CreateDate     : 2026/9/3
+@Description    : FastAPI 应用工厂
+                  - 统一创建 FastAPI 实例并自动装配生命周期、数据库、中间件、路由
+                  - 支持可选的 SQLAlchemy 读写分离集成
+                  - 通过 app.state 暴露核心组件供业务层使用
+"""
+from __future__ import annotations
+
+from time import time
+from typing import Sequence, Callable, Any
+
+from fastapi import FastAPI, APIRouter
+from starlette.middleware import Middleware
+
+from .common.exception_handlers import register_exception_handlers
+from .lifespan import HookRegistry, fastapi_lifespan
+from .middlewares import RequestIdMiddleware
+from .openapi import configure_openapi_schema, OpenAPICustomConfig
+from .health import create_health_router
+
+# -------------------------- 类型别名 --------------------------
+# 路由注册回调：接收 app 实例，负责 include_router 等操作
+_RouteRegistrar = Callable[[FastAPI], None]
+
+
+def create_app(
+        *,
+        title: str = 'FastAPI',
+        version: str = '0.1.0',
+        description: str = '',
+        debug: bool = False,
+        docs_url: str | None = '/docs',
+        redoc_url: str | None = '/redoc',
+        openapi_url: str | None = '/openapi.json',
+        # CORS配置：None=不启用；传入非空序列才启用CORS中间件
+        cors_allow_origins: Sequence[str] | None = None,
+        cors_allow_methods: Sequence[str] | None = None,
+        cors_allow_headers: Sequence[str] | None = None,
+        # 生命周期钩子
+        registries: Sequence[HookRegistry] | None = None,
+        # 中间件
+        middlewares: Sequence[Middleware] | None = None,
+        # 路由注册
+        routers: Sequence[APIRouter | tuple[APIRouter, dict[str, Any]]] | None = None,
+        route_registrars: Sequence[_RouteRegistrar] | None = None,
+        # OpenAPI 自定义参数新增
+        openapi_remove_422: bool = True,
+        openapi_remove_validation_error: bool = True,
+        openapi_enable_bearer_auth: bool = False,
+        # 异常处理器
+        register_exceptions: bool = True,
+        # 数据库集成（可选）
+        engine_manager: Any | None = None,
+        session_factory: Any | None = None,
+        # 健康检查
+        health_check: bool = False,
+        # 额外 FastAPI 参数
+        **kwargs: Any,
+) -> FastAPI:
+    """创建并配置 FastAPI 应用实例。
+
+    工厂函数将以下组件统一装配到应用上：
+        1. 生命周期管理 — 自动接入 :func:`fastapi_lifespan`，合并用户注册表与 core_registry
+        2. 中间件 — 按列表顺序添加（先添加的在内层）
+        3. 路由 — 支持直接传入 APIRouter 或 (router, kwargs) 元组
+        4. 异常处理器 — 自动注册统一异常处理，返回标准 APIResponse 格式
+        5. 数据库 — 可选地将 EngineManager / SessionFactory 挂载到 app.state
+
+    装配完成后，``app.state`` 上可访问以下属性：
+        - ``app.state.registries``  — 生命周期注册表列表
+        - ``app.state.engine_manager`` — 数据库引擎管理器（若提供）
+        - ``app.state.session_factory``  — 会话工厂（若提供）
+
+    Example::
+
+        from fastapi_augment.factory import create_app
+        from fastapi_augment.lifespan import HookRegistry
+        from fastapi_augment.db.sqlalchemy import EngineManager, SessionFactory, ClusterTopology, NodeConfig
+
+        # 数据库
+        topology = ClusterTopology(primary=NodeConfig(url='sqlite+aiosqlite:///app.db'))
+        manager = EngineManager(topology).start()
+        sessions = SessionFactory(manager)
+
+        # 自定义钩子
+        my_registry = HookRegistry()
+
+        @my_registry.on_startup
+        async def init_cache() -> None:
+            ...
+
+        app = create_app(
+            title='My Service',
+            registries=[my_registry],
+            engine_manager=manager,
+            session_factory=sessions,
+        )
+
+    Args:
+        title: 应用标题
+        version: 应用版本
+        description: 应用描述
+        debug: 是否开启调试模式
+        docs_url: Swagger UI 路径，None 禁用
+        redoc_url: ReDoc 路径，None 禁用
+        openapi_url: OpenAPI schema 路径，None 禁用
+        cors_allow_origins: CORS 允许的源列表，None=不启用CORS
+        cors_allow_methods: CORS 允许的 HTTP 方法列表，None=不启用CORS
+        cors_allow_headers: CORS 允许的 HTTP 头列表，None=不启用CORS
+        registries: 生命周期钩子注册表
+        middlewares: Starlette 中间件列表
+        routers: 路由列表，元素可以是 APIRouter 或 (router, kwargs) 元组
+        route_registrars: 路由注册回调列表，接收 app 参数
+        openapi_remove_422: 是否移除 422 验证错误响应
+        openapi_remove_validation_error: 是否移除验证错误参数
+        openapi_enable_bearer_auth: 是否启用 Bearer 认证
+        register_exceptions: 是否自动注册统一异常处理器，默认 True
+        engine_manager: 数据库引擎管理器实例（可选）
+        session_factory: 会话工厂实例（可选）
+        health_check: 是否启用健康检查端点（默认 ``/health``）；
+            当传入 ``engine_manager`` 时自动包含数据库连通性检查
+        **kwargs: 传递给 FastAPI() 构造函数的额外参数（不允许传 lifespan）
+
+    Returns:
+        配置完成的 FastAPI 应用实例
+
+    Raises:
+        ValueError: 当 kwargs 中包含 lifespan 时抛出
+    """
+    # ---- 前置校验：禁止外部传入 lifespan ----
+    if 'lifespan' in kwargs:
+        raise ValueError(
+            '不允许通过 kwargs 传递 lifespan，'
+            '请使用 registries 参数注册生命周期钩子，工厂会自动管理 lifespan'
+        )
+
+    # ---- 1. 构建 FastAPI 实例 ----
+    app = FastAPI(
+        title=title,
+        version=version,
+        description=description,
+        debug=debug,
+        docs_url=docs_url,
+        redoc_url=redoc_url,
+        openapi_url=openapi_url,
+        middleware=middlewares,
+        **kwargs,
+    )
+
+    # ---- 2. 添加 RequestId 中间件 ----
+    app.add_middleware(RequestIdMiddleware)
+
+    # ---- 3. CORS 配置 ----
+    if cors_allow_origins:
+        from fastapi.middleware.cors import CORSMiddleware
+
+        opts = dict(
+            allow_origins=cors_allow_origins,
+            allow_credentials='*' not in cors_allow_origins,
+            allow_methods=cors_allow_methods or ['*'],
+            allow_headers=cors_allow_headers or ['*']
+        )
+        app.add_middleware(CORSMiddleware, **opts)  # type: ignore
+
+    # ---- 4. 生命周期注册表 ----
+    resolved_registries = _resolve_registries(registries)
+    app.state.registries = resolved_registries
+
+    # 设置 lifespan（使用模块级的 fastapi_lifespan，它会自动将 core_registry 插入首位）
+    app.router.lifespan_context = fastapi_lifespan
+
+    # ---- 5. 路由 ----
+    if routers:
+        for item in routers:
+            if isinstance(item, APIRouter):
+                app.include_router(item)
+            elif isinstance(item, tuple) and len(item) == 2:
+                router, router_kwargs = item
+                app.include_router(router, **router_kwargs)
+            else:
+                raise TypeError(
+                    f'routers 元素必须为 APIRouter 或 (APIRouter, dict) 元组，'
+                    f'实际为 {type(item).__name__}'
+                )
+
+    # ---- 6. 路由注册回调 ----
+    if route_registrars:
+        for registrar in route_registrars:
+            registrar(app)
+
+    # ---- 7. OpenAPI 配置 ----
+    configure_openapi_schema(
+        app,
+        config=OpenAPICustomConfig(
+            remove_422=openapi_remove_422,
+            remove_validation_error_schema=openapi_remove_validation_error,
+            enable_bearer_auth=openapi_enable_bearer_auth
+        )
+    )
+
+    # ---- 8. 异常处理器 ----
+    if register_exceptions:
+        register_exception_handlers(app)
+
+    # ---- 9. 数据库组件挂载到 app.state ----
+    if engine_manager is not None:
+        app.state.engine_manager = engine_manager
+
+    if session_factory is not None:
+        app.state.session_factory = session_factory
+
+    # ---- 10. 健康检查 ----
+    if health_check:
+        app.state.start_time = time()
+        app.include_router(create_health_router(include_db_check=engine_manager is not None))
+
+    return app
+
+
+# -------------------------- 内部辅助 --------------------------
+
+def _resolve_registries(
+        registries: Sequence[HookRegistry] | None,
+) -> list[HookRegistry]:
+    """将 registries 参数规范化为列表。
+
+    Args:
+        registries: 用户传入的注册表参数
+
+    Returns:
+        规范化后的 HookRegistry 列表
+    """
+    if registries is None:
+        return []
+
+    return list(registries)

fastapi_augment/health/__init__.py

@@ -0,0 +1,34 @@
+"""
+@Author     : zarkhan
+@CreateDate : 2026/9/6
+@Description: 健康检查模块
+              - 可扩展的检查器模式（继承 BaseChecker）
+              - 内置 AppChecker（应用状态）和 DatabaseChecker（数据库连通性）
+              - create_health_router() 一键创建健康检查路由
+"""
+from .checker import (
+    BaseChecker,
+    CheckResult,
+    HealthResponse,
+    STATUS_HEALTHY,
+    STATUS_DEGRADED,
+    STATUS_UNHEALTHY,
+)
+from .checkers import AppChecker, DatabaseChecker
+from .router import create_health_router
+
+__all__ = [
+    # 基础
+    'BaseChecker',
+    'CheckResult',
+    'HealthResponse',
+    # 状态常量
+    'STATUS_HEALTHY',
+    'STATUS_DEGRADED',
+    'STATUS_UNHEALTHY',
+    # 内置检查器
+    'AppChecker',
+    'DatabaseChecker',
+    # 路由工厂
+    'create_health_router',
+]

fastapi_augment/health/checker.py

@@ -0,0 +1,101 @@
+"""
+@Author     : zarkhan
+@CreateDate : 2026/9/6
+@Description: 健康检查基础模型与检查器抽象基类
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from pydantic import Field
+
+from ..schemas import SchemaBase
+
+
+# ── 结果模型 ──────────────────────────────────────────────────────────
+
+
+class CheckResult(SchemaBase):
+    """单项检查结果。
+
+    Attributes:
+        name: 检查项名称
+        status: 状态（healthy / degraded / unhealthy）
+        latency_ms: 检查耗时（毫秒）
+        details: 附加详情
+    """
+    name: str = Field(description='检查项名称')
+    status: str = Field(description='状态：healthy / degraded / unhealthy')
+    latency_ms: float = Field(default=0.0, description='检查耗时（毫秒）')
+    details: dict[str, Any] | None = Field(default=None, description='附加详情')
+
+
+class HealthResponse(SchemaBase):
+    """健康检查响应体。
+
+    Attributes:
+        status: 总体状态（取所有检查项中最差的状态）
+        checks: 各检查项结果列表
+    """
+    status: str = Field(description='总体状态：healthy / degraded / unhealthy')
+    checks: list[CheckResult] = Field(default_factory=list, description='各检查项结果')
+
+
+# ── 状态常量 ──────────────────────────────────────────────────────────
+
+STATUS_HEALTHY = 'healthy'
+STATUS_DEGRADED = 'degraded'
+STATUS_UNHEALTHY = 'unhealthy'
+
+# 状态严重程度排序，用于聚合时取最差状态
+_STATUS_SEVERITY = {STATUS_HEALTHY: 0, STATUS_DEGRADED: 1, STATUS_UNHEALTHY: 2}
+
+
+def _worst_status(*statuses: str) -> str:
+    """从多个状态中取最差的一个。
+
+    Args:
+        *statuses: 待比较的状态字符串
+
+    Returns:
+        最差状态
+    """
+    return max(statuses, key=lambda s: _STATUS_SEVERITY.get(s, 99))
+
+
+# ── 检查器基类 ────────────────────────────────────────────────────────
+
+
+class BaseChecker(ABC):
+    """健康检查器抽象基类。
+
+    所有检查器必须实现 ``name`` 属性和 ``check`` 异步方法::
+
+        class RedisChecker(BaseChecker):
+            @property
+            def name(self) -> str:
+                return 'redis'
+
+            async def check(self, app: FastAPI) -> CheckResult:
+                # 执行检查逻辑
+                return CheckResult(name=self.name, status='healthy', latency_ms=1.2)
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """检查项名称，用于标识和展示。"""
+        ...
+
+    @abstractmethod
+    async def check(self, app: Any) -> CheckResult:
+        """执行健康检查。
+
+        Args:
+            app: FastAPI 应用实例，可从中读取 app.state 等
+
+        Returns:
+            单项检查结果
+        """
+        ...

fastapi_augment/health/checkers.py

@@ -0,0 +1,109 @@
+"""
+@Author     : zarkhan
+@CreateDate : 2026/9/6
+@Description: 内置健康检查器
+              - AppChecker：应用基本信息（状态、版本、运行时长）
+              - DatabaseChecker：数据库连通性（基于 EngineManager）
+"""
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from fastapi import FastAPI
+from sqlalchemy import text
+
+from .checker import (
+    BaseChecker,
+    CheckResult,
+    STATUS_HEALTHY,
+    STATUS_UNHEALTHY,
+)
+
+
+class AppChecker(BaseChecker):
+    """应用基本健康检查。
+
+    返回应用运行状态、版本、运行时长。
+    通过 ``app.state.start_time`` 获取启动时间戳（由工厂或用户设置）。
+    """
+
+    @property
+    def name(self) -> str:
+        return 'app'
+
+    async def check(self, app: Any) -> CheckResult:
+        """检查应用状态。
+
+        Args:
+            app: FastAPI 应用实例
+
+        Returns:
+            包含版本、状态、运行时长的检查结果
+        """
+        start_time = getattr(app.state, 'start_time', None)
+        uptime = time.time() - start_time if start_time else 0.0
+
+        return CheckResult(
+            name=self.name,
+            status=STATUS_HEALTHY,
+            latency_ms=0.0,
+            details={
+                'status': 'running',
+                'version': app.version if isinstance(app, FastAPI) else '',
+                'uptime_seconds': round(uptime, 2),
+            },
+        )
+
+
+class DatabaseChecker(BaseChecker):
+    """数据库连通性健康检查。
+
+    通过 ``app.state.engine_manager`` 获取写引擎，执行 ``SELECT 1`` 验证连通性。
+    若未挂载 ``engine_manager``，则跳过检查并返回 ``unhealthy``。
+    """
+
+    def __init__(self, name: str = 'database') -> None:
+        self._name = name
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    async def check(self, app: Any) -> CheckResult:
+        """检查数据库连通性。
+
+        Args:
+            app: FastAPI 应用实例
+
+        Returns:
+            数据库连通性检查结果
+        """
+        engine_manager = getattr(app.state, 'engine_manager', None)
+        if engine_manager is None:
+            return CheckResult(
+                name=self.name,
+                status=STATUS_UNHEALTHY,
+                details={'error': 'engine_manager not found on app.state'},
+            )
+
+        start = time.perf_counter()
+        try:
+            engine = engine_manager.write_engine
+            async with engine.connect() as conn:
+                await conn.execute(text('SELECT 1'))
+            latency = (time.perf_counter() - start) * 1000
+
+            return CheckResult(
+                name=self.name,
+                status=STATUS_HEALTHY,
+                latency_ms=round(latency, 2),
+            )
+        except Exception as e:
+            latency = (time.perf_counter() - start) * 1000
+            return CheckResult(
+                name=self.name,
+                status=STATUS_UNHEALTHY,
+                latency_ms=round(latency, 2),
+                details={'error': str(e)},
+            )