feat:新增投顾agent和nl2sqlagent

2026-09-13 16:19:24 +08:00
parent c80c6acac0
commit 163192bf55
122 changed files with 7488 additions and 362 deletions
@@ -0,0 +1,97 @@
+"""NL2SQL 只读 SQL 的 AST 安全校验。"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from sqlglot import exp, parse
+from sqlglot.errors import ParseError
+
+
+_DANGEROUS_FUNCTIONS = {
+    "LOAD_FILE",
+    "UUID_FILE_NAME",
+    "BENCHMARK",
+    "SLEEP",
+}
+
+class SqlSecurityError(ValueError):
+    """SQL 未通过只读和权限校验。"""
+
+
+@dataclass(frozen=True)
+class ValidatedSql:
+    sql: str
+    access_tables: set[str]
+
+
+def _read_limit(statement: exp.Expression) -> int | None:
+    limit = statement.args.get("limit")
+    if limit is None:
+        return None
+    expression = limit.args.get("expression")
+    if not isinstance(expression, exp.Literal) or not expression.is_number:
+        raise SqlSecurityError("只允许使用数字 LIMIT")
+    value = int(expression.this)
+    if value < 0:
+        raise SqlSecurityError("LIMIT 不能为负数")
+    return value
+
+
+def validate_select_sql(
+    sql: str,
+    *,
+    authorized_tables: set[str],
+    authorized_columns: dict[str, set[str]] | None = None,
+    max_rows: int,
+    max_joins: int = 5,
+    max_columns: int = 100,
+) -> ValidatedSql:
+    """校验单条 SELECT，检查授权表并将 LIMIT 控制在最大行数内。"""
+    if not isinstance(sql, str) or not sql.strip():
+        raise SqlSecurityError("SQL 不能为空")
+    if max_rows <= 0:
+        raise SqlSecurityError("最大行数必须为正数")
+    if max_joins < 0 or max_columns <= 0:
+        raise SqlSecurityError("SQL 复杂度限制参数无效")
+    try:
+        statements = parse(sql, read="mysql")
+    except ParseError as exc:
+        raise SqlSecurityError("SQL 解析失败") from exc
+    if len(statements) != 1 or not isinstance(statements[0], exp.Select):
+        raise SqlSecurityError("只允许执行单条 SELECT")
+
+    statement = statements[0]
+    join_count = len(list(statement.find_all(exp.Join)))
+    if join_count > max_joins:
+        raise SqlSecurityError("SQL JOIN 深度超过限制")
+    if len(statement.expressions) > max_columns:
+        raise SqlSecurityError("SQL 返回列数超过限制")
+    for function in statement.find_all(exp.Func):
+        function_name = getattr(function, "name", "") or function.sql_name()
+        if function_name.upper() in _DANGEROUS_FUNCTIONS:
+            raise SqlSecurityError("SQL 包含危险函数")
+    access_tables: set[str] = set()
+    aliases: dict[str, str] = {}
+    for table in statement.find_all(exp.Table):
+        if table.db or table.catalog:
+            raise SqlSecurityError("禁止跨库访问")
+        access_tables.add(table.name)
+        aliases[table.alias_or_name] = table.name
+    if not access_tables.issubset(authorized_tables):
+        raise SqlSecurityError("SQL 访问了未授权表")
+    if authorized_columns is not None:
+        if statement.find(exp.Star):
+            raise SqlSecurityError("配置字段权限时禁止 SELECT *")
+        default_table = next(iter(access_tables), None) if len(access_tables) == 1 else None
+        for column in statement.find_all(exp.Column):
+            table_name = aliases.get(column.table, column.table) or default_table
+            if table_name is None or column.name not in authorized_columns.get(table_name, set()):
+                raise SqlSecurityError("SQL 访问了未授权字段")
+
+    current_limit = _read_limit(statement)
+    if current_limit is None or current_limit > max_rows:
+        statement = statement.limit(max_rows)
+    return ValidatedSql(
+        sql=statement.sql(dialect="mysql"),
+        access_tables=access_tables,
+    )