43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
"""NL2SQL 行级权限条件构造与注入。"""
|
|
from __future__ import annotations
|
|
|
|
from sqlglot import exp, parse_one
|
|
from sqlglot.errors import ParseError
|
|
|
|
|
|
class RowScopeError(ValueError):
|
|
"""行级权限范围缺失或格式不合法。"""
|
|
|
|
|
|
def _values(data_scope: dict, scope_type: str) -> list[int]:
|
|
values = data_scope.get(scope_type)
|
|
if not isinstance(values, list) or not values or any(
|
|
isinstance(value, bool) or not isinstance(value, int) for value in values
|
|
):
|
|
raise RowScopeError(f"缺少有效的 {scope_type} 行范围")
|
|
return sorted(set(values))
|
|
|
|
|
|
def apply_row_scope(sql: str, permission: dict, data_scope: dict | None) -> str:
|
|
"""按照服务端权限快照为每个受限表注入 IN 条件。"""
|
|
scopes = permission.get("row_scopes", {})
|
|
if not scopes:
|
|
return sql
|
|
if not isinstance(data_scope, dict):
|
|
raise RowScopeError("缺少行范围参数")
|
|
try:
|
|
statement = parse_one(sql, read="mysql")
|
|
except ParseError as exc:
|
|
raise RowScopeError("SQL 解析失败") from exc
|
|
for table in statement.find_all(exp.Table):
|
|
scope = scopes.get(table.name)
|
|
if not scope:
|
|
continue
|
|
values = _values(data_scope, scope["type"])
|
|
condition = exp.In(
|
|
this=exp.column(scope["column"]),
|
|
expressions=[exp.Literal.number(value) for value in values],
|
|
)
|
|
statement = statement.where(condition)
|
|
return statement.sql(dialect="mysql")
|