Mutual_Fund/nl2sql/row_scope.py

"""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 list(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"], table=table.alias_or_name),
            expressions=[exp.Literal.number(value) for value in values],
        )
        owner = table.find_ancestor(exp.Select) or statement
        owner.where(condition, copy=False)
    return statement.sql(dialect="mysql")