From 051857e0aa3d659bb40cff7c440af4a900722ff4 Mon Sep 17 00:00:00 2001 From: ulleo Date: Thu, 12 Nov 2026 18:09:00 +0800 Subject: [PATCH] feat: add startup option to skip SQL table checks --- backend/README.md | 16 +++- backend/apps/chat/task/llm.py | 29 +++--- backend/common/core/config.py | 4 + backend/tests/test_table_check.py | 146 ++++++++++++++++++++++++++++++ docs/agents/security.md | 3 +- 5 files changed, 182 insertions(+), 16 deletions(-) create mode 100644 backend/tests/test_table_check.py diff --git a/backend/README.md b/backend/README.md index b6dbbb6d3..0b2efd9b3 100644 --- a/backend/README.md +++ b/backend/README.md @@ -1 +1,15 @@ -# FastAPI Project - Backend \ No newline at end of file +# FastAPI Project - Backend + +## 表检查启动参数 + +`SQLBOT_SKIP_TABLE_CHECK` 默认为 `false`。不配置时，生成的 SQL 仍通过真实表名提取和本轮允许表集合检查。 + +仅在需要跳过这组检查时设置： + +```dotenv +SQLBOT_SKIP_TABLE_CHECK=true +``` + +通过进程环境变量传入；本地开发也可写入仓库根目录 `.env`，安装版可写入安装目录的 `conf/sqlbot.conf`。Docker Compose 部署可在 SQLBot 服务的 `environment` 中设置。修改后需重启服务或重建容器，所有问数入口共用该启动配置，不提供配置页面或运行时修改。 + +开启后跳过表名提取、无法提取表名时的拒绝及表白名单检查，可能允许查询数据源账号可读取的未勾选表。SQL 只读、危险函数、元数据查询控制、行权限处理及行数限制仍按原有流程执行。 diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 4eca16d1e..245f3dc92 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -1578,20 +1578,21 @@ def run_task(self, in_chat: bool = True, stream: bool = True, sql_operate = OperationEnum.GENERATE_SQL sql, tables = self.check_sql(session=_session, res=full_sql_text, operate=sql_operate) - # 表名安全检查：用 sqlglot 解析真实 SQL，不信任 AI 返回的 tables - actual_tables = extract_tables_from_sql(sql, ds_type=self.ds.type) - if not actual_tables: - raise SingleMessageError( - "SQL parsing failed: unable to extract table names. " - "This may indicate an unsupported SQL syntax or a security issue." - ) - allowed_tables = set(self.table_name_list) - unauthorized_tables = actual_tables - allowed_tables - if unauthorized_tables: - raise SingleMessageError( - f"SQL contains unauthorized tables: {', '.join(unauthorized_tables)}. " - f"Allowed tables: {', '.join(allowed_tables)}" - ) + if not settings.SQLBOT_SKIP_TABLE_CHECK: + # 表名安全检查：用 sqlglot 解析真实 SQL，不信任 AI 返回的 tables + actual_tables = extract_tables_from_sql(sql, ds_type=self.ds.type) + if not actual_tables: + raise SingleMessageError( + "SQL parsing failed: unable to extract table names. " + "This may indicate an unsupported SQL syntax or a security issue." + ) + allowed_tables = set(self.table_name_list) + unauthorized_tables = actual_tables - allowed_tables + if unauthorized_tables: + raise SingleMessageError( + f"SQL contains unauthorized tables: {', '.join(unauthorized_tables)}. " + f"Allowed tables: {', '.join(allowed_tables)}" + ) if ((not self.current_assistant or is_page_embedded) and is_normal_user( self.current_user)) or use_dynamic_ds: diff --git a/backend/common/core/config.py b/backend/common/core/config.py index 0373728f9..8305c2c8a 100644 --- a/backend/common/core/config.py +++ b/backend/common/core/config.py @@ -117,6 +117,9 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str: # 默认关闭，防止通过元数据查询泄露数据库结构 SQLBOT_ALLOW_METADATA_QUERIES: bool = False + # 启动配置：显式开启时跳过生成 SQL 的表名提取和白名单检查 + SQLBOT_SKIP_TABLE_CHECK: bool = False + PARSE_REASONING_BLOCK_ENABLED: bool = True DEFAULT_REASONING_CONTENT_START: str = '' DEFAULT_REASONING_CONTENT_END: str = '' @@ -136,6 +139,7 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str: ORACLE_CLIENT_PATH: str = '/opt/sqlbot/db_client/oracle_instant_client' @field_validator('SQL_DEBUG', + 'SQLBOT_SKIP_TABLE_CHECK', 'EMBEDDING_ENABLED', 'GENERATE_SQL_QUERY_LIMIT_ENABLED', 'PARSE_REASONING_BLOCK_ENABLED', diff --git a/backend/tests/test_table_check.py b/backend/tests/test_table_check.py new file mode 100644 index 000000000..1adb7f2aa --- /dev/null +++ b/backend/tests/test_table_check.py @@ -0,0 +1,146 @@ +"""Offline task regressions with model, persistence and database I/O replaced.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import sqlbot_xpack # noqa: F401 # Initialize the published wheel before app imports. +from pydantic import ValidationError + +from apps.chat.models.chat_model import ChatFinishStep +from apps.chat.task import llm +from common.core.config import Settings + + +@pytest.mark.parametrize( + "value,expected", + [(None, False), ("false", False), ("true", True), (" TRUE ", True)], +) +def test_startup_configuration(monkeypatch, value, expected): + monkeypatch.delenv("SQLBOT_SKIP_TABLE_CHECK", raising=False) + if value is not None: + monkeypatch.setenv("SQLBOT_SKIP_TABLE_CHECK", value) + config = Settings(_env_file=None) + assert config.SQLBOT_SKIP_TABLE_CHECK is expected + monkeypatch.setenv("SQLBOT_SKIP_TABLE_CHECK", str(not expected)) + assert config.SQLBOT_SKIP_TABLE_CHECK is expected + + +def test_invalid_startup_configuration_is_rejected(monkeypatch): + monkeypatch.setenv("SQLBOT_SKIP_TABLE_CHECK", "invalid") + with pytest.raises(ValidationError, match="SQLBOT_SKIP_TABLE_CHECK"): + Settings(_env_file=None) + + +@pytest.fixture +def task(monkeypatch): + monkeypatch.setattr(llm.settings, "SQLBOT_SKIP_TABLE_CHECK", False) + monkeypatch.setattr(llm, "session_maker", Mock()) + monkeypatch.setattr(llm, "check_connection", Mock(return_value=True)) + monkeypatch.setattr(llm, "start_log", Mock()) + monkeypatch.setattr(llm, "end_log", Mock()) + monkeypatch.setattr(llm, "get_chat_chart_data", Mock(return_value=[])) + service = Mock(spec=llm.LLMService) + service.ds = SimpleNamespace(id=1, type="pg") + service.chat_question = SimpleNamespace(question="sales") + service.record = SimpleNamespace(id=1) + service.get_record.return_value = service.record + service.current_user = SimpleNamespace(id=2, isAdmin=False, weight=0) + service.current_assistant = None + service.change_title = False + service.table_name_list = ["sales"] + service.current_logs = {} + service.generate_sql.return_value = [{"content": "{}"}] + service.generate_filter.return_value = None + service.execute_sql.return_value = {"data": []} + return service + + +def run(task, sql, reported_tables=None): + # Model-reported tables are deliberately independent of the real SQL. + task.check_sql.return_value = sql, reported_tables or ["sales"] + task.check_save_sql.return_value = sql + results = list( + llm.LLMService.run_task( + task, in_chat=False, stream=False, finish_step=ChatFinishStep.QUERY_DATA + ) + ) + assert len(results) == 1 + return results[0] + + +@pytest.mark.parametrize( + "sql", + [ + "SELECT id FROM sales", + "WITH totals AS (SELECT id FROM sales) SELECT * FROM totals", + ], +) +def test_default_allows_known_tables_and_cte_aliases(task, sql): + result = run(task, sql) + assert result["success"] is True + task.execute_sql.assert_called_once_with(sql=sql) + task.check_save_sql.assert_called_once() + task.save_error.assert_not_called() + + +@pytest.mark.parametrize( + "sql,error", + [ + ("SELECT * FROM secret_payroll", "unauthorized tables: secret_payroll"), + ( + "SELECT id FROM sales UNION ALL SELECT id FROM secret_payroll", + "unauthorized tables: secret_payroll", + ), + ("SELECT 1", "unable to extract table names"), + ("SELECT (", "unable to extract table names"), + ], +) +def test_default_rejects_invalid_or_unknown_tables_before_save_and_execute( + task, sql, error +): + result = run(task, sql, reported_tables=["sales"]) + assert result["success"] is False + assert error in result["message"] + task.save_error.assert_called_once() + task.check_save_sql.assert_not_called() + task.execute_sql.assert_not_called() + + +@pytest.mark.parametrize("sql", ["SELECT * FROM secret_payroll", "SELECT 1"]) +def test_explicit_skip_bypasses_extraction_and_keeps_remaining_flow( + task, monkeypatch, sql +): + monkeypatch.setattr(llm.settings, "SQLBOT_SKIP_TABLE_CHECK", True) + extract = Mock(side_effect=AssertionError("table extraction must be skipped")) + monkeypatch.setattr(llm, "extract_tables_from_sql", extract) + result = run(task, sql) + assert result["success"] is True + extract.assert_not_called() + task.generate_filter.assert_called_once() + task.check_save_sql.assert_called_once() + task.execute_sql.assert_called_once_with(sql=sql) + task.save_error.assert_not_called() + + +@pytest.mark.parametrize( + "sql", + [ + "DELETE FROM sales", + "SELECT pg_read_file('/etc/passwd')", + "SHOW TABLES", + ], +) +def test_skip_keeps_database_execution_guards(task, monkeypatch, sql): + from apps.db import db + + monkeypatch.setattr(llm.settings, "SQLBOT_SKIP_TABLE_CHECK", True) + monkeypatch.setattr(llm.settings, "SQLBOT_ALLOW_METADATA_QUERIES", False) + connect = Mock(side_effect=AssertionError("database must not be contacted")) + monkeypatch.setattr(db, "get_session", connect) + task.execute_sql.side_effect = lambda sql: llm.LLMService.execute_sql(task, sql) + result = run(task, sql) + assert result["success"] is False + assert "SQL can only contain read operations" in result["message"] + task.execute_sql.assert_called_once_with(sql=sql) + connect.assert_not_called() diff --git a/docs/agents/security.md b/docs/agents/security.md index 3fbcaacd4..81396a3ec 100644 --- a/docs/agents/security.md +++ b/docs/agents/security.md @@ -33,7 +33,7 @@ 1. 只把允许的表和字段元数据提供给模型； 2. 解析模型返回 JSON 并提取 SQL； 3. 用 `sqlglot` 提取真实表名并排除 CTE 别名； -4. 对照允许表集合拒绝未授权表； +4. 默认对照允许表集合拒绝未授权表；部署显式设置 `SQLBOT_SKIP_TABLE_CHECK=true` 时，跳过第 3、4 步及无法提取表名时的拒绝； 5. 按用户、助手和页面嵌入场景应用行权限或动态 SQL 变换； 6. 校验并保存最终 SQL； 7. 执行最终 SQL 并保存结果。 @@ -43,6 +43,7 @@ - 不要信任模型返回的 `tables` 字段替代真实 SQL 解析。 - 行权限拼接沿用白名单逻辑运算符和 `_escape_sql_value`，不要引入新的字符串拼接入口。 - 保持 SQL 行数限制和 `SQLBOT_ALLOW_METADATA_QUERIES` 默认关闭语义；不要为了调试全局开启元数据查询。 +- `SQLBOT_SKIP_TABLE_CHECK` 默认为 `false`，仅在启动时读取，不提供参数页面或运行时修改。开启后可能访问数据源账号可读取的未勾选表；SQL 只读、危险函数、元数据查询控制与既有权限处理不随此开关关闭。配置方式见 [`backend/README.md`](../../backend/README.md)。 - 动态 SQL 生成后仍必须走同一解析、校验和保存路径。 - 数据源连接配置解密只发生在服务端执行路径，不进入日志或 API 响应。