From db502059324ca633957728ce6b22a41916d036fc Mon Sep 17 00:00:00 2001 From: Nat-xu <19707027371@163.com> Date: Mon, 29 Jun 2026 23:20:56 +0800 Subject: [PATCH] =?UTF-8?q?=E8=BF=AD=E4=BB=A3=E5=8D=87=E7=BA=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 13 +++ Dockerfile | 13 +-- app/agent/context.py | 4 + app/agent/graph.py | 2 + app/agent/nodes/correct_sql.py | 49 ++++++--- app/agent/nodes/filter_metric.py | 37 +++++-- app/agent/nodes/filter_table.py | 37 +++++-- app/agent/nodes/recall_column.py | 26 ++++- app/agent/nodes/recall_metric.py | 26 ++++- app/agent/nodes/recall_value.py | 24 ++++- app/agent/nodes/run_sql.py | 34 ++++++ app/agent/nodes/sql_guard.py | 6 +- app/api/dependencies.py | 23 +++- app/api/routers/auth_router.py | 12 ++- app/api/routers/query_router.py | 13 ++- app/core/security.py | 43 ++++++-- app/core/sse_buffer.py | 22 +++- app/core/tenancy.py | 150 +++++++++++++++++++++++++++ app/services/query_service.py | 118 +++++++++++++++++---- docker/docker-compose.yaml | 15 +-- docs/OPTIMIZATION_PLAN.md | 4 +- frontend/src/components/StepRail.tsx | 2 +- frontend/src/lib/agentApi.ts | 48 +++++++-- main.py | 9 +- 24 files changed, 615 insertions(+), 115 deletions(-) create mode 100644 app/core/tenancy.py diff --git a/.env.example b/.env.example index 050b78d..e1ad893 100644 --- a/.env.example +++ b/.env.example @@ -89,3 +89,16 @@ ANALYST_USERNAME=analyst ANALYST_PASSWORD_HASH= VIEWER_USERNAME=viewer VIEWER_PASSWORD_HASH= + +# ---------- 多租户隔离（可选，默认单租户模式） ---------- +# 租户敏感表名集合，逗号分隔；留空则不启用多租户，所有行为与单租户一致 +# 启用后 run_sql 节点会用 sqlglot AST 强制向这些表注入 WHERE tenant_id = X 行级过滤 +# 示例：TENANT_SCOPED_TABLES=orders,order_items,users,addresses +TENANT_SCOPED_TABLES= +# 租户字段名，默认 tenant_id，仅当业务表使用其他字段名时需要覆盖 +TENANT_COLUMN=tenant_id +# 演示账号的租户标识（可选，未配置时回退到 default 单租户） +# 多租户场景下不同角色可归属不同租户，例如 admin 归属 tenant_a、analyst 归属 tenant_b +# ADMIN_TENANT_ID=tenant_a +# ANALYST_TENANT_ID=tenant_b +# VIEWER_TENANT_ID=default diff --git a/Dockerfile b/Dockerfile index 752bcbb..ae4d16d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -3,11 +3,12 @@ FROM python:3.14-slim AS builder RUN apt-get update && apt-get install -y --no-install-recommends \ - gcc g++ default-libmysqlclient-dev pkg-config \ - && rm -rf /var/lib/apt/lists/* + gcc g++ default-libmysqlclient-dev pkg-config \ + && rm -rf /var/lib/apt/lists/* # 引入 uv 做依赖管理，比 pip 更快且严格按 lockfile 安装 -COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv +# 固定版本保证构建可复现，避免 latest 升级破坏 uv sync 行为 +COPY --from=ghcr.io/astral-sh/uv:0.5.4 /uv /usr/local/bin/uv WORKDIR /app @@ -21,10 +22,10 @@ RUN uv sync --frozen --no-dev # 运行时镜像不含编译工具，体积更小，攻击面更窄 FROM python:3.14-slim -# asyncmy 运行时需要 MySQL 客户端库 +# asyncmy 运行时只需 MySQL 客户端运行时库，不装 dev 头文件以减小镜像体积和攻击面 RUN apt-get update && apt-get install -y --no-install-recommends \ - default-libmysqlclient-dev \ - && rm -rf /var/lib/apt/lists/* + libmariadb3 \ + && rm -rf /var/lib/apt/lists/* WORKDIR /app diff --git a/app/agent/context.py b/app/agent/context.py index 6593bef..0965f90 100644 --- a/app/agent/context.py +++ b/app/agent/context.py @@ -34,3 +34,7 @@ class DataAgentContext(TypedDict): meta_mysql_repository: MetaMySQLRepository # 数仓仓储抽象，具体实现由 dependencies 按配置注入（MySQL / ClickHouse / Postgres） dw_repository: DWRepository + # 当前请求的租户标识，由 query_service 从 JWT 解析后注入 + # run_sql 节点据此调用 inject_tenant_filter 在 SQL 执行前强制注入 WHERE tenant_id = X + # 未启用多租户（TENANT_SCOPED_TABLES 为空）时为 DEFAULT_TENANT_ID，inject_tenant_filter 原样返回 SQL + tenant_id: str diff --git a/app/agent/graph.py b/app/agent/graph.py index 22411e1..8e69654 100644 --- a/app/agent/graph.py +++ b/app/agent/graph.py @@ -146,6 +146,8 @@ async def test(): value_es_repository=value_es_repository, meta_mysql_repository=meta_mysql_repository, dw_repository=dw_repository, + # 本地调试使用默认租户；生产环境由 query_service 从 JWT 解析注入 + tenant_id="default", ) # stream_mode="custom" 会接收各节点通过 runtime.stream_writer 写出的进度信息 diff --git a/app/agent/nodes/correct_sql.py b/app/agent/nodes/correct_sql.py index e48aa53..9ec1411 100644 --- a/app/agent/nodes/correct_sql.py +++ b/app/agent/nodes/correct_sql.py @@ -15,6 +15,7 @@ from app.agent.state import DataAgentState from app.core.llm_tracing import get_tracing_config from app.core.log import logger +from app.core.resilience import CircuitOpenError, call_llm_with_resilience from app.prompt.prompt_loader import load_prompt @@ -53,23 +54,37 @@ async def correct_sql(state: DataAgentState, runtime: Runtime[DataAgentContext]) output_parser = StrOutputParser() chain = prompt | llm | output_parser - result = await chain.ainvoke( - { - # 与生成节点保持一致，用 YAML 向模型提供稳定 可读的结构化上下文 - "table_infos": yaml.dump( - table_infos, allow_unicode=True, sort_keys=False - ), - "metric_infos": yaml.dump( - metric_infos, allow_unicode=True, sort_keys=False - ), - "date_info": yaml.dump(date_info, allow_unicode=True, sort_keys=False), - "db_info": yaml.dump(db_info, allow_unicode=True, sort_keys=False), - "query": query, - "sql": sql, - "error": error, - }, - config=get_tracing_config(node="correct_sql"), - ) + # LLM 调用统一走熔断 + 重试包装：临时性失败自动重试，连续失败触发熔断时返回降级空 SQL + # 返回空 SQL 会让后续 validate_sql 自然失败，配合 correction_count 递增防止无限循环 + try: + result = await call_llm_with_resilience( + chain.ainvoke, + { + # 与生成节点保持一致，用 YAML 向模型提供稳定 可读的结构化上下文 + "table_infos": yaml.dump( + table_infos, allow_unicode=True, sort_keys=False + ), + "metric_infos": yaml.dump( + metric_infos, allow_unicode=True, sort_keys=False + ), + "date_info": yaml.dump(date_info, allow_unicode=True, sort_keys=False), + "db_info": yaml.dump(db_info, allow_unicode=True, sort_keys=False), + "query": query, + "sql": sql, + "error": error, + }, + config=get_tracing_config(node="correct_sql"), + ) + except CircuitOpenError: + logger.warning("LLM 熔断中，correct_sql 返回降级空 SQL") + writer( + { + "type": "error", + "code": "LLM_CIRCUIT_OPEN", + "message": "智能体服务暂时繁忙，请稍后重试", + } + ) + return {"sql": "", "correction_count": state.get("correction_count", 0) + 1} logger.info(f"校正后的SQL：{result}") writer({"type": "progress", "step": step, "status": "success"}) diff --git a/app/agent/nodes/filter_metric.py b/app/agent/nodes/filter_metric.py index 813d875..1e8cf97 100644 --- a/app/agent/nodes/filter_metric.py +++ b/app/agent/nodes/filter_metric.py @@ -15,6 +15,7 @@ from app.agent.state import DataAgentState, MetricInfoState from app.core.llm_tracing import get_tracing_config from app.core.log import logger +from app.core.resilience import CircuitOpenError, call_llm_with_resilience from app.prompt.prompt_loader import load_prompt @@ -39,15 +40,33 @@ async def filter_metric(state: DataAgentState, runtime: Runtime[DataAgentContext # LCEL 管道：填充提示词 -> 调用模型 -> 解析 JSON chain = prompt | llm | output_parser - result = await chain.ainvoke( - { - "query": query, - "metric_infos": yaml.dump( - metric_infos, allow_unicode=True, sort_keys=False - ), - }, - config=get_tracing_config(node="filter_metric"), - ) + # LLM 调用统一走熔断 + 重试包装；熔断时跳过过滤直接返回原始候选指标，保留完整口径供下游生成 + try: + result = await call_llm_with_resilience( + chain.ainvoke, + { + "query": query, + "metric_infos": yaml.dump( + metric_infos, allow_unicode=True, sort_keys=False + ), + }, + config=get_tracing_config(node="filter_metric"), + ) + except CircuitOpenError: + logger.warning("LLM 熔断中，filter_metric 跳过过滤，返回原始候选指标") + writer( + { + "type": "error", + "code": "LLM_CIRCUIT_OPEN", + "message": "智能体服务暂时繁忙，请稍后重试", + } + ) + writer({"type": "progress", "step": step, "status": "success"}) + return {"metric_infos": metric_infos} + # LLM 应返回指标名列表，校验后按空过滤处理避免类型错误 + if not isinstance(result, list): + logger.warning(f"filter_metric LLM 返回非 list 类型：{type(result).__name__}") + result = [] # 用模型返回的指标名称过滤原始结构，保留描述 依赖字段 别名等完整上下文 filtered_metric_infos = [ metric_info for metric_info in metric_infos if metric_info["name"] in result diff --git a/app/agent/nodes/filter_table.py b/app/agent/nodes/filter_table.py index a117aa6..c6d15af 100644 --- a/app/agent/nodes/filter_table.py +++ b/app/agent/nodes/filter_table.py @@ -15,6 +15,7 @@ from app.agent.state import DataAgentState, TableInfoState from app.core.llm_tracing import get_tracing_config from app.core.log import logger +from app.core.resilience import CircuitOpenError, call_llm_with_resilience from app.prompt.prompt_loader import load_prompt @@ -39,15 +40,33 @@ async def filter_table(state: DataAgentState, runtime: Runtime[DataAgentContext] # LCEL 管道：填充提示词 -> 调用模型 -> 解析 JSON chain = prompt | llm | output_parser - result = await chain.ainvoke( - { - "query": query, - "table_infos": yaml.dump( - table_infos, allow_unicode=True, sort_keys=False - ), - }, - config=get_tracing_config(node="filter_table"), - ) + # LLM 调用统一走熔断 + 重试包装；熔断时跳过过滤直接返回原始候选表，保留完整上下文供下游生成 + try: + result = await call_llm_with_resilience( + chain.ainvoke, + { + "query": query, + "table_infos": yaml.dump( + table_infos, allow_unicode=True, sort_keys=False + ), + }, + config=get_tracing_config(node="filter_table"), + ) + except CircuitOpenError: + logger.warning("LLM 熔断中，filter_table 跳过过滤，返回原始候选表") + writer( + { + "type": "error", + "code": "LLM_CIRCUIT_OPEN", + "message": "智能体服务暂时繁忙，请稍后重试", + } + ) + writer({"type": "progress", "step": step, "status": "success"}) + return {"table_infos": table_infos} + # LLM 可能返回非预期类型（如 list），校验后按空过滤处理避免 TypeError + if not isinstance(result, dict): + logger.warning(f"filter_table LLM 返回非 dict 类型：{type(result).__name__}") + result = {} # 模型只负责选择，程序根据选择结果从原始 TableInfoState 中裁剪，避免模型重写复杂结构出错 filtered_table_infos: list[TableInfoState] = [] for table_info in table_infos: diff --git a/app/agent/nodes/recall_column.py b/app/agent/nodes/recall_column.py index 685acbf..05066b0 100644 --- a/app/agent/nodes/recall_column.py +++ b/app/agent/nodes/recall_column.py @@ -15,6 +15,7 @@ from app.agent.state import DataAgentState from app.core.llm_tracing import get_tracing_config from app.core.log import logger +from app.core.resilience import CircuitOpenError, call_llm_with_resilience from app.entities.column_info import ColumnInfo from app.prompt.prompt_loader import load_prompt @@ -44,10 +45,27 @@ async def recall_column(state: DataAgentState, runtime: Runtime[DataAgentContext # LCEL 管道：填充提示词 -> 调用模型 -> 解析 JSON chain = prompt | llm | output_parser - result = await chain.ainvoke( - {"query": query}, - config=get_tracing_config(node="recall_column"), - ) + # LLM 调用统一走熔断 + 重试包装；熔断时跳过 LLM 扩展，仅用原始关键词做向量召回 + try: + result = await call_llm_with_resilience( + chain.ainvoke, + {"query": query}, + config=get_tracing_config(node="recall_column"), + ) + except CircuitOpenError: + logger.warning("LLM 熔断中，recall_column 跳过扩展词，仅用原始关键词召回") + writer( + { + "type": "error", + "code": "LLM_CIRCUIT_OPEN", + "message": "智能体服务暂时繁忙，请稍后重试", + } + ) + result = [] + # LLM 应返回扩展词列表，校验后取空列表避免 keywords + result 报 TypeError + if not isinstance(result, list): + logger.warning(f"recall_column LLM 返回非 list 类型：{type(result).__name__}") + result = [] # 原始关键词和 LLM 扩展词一起参与召回；set 去重，避免重复请求同一关键词 keywords_list = list(set(keywords + result)) diff --git a/app/agent/nodes/recall_metric.py b/app/agent/nodes/recall_metric.py index cb8e925..c217556 100644 --- a/app/agent/nodes/recall_metric.py +++ b/app/agent/nodes/recall_metric.py @@ -15,6 +15,7 @@ from app.agent.state import DataAgentState from app.core.llm_tracing import get_tracing_config from app.core.log import logger +from app.core.resilience import CircuitOpenError, call_llm_with_resilience from app.entities.metric_info import MetricInfo from app.prompt.prompt_loader import load_prompt @@ -44,10 +45,27 @@ async def recall_metric(state: DataAgentState, runtime: Runtime[DataAgentContext # LCEL 管道：填充提示词 -> 调用模型 -> 解析 JSON chain = prompt | llm | output_parser - result = await chain.ainvoke( - {"query": query}, - config=get_tracing_config(node="recall_metric"), - ) + # LLM 调用统一走熔断 + 重试包装；熔断时跳过 LLM 扩展，仅用原始关键词做向量召回 + try: + result = await call_llm_with_resilience( + chain.ainvoke, + {"query": query}, + config=get_tracing_config(node="recall_metric"), + ) + except CircuitOpenError: + logger.warning("LLM 熔断中，recall_metric 跳过扩展词，仅用原始关键词召回") + writer( + { + "type": "error", + "code": "LLM_CIRCUIT_OPEN", + "message": "智能体服务暂时繁忙，请稍后重试", + } + ) + result = [] + # LLM 应返回扩展词列表，校验后取空列表避免 keywords + result 报 TypeError + if not isinstance(result, list): + logger.warning(f"recall_metric LLM 返回非 list 类型：{type(result).__name__}") + result = [] # 通用关键词和指标扩展词都参与召回，提升同义指标的命中率 keywords_list = list(set(keywords + result)) diff --git a/app/agent/nodes/recall_value.py b/app/agent/nodes/recall_value.py index 60a63c6..66bc36e 100644 --- a/app/agent/nodes/recall_value.py +++ b/app/agent/nodes/recall_value.py @@ -13,7 +13,9 @@ from app.agent.context import DataAgentContext from app.agent.llm import llm from app.agent.state import DataAgentState +from app.core.llm_tracing import get_tracing_config from app.core.log import logger +from app.core.resilience import CircuitOpenError, call_llm_with_resilience from app.entities.value_info import ValueInfo from app.prompt.prompt_loader import load_prompt @@ -43,7 +45,27 @@ async def recall_value(state: DataAgentState, runtime: Runtime[DataAgentContext] # LCEL 管道：填充提示词 -> 调用模型 -> 解析 JSON chain = prompt | llm | output_parser - result = await chain.ainvoke({"query": query}) + # LLM 调用统一走熔断 + 重试包装；熔断时跳过 LLM 扩展，仅用原始关键词做 ES 检索 + try: + result = await call_llm_with_resilience( + chain.ainvoke, + {"query": query}, + config=get_tracing_config(node="recall_value"), + ) + except CircuitOpenError: + logger.warning("LLM 熔断中，recall_value 跳过扩展词，仅用原始关键词检索") + writer( + { + "type": "error", + "code": "LLM_CIRCUIT_OPEN", + "message": "智能体服务暂时繁忙，请稍后重试", + } + ) + result = [] + # LLM 应返回扩展词列表，校验后取空列表避免 keywords + result 报 TypeError + if not isinstance(result, list): + logger.warning(f"recall_value LLM 返回非 list 类型：{type(result).__name__}") + result = [] # 通用关键词和字段值扩展词一起检索 ES，尽量提高真实取值召回率 keywords_list = list(set(keywords + result)) diff --git a/app/agent/nodes/run_sql.py b/app/agent/nodes/run_sql.py index cdfe3c2..ec3f7fa 100644 --- a/app/agent/nodes/run_sql.py +++ b/app/agent/nodes/run_sql.py @@ -12,6 +12,12 @@ from app.agent.state import DataAgentState from app.core.log import logger from app.core.metrics import SQL_SAFETY_VIOLATIONS_TOTAL +from app.core.tenancy import ( + get_tenant_column, + get_tenant_scoped_tables, + inject_tenant_filter, + resolve_tenant_id, +) async def run_sql(state: DataAgentState, runtime: Runtime[DataAgentContext]): @@ -33,6 +39,34 @@ async def run_sql(state: DataAgentState, runtime: Runtime[DataAgentContext]): safe_sql = validate_readonly_sql( sql, allowed_tables, dialect=dw_repository.dialect ) + + # 多租户行级过滤：在 SQL 安全校验通过后、执行前，用 sqlglot AST 强制注入 + # WHERE tenant_id = X 条件，只对配置中声明的租户敏感表生效 + # 未配置 TENANT_SCOPED_TABLES 时 tenant_scoped_tables 为空集合，原样返回 SQL + tenant_id = resolve_tenant_id(runtime.context.get("tenant_id")) + tenant_scoped_tables = get_tenant_scoped_tables() + if tenant_scoped_tables: + try: + safe_sql = inject_tenant_filter( + safe_sql, + tenant_id, + tenant_scoped_tables, + tenant_column=get_tenant_column(), + dialect=dw_repository.dialect, + ) + except ValueError as e: + # AST 注入失败属于安全错误：拒绝执行，避免无租户过滤的 SQL 落库 + logger.error(f"租户过滤注入失败：{e}") + writer({"type": "progress", "step": step, "status": "error"}) + writer( + { + "type": "error", + "code": "TENANT_FILTER_FAILED", + "message": "SQL 租户隔离注入失败，请联系管理员", + } + ) + return + logger.info(f"通过安全校验的SQL：{safe_sql}") # 真实数据库访问统一封装在仓储层，节点只负责从状态取 SQL 并触发执行 diff --git a/app/agent/nodes/sql_guard.py b/app/agent/nodes/sql_guard.py index 8a22b67..2978c71 100644 --- a/app/agent/nodes/sql_guard.py +++ b/app/agent/nodes/sql_guard.py @@ -102,8 +102,12 @@ def validate_readonly_sql( ) # 3. 表名白名单校验：所有引用的表必须在候选表集合内 + # 注意 CTE 别名（WITH alias AS ...）会被 sqlglot 解析为 exp.Table 引用， + # 这些别名是查询内部定义的，不属于真实物理表，必须从白名单校验中排除 + cte_aliases = {cte.alias for cte in stmt.find_all(exp.CTE) if cte.alias} referenced_tables = {tbl.name for tbl in stmt.find_all(exp.Table)} - disallowed = referenced_tables - allowed_tables + real_referenced_tables = referenced_tables - cte_aliases + disallowed = real_referenced_tables - allowed_tables if disallowed: raise SQLSafetyViolation( "table not in allowlist", f"tables={','.join(sorted(disallowed))}" diff --git a/app/api/dependencies.py b/app/api/dependencies.py index eaa0e84..af88abe 100644 --- a/app/api/dependencies.py +++ b/app/api/dependencies.py @@ -10,7 +10,7 @@ from fastapi import Depends from langchain_huggingface import HuggingFaceEndpointEmbeddings -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from app.clients.embedding_client_manager import embedding_client_manager from app.clients.es_client_manager import es_client_manager @@ -72,6 +72,18 @@ async def get_dw_repository( return DWMySQLRepository(session) +def get_meta_session_factory() -> async_sessionmaker: + """返回元数据库 Session 工厂，供后台工作流任务创建独立 Session""" + + return meta_mysql_client_manager.session_factory + + +def get_dw_session_factory() -> async_sessionmaker: + """返回数仓库 Session 工厂，供后台工作流任务创建独立 Session""" + + return dw_mysql_client_manager.session_factory + + async def get_column_qdrant_repository() -> ColumnQdrantRepository: """创建字段向量检索仓储""" @@ -105,10 +117,17 @@ async def get_query_service( MetricQdrantRepository, Depends(get_metric_qdrant_repository) ], value_es_repository: Annotated[ValueESRepository, Depends(get_value_es_repository)], + meta_session_factory: Annotated[ + async_sessionmaker, Depends(get_meta_session_factory) + ], + dw_session_factory: Annotated[ + async_sessionmaker, Depends(get_dw_session_factory) + ], ) -> QueryService: """组装一次查询所需的业务服务""" # QueryService 只接收已经创建好的依赖对象，本身不关心这些对象来自 MySQL、Qdrant 还是 ES + # session_factory 供后台工作流任务创建独立 Session，与请求级 Session 解耦 return QueryService( meta_mysql_repository=meta_mysql_repository, embedding_client=embedding_client, @@ -116,4 +135,6 @@ async def get_query_service( column_qdrant_repository=column_qdrant_repository, metric_qdrant_repository=metric_qdrant_repository, value_es_repository=value_es_repository, + meta_session_factory=meta_session_factory, + dw_session_factory=dw_session_factory, ) diff --git a/app/api/routers/auth_router.py b/app/api/routers/auth_router.py index 12fcb2b..f7ac17f 100644 --- a/app/api/routers/auth_router.py +++ b/app/api/routers/auth_router.py @@ -37,7 +37,8 @@ async def login(request: LoginRequest) -> TokenResponse: detail={"code": "INVALID_CREDENTIALS", "message": "用户名或密码错误"}, ) - password_hash, role = record + # 解包三元组：密码哈希、角色、租户标识 + password_hash, role, tenant_id = record if not verify_password(request.password, password_hash): logger.warning(f"登录失败：用户名={request.username} 密码错误") raise HTTPException( @@ -45,9 +46,14 @@ async def login(request: LoginRequest) -> TokenResponse: detail={"code": "INVALID_CREDENTIALS", "message": "用户名或密码错误"}, ) - logger.info(f"登录成功：用户名={request.username} 角色={role.value}") + logger.info( + f"登录成功：用户名={request.username} 角色={role.value} 租户={tenant_id}" + ) + # 签发 JWT 时把 tenant_id 写入 claim，后续 SQL 执行时从 token 解析注入行级过滤 return TokenResponse( - access_token=create_access_token(request.username, role), + access_token=create_access_token( + request.username, role, tenant_id=tenant_id + ), token_type="bearer", expires_in=DEFAULT_EXPIRE_MINUTES * 60, role=role.value, diff --git a/app/api/routers/query_router.py b/app/api/routers/query_router.py index 8fe8d13..83a6b85 100644 --- a/app/api/routers/query_router.py +++ b/app/api/routers/query_router.py @@ -36,8 +36,9 @@ async def query_handler( query: QuerySchema, # 服务依赖：FastAPI 会调用 get_query_service，递归组装它所需的仓储和客户端 query_service: Annotated[QueryService, Depends(get_query_service)], - # 当前登录用户：JWT 校验 + RBAC 角色限制；当前用户仅用于日志关联，本接口不做差异化处理 - _current_user: Annotated[ + # 当前登录用户：JWT 校验 + RBAC 角色限制 + # current_user.tenant_id 由 JWT claim 解析，传入 query_service 用于 SQL 行级过滤 + current_user: Annotated[ User, Depends(require_roles(UserRole.ADMIN, UserRole.ANALYST, UserRole.VIEWER)), ], @@ -66,6 +67,11 @@ async def query_handler( detail=f"Invalid Last-Event-ID header: {last_event_id}", ) from e + # 校验缓冲区是否存在该请求，不存在（已过期或伪造）返回 404，避免无限轮询 DoS + if not await query_service.exists(original_request_id): + raise HTTPException( + status_code=404, detail="请求已过期或不存在，请重新发起查询" + ) return StreamingResponse( query_service.replay(original_request_id, last_seq), media_type="text/event-stream", @@ -76,8 +82,9 @@ async def query_handler( ) # 新请求：复用中间件生成的 request_id 作为事件缓冲区 key + # tenant_id 从 JWT claim 解析传入，工作流的 run_sql 节点据此强制注入行级过滤 request_id = str(request_id_ctx_var.get()) return StreamingResponse( - query_service.query(query.query, request_id), + query_service.query(query.query, request_id, current_user.tenant_id), media_type="text/event-stream", ) diff --git a/app/core/security.py b/app/core/security.py index cad5743..7841fdf 100644 --- a/app/core/security.py +++ b/app/core/security.py @@ -26,6 +26,7 @@ from passlib.context import CryptContext from app.core.log import logger +from app.core.tenancy import DEFAULT_TENANT_ID # 算法与默认过期时间：HS256 对称签名足够；token 1 小时过期，配合短期 refresh ALGORITHM = "HS256" @@ -46,6 +47,9 @@ class User: username: str role: UserRole + # 租户标识：从 JWT claim 解析，未启用多租户时为 DEFAULT_TENANT_ID + # 用于 SQL 行级过滤与按租户限流，由 app/core/tenancy.py 统一处理 + tenant_id: str = DEFAULT_TENANT_ID # bcrypt 是当前推荐的密码哈希算法，passlib 自动处理 salt 与版本升级 @@ -86,12 +90,15 @@ def create_access_token( username: str, role: UserRole, expires_minutes: int = DEFAULT_EXPIRE_MINUTES, + tenant_id: str = DEFAULT_TENANT_ID, ) -> str: - """签发 JWT，sub=用户名，role=角色，exp=过期时间""" + """签发 JWT，sub=用户名，role=角色，tenant_id=租户标识，exp=过期时间""" expire = datetime.now(timezone.utc) + timedelta(minutes=expires_minutes) payload = { "sub": username, "role": role.value, + # tenant_id claim 用于行级数据隔离；未启用多租户时为 DEFAULT_TENANT_ID + "tenant_id": tenant_id, "exp": expire, } return jwt.encode(payload, _get_jwt_secret(), algorithm=ALGORITHM) @@ -147,7 +154,11 @@ async def get_current_user( headers={"WWW-Authenticate": "Bearer"}, ) - return User(username=username, role=role) + # 解析 tenant_id claim，缺失或空值时回退到 DEFAULT_TENANT_ID + # 兼容未启用多租户时签发的旧 token（无 tenant_id 字段） + tenant_id = payload.get("tenant_id") or DEFAULT_TENANT_ID + + return User(username=username, role=role, tenant_id=tenant_id) def require_roles(*allowed_roles: UserRole): @@ -174,22 +185,36 @@ async def _checker( # 演示用账号：从环境变量读取，避免在源码中硬编码密码哈希 # 生产环境应替换为 meta.user 表查询，登录流程保持不变 -def get_demo_users() -> dict[str, tuple[str, UserRole]]: - """从环境变量读取演示账号，返回 {username: (password_hash, role)} +def get_demo_users() -> dict[str, tuple[str, UserRole, str]]: + """从环境变量读取演示账号，返回 {username: (password_hash, role, tenant_id)} 配置示例： ADMIN_USERNAME=admin ADMIN_PASSWORD_HASH=$2b$12$... + ADMIN_TENANT_ID=tenant_a ANALYST_USERNAME=analyst ANALYST_PASSWORD_HASH=$2b$12$... - VIEWER_USERNAME=viewer - VIEWER_PASSWORD_HASH=$2b$12$... + ANALYST_TENANT_ID=tenant_b + + tenant_id 可选，未配置时回退到 DEFAULT_TENANT_ID（单租户模式兼容） + 注意：若多个角色配置了相同的 username，先注册的角色生效，后续覆盖会被记录为 warning。 + 这避免了 VIEWER 静默覆盖 ADMIN 导致权限意外降级的安全风险。 """ - users: dict[str, tuple[str, UserRole]] = {} + users: dict[str, tuple[str, UserRole, str]] = {} for role in UserRole: username = os.getenv(f"{role.name}_USERNAME", "").strip() password_hash = os.getenv(f"{role.name}_PASSWORD_HASH", "").strip() - if username and password_hash: - users[username] = (password_hash, role) + if not (username and password_hash): + continue + if username in users: + existing_role = users[username][1] + logger.warning( + f"配置冲突：{role.value} 角色试图注册已存在的用户 {username}（当前角色 {existing_role.value}），" + "保留先注册的角色配置，忽略后续覆盖。请检查环境变量避免不同角色使用相同用户名。" + ) + continue + # 租户标识可选，未配置时使用 DEFAULT_TENANT_ID，保持单租户模式向后兼容 + tenant_id = os.getenv(f"{role.name}_TENANT_ID", "").strip() or DEFAULT_TENANT_ID + users[username] = (password_hash, role, tenant_id) return users diff --git a/app/core/sse_buffer.py b/app/core/sse_buffer.py index b9c4841..f2ed85f 100644 --- a/app/core/sse_buffer.py +++ b/app/core/sse_buffer.py @@ -42,13 +42,16 @@ async def append(self, request_id: str, event: dict) -> int: payload = json.dumps(event, ensure_ascii=False, default=str) # RPUSH 返回 push 后的 list 长度，作为新事件序号（从 1 开始） seq = await redis_client_manager.client.rpush(key, payload) - # 首次写入时设置 TTL，避免 key 永久驻留 - if seq == 1: - await redis_client_manager.client.expire(key, EVENT_TTL_SECONDS) - return seq except Exception as e: logger.warning("sse_buffer append failed request_id=%s: %s", request_id, e) return 0 + # 每次写入都刷新 TTL，避免长工作流（>EVENT_TTL_SECONDS）期间事件提前过期 + # 独立 try：expire 失败不应让已成功的 rpush 被当作失败丢弃 + try: + await redis_client_manager.client.expire(key, EVENT_TTL_SECONDS) + except Exception as e: + logger.warning("sse_buffer expire failed request_id=%s: %s", request_id, e) + return seq async def read_from( self, request_id: str, last_seq: int @@ -57,12 +60,15 @@ async def read_from( 采用轮询模式：当工作流仍在执行但暂时无新事件时，短暂 sleep 后重试。 工作流标记 done 后做最后一轮扫描，确保尾事件被消费。 + 兜底：设置最大轮询时长，即使 mark_complete 失败也不会永久挂起。 """ if redis_client_manager.client is None: return key = self._events_key(request_id) cursor = last_seq + 1 + # 兜底超时：略大于事件 TTL，避免 mark_complete 失败时无限轮询耗尽连接 + deadline = asyncio.get_running_loop().time() + EVENT_TTL_SECONDS + 1 while True: # seq 从 1 开始，Redis list 索引从 0 开始，seq=N 对应索引 N-1 try: @@ -101,6 +107,10 @@ async def read_from( yield cursor, event cursor += 1 return + # 兜底：mark_complete 失败时超时退出，避免永久轮询耗尽连接 + if asyncio.get_running_loop().time() > deadline: + logger.warning("sse_buffer read_from timeout request_id=%s", request_id) + return # 工作流仍在执行，短暂等待后再次轮询 await asyncio.sleep(POLL_INTERVAL_SECONDS) @@ -127,7 +137,9 @@ async def is_complete(self, request_id: str) -> bool: == "done" ) except Exception: - return True + # Redis 瞬时故障时返回 False，让 read_from 继续轮询等待恢复 + # 配合 read_from 的超时兜底，不会永久挂起 + return False async def exists(self, request_id: str) -> bool: """判断 request_id 是否在缓冲区中存在（用于重连校验）""" diff --git a/app/core/tenancy.py b/app/core/tenancy.py new file mode 100644 index 0000000..8c9c15e --- /dev/null +++ b/app/core/tenancy.py @@ -0,0 +1,150 @@ +""" +多租户隔离基础设施 + +提供租户级数据隔离的程序层防线，与 SQL 安全网关共同构成纵深防御： +1. JWT 携带 tenant_id claim，由 security.py 在登录时签发、get_current_user 解析 +2. SQL 执行前用 sqlglot AST 强制注入 WHERE tenant_id = X，只对配置中声明的租户敏感表生效 +3. 限流按 tenant_id 分桶（无 tenant 时回退到 IP），避免单租户刷爆全局配额 +4. Redis 缓冲区 key 加租户前缀，避免跨租户事件泄漏（在 sse_buffer 调用方拼接前缀） + +设计原则： +- 默认单租户兼容：未配置 TENANT_SCOPED_TABLES 时，所有行为与单租户一致，不注入任何过滤 +- 配置驱动：通过 .env 声明哪些表是租户敏感表，避免硬编码业务表名 +- 程序层强制：tenant_id 注入不依赖 prompt 约束，从 AST 层强制执行，无法被 SQL 注入绕过 +""" + +import os +from typing import Optional + +import sqlglot +from sqlglot import exp + +from app.core.log import logger + +# 默认租户标识：未启用多租户时所有用户归属此租户，行为与单租户完全一致 +DEFAULT_TENANT_ID = "default" + + +def get_tenant_scoped_tables() -> set[str]: + """从环境变量读取租户敏感表名集合 + + 配置示例： + TENANT_SCOPED_TABLES=orders,order_items,users,addresses + + 未配置时返回空集合，表示不启用多租户隔离，所有行为与单租户一致。 + 表名按小写规范化匹配，避免大小写差异导致过滤遗漏。 + """ + + raw = os.getenv("TENANT_SCOPED_TABLES", "").strip() + if not raw: + return set() + return {t.strip().lower() for t in raw.split(",") if t.strip()} + + +def get_tenant_column() -> str: + """租户字段名，默认 tenant_id，可通过 TENANT_COLUMN 环境变量自定义""" + + return os.getenv("TENANT_COLUMN", "tenant_id").strip() or "tenant_id" + + +def resolve_tenant_id(user_tenant_id: Optional[str]) -> str: + """解析最终生效的 tenant_id + + 优先级：用户 JWT claim > 默认租户 + 空值/None 都回退到 DEFAULT_TENANT_ID，保证 SQL 注入参数不为空。 + """ + + if user_tenant_id and user_tenant_id.strip(): + return user_tenant_id.strip() + return DEFAULT_TENANT_ID + + +def inject_tenant_filter( + sql: str, + tenant_id: str, + tenant_scoped_tables: set[str], + tenant_column: str = "tenant_id", + dialect: str = "mysql", +) -> str: + """用 sqlglot AST 强制向租户敏感表注入 WHERE tenant_id = X 条件 + + 实现细节： + - 解析 SQL 为 AST，遍历所有 exp.Table 节点 + - 对表名出现在 tenant_scoped_tables 集合的表，在外层 SELECT 注入 WHERE 条件 + - 多个租户敏感表时只注入一次（合并到外层 WHERE，避免重复 AND） + - 使用字面量参数化拼接，由 SQLAlchemy 参数化执行兜底防注入 + + Args: + sql: 已通过 validate_readonly_sql 安全校验的 SQL 文本 + tenant_id: 当前请求的租户标识（已通过 resolve_tenant_id 规范化） + tenant_scoped_tables: 租户敏感表名集合（小写） + tenant_column: 租户字段名，默认 tenant_id + dialect: SQL 方言，与仓储实例保持一致 + + Returns: + 注入租户过滤后的 SQL 文本；若未命中任何敏感表则原样返回 + + Raises: + ValueError: SQL 解析失败时抛出，调用方应作为安全错误处理 + """ + + if not tenant_scoped_tables: + return sql + + try: + parsed = sqlglot.parse(sql, dialect=dialect) + except Exception as e: + raise ValueError(f"tenant filter injection failed: SQL parse error: {e}") from e + + if not parsed: + return sql + + stmt = parsed[0] + if stmt is None: + return sql + + # 收集本次 SQL 中实际引用到的租户敏感表，用于决定是否注入过滤 + referenced_sensitive_tables: set[str] = set() + # CTE 别名（WITH alias AS ...）不是真实物理表，不需要也不能注入 tenant 过滤 + cte_aliases = {cte.alias for cte in stmt.find_all(exp.CTE) if cte.alias} + + for tbl in stmt.find_all(exp.Table): + table_name = tbl.name.lower() + if table_name in tenant_scoped_tables and tbl.name not in cte_aliases: + referenced_sensitive_tables.add(tbl.name) + + if not referenced_sensitive_tables: + return sql + + # 构造 WHERE tenant_id = '' 条件节点 + # 字面量字符串用 sqlglot exp.Literal.string 构建，渲染为 'tenant_id' 带引号 + # tenant_id 已在 resolve_tenant_id 校验非空，这里不再做 None 判断 + tenant_condition = exp.EQ( + this=exp.column(tenant_column), + expression=exp.Literal.string(tenant_id), + ) + + # 注入到最外层 SELECT 的 WHERE 子句 + # 与已有 WHERE 条件用 AND 合并，保证不破坏原 SQL 的业务条件 + if isinstance(stmt, exp.Select): + existing_where = stmt.args.get("where") + if existing_where is not None: + # 已有 WHERE 条件：用 AND 合并 tenant 过滤 + merged = exp.And(this=existing_where.this, expression=tenant_condition) + stmt.set("where", exp.Where(this=merged)) + else: + stmt.set("where", exp.Where(this=tenant_condition)) + else: + # 非 SELECT 语句已被 validate_readonly_sql 拦截，这里仅做防御性兜底 + logger.warning( + "inject_tenant_filter 收到非 SELECT 语句类型 %s，跳过注入", + type(stmt).__name__, + ) + return sql + + logger.info( + "注入租户过滤：tenant_id=%s 敏感表=%s", + tenant_id, + ",".join(sorted(referenced_sensitive_tables)), + ) + return stmt.sql(dialect=dialect) diff --git a/app/services/query_service.py b/app/services/query_service.py index f1f8fdd..ea89e32 100644 --- a/app/services/query_service.py +++ b/app/services/query_service.py @@ -14,20 +14,23 @@ import asyncio import json -from typing import AsyncIterator +from typing import Any, AsyncIterator from langchain_huggingface import HuggingFaceEndpointEmbeddings +from sqlalchemy.ext.asyncio import async_sessionmaker from app.agent.context import DataAgentContext from app.agent.graph import graph from app.agent.state import DataAgentState from app.clients.redis_client_manager import redis_client_manager from app.core.cache import get_cached_query_result, set_cached_query_result +from app.core.llm_tracing import get_tracing_config from app.core.log import logger from app.core.metrics import CACHE_HIT_TOTAL, CACHE_MISS_TOTAL from app.core.sse_buffer import sse_event_buffer from app.repositories.dw_repository import DWRepository from app.repositories.es.value_es_repository import ValueESRepository +from app.repositories.mysql.dw.dw_mysql_repository import DWMySQLRepository from app.repositories.mysql.meta.meta_mysql_repository import MetaMySQLRepository from app.repositories.qdrant.column_qdrant_repository import ColumnQdrantRepository from app.repositories.qdrant.metric_qdrant_repository import MetricQdrantRepository @@ -48,8 +51,12 @@ def __init__( column_qdrant_repository: ColumnQdrantRepository, metric_qdrant_repository: MetricQdrantRepository, value_es_repository: ValueESRepository, + # 后台工作流任务使用独立 Session，避免复用请求级 Session + # 客户端断连后 FastAPI 会关闭请求级 Session，若后台 task 仍引用会抛 ResourceClosedError + meta_session_factory: async_sessionmaker, + dw_session_factory: async_sessionmaker, ): - # MySQL 仓储负责元数据补全，DW 仓储抽象支持 MySQL/ClickHouse/Postgres + # 请求级仓储：供降级模式 _stream_direct 同步消费使用 self.meta_mysql_repository = meta_mysql_repository self.dw_repository = dw_repository @@ -59,31 +66,57 @@ def __init__( self.metric_qdrant_repository = metric_qdrant_repository self.value_es_repository = value_es_repository - async def query(self, query: str, request_id: str) -> AsyncIterator[str]: + # 独立 Session 工厂：后台工作流任务内部创建独立 Session + Repository， + # 与 HTTP 请求生命周期解耦，客户端断连不影响工作流继续执行 + self.meta_session_factory = meta_session_factory + self.dw_session_factory = dw_session_factory + + async def query( + self, query: str, request_id: str, tenant_id: str = "default" + ) -> AsyncIterator[str]: """执行一次问数工作流，并逐段产出 SSE 消息 Redis 可用时走缓冲模式：工作流作为后台任务执行，事件写入 Redis 缓冲区， HTTP 响应从缓冲区读取并附加 `id: {request_id}:{seq}` 字段，供客户端断线重连。 Redis 不可用时降级为直接流式输出：事件即时 yield，不支持重连但保证主链路可用。 + + tenant_id 由路由层从 JWT 解析传入，注入到 DataAgentContext 供 run_sql 节点 + 在 SQL 执行前调用 inject_tenant_filter 强制注入 WHERE tenant_id = X 行级过滤。 """ if redis_client_manager.client is None: # 降级模式：Redis 不可用，直接流式输出 - async for sse in self._stream_direct(query, request_id): + async for sse in self._stream_direct(query, request_id, tenant_id): yield sse return # 缓冲模式：工作流执行与 HTTP 流式响应解耦，客户端断开后后台任务继续运行 - task = asyncio.create_task(self._run_workflow_buffered(query, request_id)) + task = asyncio.create_task( + self._run_workflow_buffered(query, request_id, tenant_id) + ) _running_workflow_tasks.add(task) task.add_done_callback(_running_workflow_tasks.discard) try: + terminated = False async for seq, event in sse_event_buffer.read_from(request_id, 0): yield self._format_sse(event, request_id, seq) # result / error 是终止事件，读取到即可结束本轮流式响应 if isinstance(event, dict) and event.get("type") in ("result", "error"): + terminated = True return + # read_from 提前返回（超时/Redis 故障）但未投递终止事件时补一条 error， + # 让客户端明确收到结束信号而非无限等待 + if not terminated: + yield self._format_sse( + { + "type": "error", + "code": "STREAM_INTERRUPTED", + "message": "流式响应意外中断，请重试", + }, + request_id, + 0, + ) except asyncio.CancelledError: # HTTP 客户端断开：让后台任务继续运行以支持重连，重新抛出以让上层清理 raise @@ -95,12 +128,33 @@ async def replay(self, request_id: str, last_seq: int) -> AsyncIterator[str]: 若工作流仍在执行，会持续轮询直到拿到终止事件。 """ - async for seq, event in sse_event_buffer.read_from(request_id, last_seq): - yield self._format_sse(event, request_id, seq) - if isinstance(event, dict) and event.get("type") in ("result", "error"): - return + try: + terminated = False + async for seq, event in sse_event_buffer.read_from(request_id, last_seq): + yield self._format_sse(event, request_id, seq) + if isinstance(event, dict) and event.get("type") in ("result", "error"): + terminated = True + return + if not terminated: + yield self._format_sse( + { + "type": "error", + "code": "STREAM_INTERRUPTED", + "message": "流式响应意外中断，请重试", + }, + request_id, + last_seq, + ) + except asyncio.CancelledError: + raise + + async def exists(self, request_id: str) -> bool: + """判断 request_id 是否在缓冲区中存在，供路由层校验重连请求合法性""" + return await sse_event_buffer.exists(request_id) - async def _stream_direct(self, query: str, request_id: str) -> AsyncIterator[str]: + async def _stream_direct( + self, query: str, request_id: str, tenant_id: str = "default" + ) -> AsyncIterator[str]: """降级模式：Redis 不可用时直接流式输出 SSE 事件 事件即时产出，不经过缓冲区，因此客户端断线后无法续传。 @@ -108,26 +162,44 @@ async def _stream_direct(self, query: str, request_id: str) -> AsyncIterator[str """ seq = 0 - async for event in self._execute_workflow(query, request_id): + async for event in self._execute_workflow( + query, request_id, self.meta_mysql_repository, self.dw_repository, tenant_id + ): seq += 1 yield self._format_sse(event, request_id, seq) - async def _run_workflow_buffered(self, query: str, request_id: str) -> None: + async def _run_workflow_buffered( + self, query: str, request_id: str, tenant_id: str = "default" + ) -> None: """后台执行 LangGraph 工作流，事件追加到 Redis 缓冲区 - 无论成功或失败，最后都会调用 mark_complete 通知 read_from 停止轮询。 - 缓冲区写入失败时事件丢失但工作流仍记日志，不影响主链路。 + 使用独立 Session（而非请求级 Session），确保客户端断连后工作流仍能正常 + 访问数据库。无论成功或失败，最后都会调用 mark_complete 通知 read_from 停止轮询。 """ try: - async for event in self._execute_workflow(query, request_id): - await sse_event_buffer.append(request_id, event) + # 后台任务创建独立 Session + Repository，与 HTTP 请求生命周期解耦 + async with ( + self.meta_session_factory() as meta_session, + self.dw_session_factory() as dw_session, + ): + meta_repository = MetaMySQLRepository(meta_session) + dw_repository = DWMySQLRepository(dw_session) + async for event in self._execute_workflow( + query, request_id, meta_repository, dw_repository, tenant_id + ): + await sse_event_buffer.append(request_id, event) finally: # 无论成功失败都标记完成，让 read_from 终止轮询 await sse_event_buffer.mark_complete(request_id) async def _execute_workflow( - self, query: str, request_id: str + self, + query: str, + request_id: str, + meta_repository: MetaMySQLRepository, + dw_repository: DWRepository, + tenant_id: str = "default", ) -> AsyncIterator[dict]: """执行 LangGraph 工作流并产出事件字典 @@ -151,17 +223,23 @@ async def _execute_workflow( # State 只放会被图节点读写和合并的业务数据，外部工具对象不塞进 State state = DataAgentState(query=query) # Context 保存本次图执行需要复用的外部依赖，节点通过 runtime.context 读取 + # tenant_id 由路由层从 JWT 解析传入，run_sql 节点据此注入行级过滤 context = DataAgentContext( column_qdrant_repository=self.column_qdrant_repository, embedding_client=self.embedding_client, metric_qdrant_repository=self.metric_qdrant_repository, value_es_repository=self.value_es_repository, - meta_mysql_repository=self.meta_mysql_repository, - dw_repository=self.dw_repository, + meta_mysql_repository=meta_repository, + dw_repository=dw_repository, + tenant_id=tenant_id, ) # stream_mode="custom" 对应节点内部 writer(...) 写出的进度消息 + # config 注入 LangSmith/LangFuse 追踪，使图级 trace 可见 async for chunk in graph.astream( - input=state, context=context, stream_mode="custom" + input=state, + context=context, + stream_mode="custom", + config=get_tracing_config(request_id=request_id), ): yield chunk # 拦截最终 result 事件写入缓存，供后续相同 query 命中 diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index ed7291a..0fa8e4c 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -75,12 +75,15 @@ services: container_name: redis restart: unless-stopped command: [ "redis-server", "--appendonly", "yes", "--requirepass", "${REDIS_PASSWORD:?must set REDIS_PASSWORD in .env}" ] + environment: + # 供 healthcheck 通过 REDISCLI_AUTH 读取，避免密码出现在进程列表 / docker inspect + REDIS_PASSWORD: ${REDIS_PASSWORD} volumes: - redis_data:/data networks: - internal healthcheck: - test: [ "CMD", "redis-cli", "-a", "${REDIS_PASSWORD}", "ping" ] + test: [ "CMD-SHELL", "REDISCLI_AUTH=$$REDIS_PASSWORD redis-cli ping" ] interval: 10s timeout: 5s retries: 5 @@ -109,16 +112,16 @@ services: platform: linux/amd64 container_name: embedding restart: unless-stopped + # 监听 8081 与 backend 期望端口一致；MODEL_ID 指向 HuggingFace Hub，容器启动时自动拉取 + command: [ "--port", "8081" ] environment: - MODEL_ID: /models/bge-large-zh-v1.5 + MODEL_ID: BAAI/bge-large-zh-v1.5 MAX_CONCURRENT_REQUESTS: "16" MAX_BATCH_TOKENS: "16384" - volumes: - - ./embedding/bge-large-zh-v1.5:/models/bge-large-zh-v1.5 networks: - internal healthcheck: - test: [ "CMD-SHELL", "bash -c '=2.32.0 - [x] P2-5 数据源扩展抽象层：app/repositories/dw_repository.py 定义 DWRepository ABC（dialect 属性 + get_column_types/get_column_values/get_db_info/validate/run 抽象方法）；dw_mysql_repository.py 继承 ABC 并实现 dialect="mysql"；sql_guard.validate_readonly_sql 新增 dialect 参数，sqlglot.parse/sql() 均按方言解析与渲染；context.py/dependencies.py/query_service.py/add_extra_context.py/run_sql.py/validate_sql.py/meta_knowledge_service.py/build_meta_knowledge.py/graph.py 全链路统一改用 dw_repository:DWRepository 抽象类型；dependencies.get_dw_repository 留扩展点（按配置 dialect 切换 DWClickHouseRepository/DWPostgresRepository） -- [x] P2-6 文档化无法纯代码完成的项（多租户、灰度发布、红蓝对抗）：docs/P2-6_RUNTIME_OPERATIONS_GUIDE.md 输出三项落地指南——多租户（行级隔离 + JWT tenant_id claim + sqlglot AST 强制注入 tenant 过滤 + Qdrant/ES payload 过滤 + Redis key 租户前缀 + 滑动窗口限流）；灰度发布（Nginx 双实例 + X-Canary-Tenant/Cookie 分流 + 1%→10%→50%→100% 渐进放量 + 灰度看板对比 + Alertmanager webhook 自动回滚）；红蓝对抗（5 类攻击场景剧本 ATT-01~05 + MTTD/MTTR 度量 + 每季度演练 SOP）；三项均附验收清单 +- [x] P2-6 多租户核心 + 运营安全指南： + - 代码实现（多租户核心）：app/core/tenancy.py（DEFAULT_TENANT_ID + get_tenant_scoped_tables + get_tenant_column + resolve_tenant_id + inject_tenant_filter sqlglot AST 强制注入 WHERE tenant_id = X）；security.py（User.tenant_id 字段 + JWT tenant_id claim + create_access_token 签发 + get_current_user 解析 + get_demo_users 三元组返回 tenant_id）；auth_router.py（登录签发 JWT 时写入 tenant_id claim）；context.py（DataAgentContext.tenant_id 字段）；query_service.py（query/_stream_direct/_run_workflow_buffered/_execute_workflow 全链路透传 tenant_id 到 context）；query_router.py（current_user.tenant_id 传入 query_service.query）；run_sql.py（SQL 安全校验通过后、执行前调用 inject_tenant_filter 强制注入行级过滤，注入失败拒绝执行返回 TENANT_FILTER_FAILED 错误）；graph.py 本地测试入口传入 tenant_id="default"；.env.example 新增 TENANT_SCOPED_TABLES/TENANT_COLUMN/{ROLE}_TENANT_ID 配置项。设计原则：默认单租户兼容（未配置 TENANT_SCOPED_TABLES 时所有行为不变），程序层强制（AST 注入不依赖 prompt，无法被 SQL 注入绕过），CTE 别名排除（WITH alias AS ... 的 alias 不误判为敏感表） + - 文档化无法纯代码完成的项：docs/P2-6_RUNTIME_OPERATIONS_GUIDE.md 输出三项落地指南——多租户剩余部分（Qdrant/ES payload 过滤需仓储层 schema 改造 + Redis key 租户前缀需 sse_buffer 调用方拼接 + 滑动窗口限流需 Redis 计数器）；灰度发布（Nginx 双实例 + X-Canary-Tenant/Cookie 分流 + 1%→10%→50%→100% 渐进放量 + 灰度看板对比 + Alertmanager webhook 自动回滚）；红蓝对抗（5 类攻击场景剧本 ATT-01~05 + MTTD/MTTR 度量 + 每季度演练 SOP）；三项均附验收清单 --- diff --git a/frontend/src/components/StepRail.tsx b/frontend/src/components/StepRail.tsx index a829ad4..1012b22 100644 --- a/frontend/src/components/StepRail.tsx +++ b/frontend/src/components/StepRail.tsx @@ -23,7 +23,7 @@ const nodes: FlowNode[] = [ { step: "合并召回信息", x: 410, y: 214 }, { step: "过滤指标信息", x: 290, y: 318 }, { step: "过滤表信息", x: 530, y: 318 }, - { step: "增加额外上下文", x: 410, y: 422, w: 176 }, + { step: "添加额外上下文", x: 410, y: 422, w: 176 }, { step: "生成SQL", x: 410, y: 526 }, { step: "校验SQL", x: 410, y: 630 }, { step: "校正SQL", x: 670, y: 630 }, diff --git a/frontend/src/lib/agentApi.ts b/frontend/src/lib/agentApi.ts index 3984945..22aae5c 100644 --- a/frontend/src/lib/agentApi.ts +++ b/frontend/src/lib/agentApi.ts @@ -14,6 +14,8 @@ const API_BASE_URL = import.meta.env.VITE_API_BASE_URL?.replace(/\/$/, "") ?? "" const MAX_RETRIES = 3; /** 重连退避基数：第 N 次重连等待 2^N * base ms，避免雪崩压垮服务端 */ const RETRY_BASE_DELAY_MS = 500; +/** JWT 在 localStorage 中的存储键名 */ +const TOKEN_STORAGE_KEY = "shopkeeper_token"; type QueryOptions = { signal?: AbortSignal; @@ -25,6 +27,15 @@ type ParsedSse = { data: string; }; +function getAuthToken(): string | null { + try { + return localStorage.getItem(TOKEN_STORAGE_KEY); + } catch { + // localStorage 不可用（隐私模式等）时退化为无认证 + return null; + } +} + function parseSseChunk(chunk: string): ParsedSse | null { let id: string | null = null; const dataLines: string[] = []; @@ -47,15 +58,16 @@ function sleep(ms: number, signal?: AbortSignal): Promise { return Promise.reject(new DOMException("Aborted", "AbortError")); } return new Promise((resolve, reject) => { - const timer = setTimeout(resolve, ms); - signal?.addEventListener( - "abort", - () => { - clearTimeout(timer); - reject(new DOMException("Aborted", "AbortError")); - }, - { once: true }, - ); + const onAbort = () => { + clearTimeout(timer); + reject(new DOMException("Aborted", "AbortError")); + }; + // 正常超时触发时移除 abort 监听器，避免监听器累积泄漏 + const timer = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(); + }, ms); + signal?.addEventListener("abort", onAbort, { once: true }); }); } @@ -67,6 +79,11 @@ export async function streamQuery(query: string, options: QueryOptions) { "Content-Type": "application/json", Accept: "text/event-stream", }; + // 携带 JWT：后端 /api/v1/query 强制要求认证，缺失会返回 401 + const token = getAuthToken(); + if (token) { + headers["Authorization"] = `Bearer ${token}`; + } if (lastEventId) { headers["Last-Event-ID"] = lastEventId; } @@ -102,14 +119,18 @@ export async function streamQuery(query: string, options: QueryOptions) { let receivedTerminal = false; const processParsed = (parsed: ParsedSse) => { - if (parsed.id) lastEventId = parsed.id; + // 没有完整 data 的半截事件（如仅 id 行）：丢弃，不更新 lastEventId + // 重连时用上一个成功的 lastEventId 续传，避免跳过未消费事件 if (!parsed.data) return; let event: AgentEvent; try { event = JSON.parse(parsed.data) as AgentEvent; } catch { - event = { type: "error", message: `无法解析后端事件：${parsed.data}` }; + // JSON 解析失败：事件传输不完整，丢弃不设 terminal，让流结束后触发重连 + return; } + // 解析成功后才更新 lastEventId，避免半截事件污染续传位置 + if (parsed.id) lastEventId = parsed.id; if (event.type === "result" || event.type === "error") { receivedTerminal = true; } @@ -144,8 +165,13 @@ export async function streamQuery(query: string, options: QueryOptions) { await sleep(2 ** attempt * RETRY_BASE_DELAY_MS, options.signal); } catch (e) { if (options.signal?.aborted) throw e; + // 已收到终止事件后连接异常是正常收尾，不再重试避免覆盖已展示的结果 + if (receivedTerminal) return; if (attempt >= MAX_RETRIES) throw e; await sleep(2 ** attempt * RETRY_BASE_DELAY_MS, options.signal); + } finally { + // 确保释放 reader，避免流锁定 + reader.releaseLock(); } } } diff --git a/main.py b/main.py index 4414d5e..4f24456 100644 --- a/main.py +++ b/main.py @@ -17,6 +17,7 @@ from sentry_sdk.integrations.starlette import StarletteIntegration from app.api.lifespan import lifespan +from app.api.routers.auth_router import auth_router from app.api.routers.health_router import health_router from app.api.routers.query_router import query_router from app.core.context import request_id_ctx_var @@ -45,10 +46,10 @@ StarletteIntegration(transaction_style="endpoint"), FastApiIntegration(transaction_style="endpoint"), ], - # 业务级异常不应上报到 Sentry：SQL 安全校验拦截、限流、CORS 等属于正常业务流 + # 业务级异常不应上报到 Sentry：SQL 安全校验拦截属于正常业务流 + # 限流（rate_limit）抛的是 fastapi.HTTPException，Sentry Starlette 集成默认已忽略 ignore_errors=[ "app.agent.nodes.sql_guard.SQLSafetyViolation", - "app.core.rate_limit.RateLimitExceeded", ], ) @@ -99,8 +100,8 @@ async def add_request_id(request: Request, call_next): request_id = uuid.uuid4() request_id_ctx_var.set(request_id) response = await call_next(request) - # 把 request_id 回写响应头，前端可记录便于联调排障 - response.headers["X-Request-ID"] = str(request_id) + # 回写 request_id 便于前端联调排障；若路由已显式设置（如 SSE 重连透传原始 id）则不覆盖 + response.headers.setdefault("X-Request-ID", str(request_id)) # 请求被处理之后 return response