ai4love-tools 0.0.6

为爱觉醒工具库

AI4Love 通用工具库

面向多项目复用的 Python 通用工具库，当前提供租户隔离模块，后续持续扩展更多模块。

模块概览

• tenant：租户隔离与上下文管理

租户模块功能特性

• 租户ID和用户ID提取：从 HTTP 请求头中自动提取租户ID和用户ID
• 上下文存储：使用 contextvars 实现线程安全的上下文变量存储
• 装饰器支持：提供 @extract_tenant_id 和 @with_tenant_context 装饰器
• SQL 自动封装：自动为 SQL 语句附加租户条件
• SQLAlchemy 集成：在不改动业务代码的前提下，自动为 ORM 查询追加租户过滤

安装

pip install -e .

快速开始

1. 使用装饰器自动提取租户ID（推荐）

from fastapi import FastAPI, Request
from ai4love_tools.tenant.decorators import extract_tenant_id
from ai4love_tools.tenant.context import get_tenant_id, get_user_id

app = FastAPI()

@app.get("/users")
@extract_tenant_id  # 自动从请求头提取租户ID和用户ID
async def get_users(request: Request):
    tenant_id = get_tenant_id()  # 从上下文获取租户ID
    user_id = get_user_id()  # 从上下文获取用户ID
    return {"tenant_id": tenant_id, "user_id": user_id}

2. 手动提取和设置

from ai4love_tools.tenant.extractor import extract_tenant_id_from_request, extract_user_id_from_request
from ai4love_tools.tenant.context import set_tenant_id, set_user_id

def manual_extract(request):
    tenant_id = extract_tenant_id_from_request(request)
    user_id = extract_user_id_from_request(request)
    set_tenant_id(tenant_id)
    set_user_id(user_id)

3. 多线程/后台任务上下文传递

from ai4love_tools.tenant.decorators import with_tenant_context
from ai4love_tools.tenant.context import get_tenant_id, get_user_id
from concurrent.futures import ThreadPoolExecutor

@with_tenant_context
def background_task(data):
    # 在新线程中，租户上下文已自动传递
    tenant_id = get_tenant_id()
    user_id = get_user_id()
    process_data(data)

executor = ThreadPoolExecutor()
executor.submit(background_task, data)

4. SQL 封装使用

from ai4love_tools.tenant.sql import tenant_execute
from ai4love_tools.tenant.context import set_tenant_id

# 设置租户ID到上下文
set_tenant_id("tenant_123")

def list_users(conn):
    # 业务侧只关心业务条件
    base_sql = "SELECT id, name FROM users WHERE status = :status"
    params = {"status": "active"}

    # 底层自动从上下文获取租户ID，并追加 tenant 条件
    def executor(sql: str, params: dict):
        with conn.cursor() as cursor:
            cursor.execute(sql, params)
            return cursor.fetchall()

    rows = tenant_execute(executor, base_sql, params)
    return rows

5. SQLAlchemy ORM 自动租户过滤

from sqlalchemy import create_engine
from ai4love_tools.tenant.sqlalchemy_integration import enable_sqlalchemy_tenant_isolation

# 在应用启动时启用
engine = create_engine("sqlite:///db.sqlite")
enable_sqlalchemy_tenant_isolation(engine)

# 业务代码无需改动，自动追加租户过滤
from sqlalchemy.orm import Session
from your_models import User

with Session(engine) as session:
    # 自动追加 tenant_id 条件
    users = session.query(User).filter(User.status == "active").all()

API 文档

上下文管理

• get_tenant_id(): 获取当前上下文中的租户ID
• get_user_id(): 获取当前上下文中的用户ID
• set_tenant_id(tenant_id): 设置租户ID到上下文
• set_user_id(user_id): 设置用户ID到上下文

装饰器

• @extract_tenant_id: 自动从请求头提取租户ID和用户ID并保存到上下文
• @with_tenant_context: 自动复制租户上下文到新线程/任务

SQL 封装

• tenant_execute(executor, sql, params, ...): 执行包含租户条件的 SQL 语句
• build_tenant_sql(sql, ...): 构建包含租户条件的 SQL 语句
• build_tenant_params(params, tenant_id, ...): 构建包含租户ID的参数

SQLAlchemy 集成

• enable_sqlalchemy_tenant_isolation(engine_or_session_factory, ...): 启用 SQLAlchemy ORM 自动租户过滤

依赖要求

• Python >= 3.12
• fastapi >= 0.115.0
• sqlalchemy >= 1.4.0

注意事项

• contextvars 无法跨进程传递，进程池场景需要显式传递租户ID作为参数
• 装饰器是给其他项目使用的，library 只负责提供装饰器函数，不负责具体的任务调度逻辑
• library 本身不包含业务逻辑，只提供工具函数和装饰器
• 尽量保持框架无关性，支持多种 Web 框架

运行测试

# 安装开发依赖
pip install -e ".[dev]"

# 运行所有测试
pytest

# 运行特定测试文件
pytest tests/test_context.py

# 运行测试并显示覆盖率
pytest --cov=src/ai4love_tools --cov-report=html

测试覆盖

测试覆盖以下模块：

• test_context.py - 上下文存储管理测试
• test_extractor.py - 租户ID和用户ID提取器测试
• test_decorators.py - 装饰器功能测试
• test_sql.py - SQL 封装函数测试
• test_sqlalchemy_integration.py - SQLAlchemy 集成测试
• test_integration.py - 集成测试（多模块协同工作）

许可证

MIT