Mutual_Fund/model/nl2sql_permission.py

"""NL2SQL 查询权限和查询历史 ORM 模型。"""
from __future__ import annotations

from datetime import datetime

from sqlalchemy import BigInteger, Boolean, CheckConstraint, DateTime, JSON, String, Text, func
from sqlalchemy.orm import Mapped, mapped_column

from model.base import Base


class Nl2SqlQueryRole(Base):
    __tablename__ = "nl2sql_query_role"
    __table_args__ = (
        CheckConstraint("can_query IN (0, 1)", name="ck_nl2sql_role_can_query"),
        CheckConstraint("status IN ('active', 'inactive')", name="ck_nl2sql_role_status"),
        {"comment": "NL2SQL 查询角色，按 sys_user.employee_role 映射"},
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    role_code: Mapped[str] = mapped_column(String(64), unique=True)
    role_name: Mapped[str] = mapped_column(String(128))
    employee_role: Mapped[str] = mapped_column(String(32), unique=True)
    can_query: Mapped[bool] = mapped_column(Boolean, default=False)
    max_rows: Mapped[int] = mapped_column(BigInteger, default=1000)
    daily_quota: Mapped[int] = mapped_column(BigInteger, default=0)
    status: Mapped[str] = mapped_column(String(16), default="active", index=True)
    create_time: Mapped[datetime] = mapped_column(DateTime, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(
        DateTime, server_default=func.now(), onupdate=func.now()
    )


class Nl2SqlRoleTablePermission(Base):
    __tablename__ = "nl2sql_role_table_permission"
    __table_args__ = (
        CheckConstraint("permission = 'SELECT'", name="ck_nl2sql_table_permission"),
        CheckConstraint(
            "row_scope_type IN ('none', 'customer_ids', 'product_ids')",
            name="ck_nl2sql_row_scope_type",
        ),
        CheckConstraint("status IN ('active', 'inactive')", name="ck_nl2sql_table_status"),
        {"comment": "NL2SQL 查询角色的表和行级权限"},
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    role_id: Mapped[int] = mapped_column(BigInteger, index=True)
    table_name: Mapped[str] = mapped_column(String(128))
    permission: Mapped[str] = mapped_column(String(16), default="SELECT")
    row_scope_type: Mapped[str] = mapped_column(String(32), default="none")
    row_scope_column: Mapped[str | None] = mapped_column(String(128))
    status: Mapped[str] = mapped_column(String(16), default="active", index=True)
    create_time: Mapped[datetime] = mapped_column(DateTime, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(
        DateTime, server_default=func.now(), onupdate=func.now()
    )


class Nl2SqlRoleColumnPermission(Base):
    __tablename__ = "nl2sql_role_column_permission"
    __table_args__ = (
        CheckConstraint(
            "access_mode IN ('allow', 'deny', 'mask')",
            name="ck_nl2sql_column_access_mode",
        ),
        CheckConstraint("status IN ('active', 'inactive')", name="ck_nl2sql_column_status"),
        {"comment": "NL2SQL 查询角色的字段权限"},
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    role_id: Mapped[int] = mapped_column(BigInteger, index=True)
    table_name: Mapped[str] = mapped_column(String(128))
    column_name: Mapped[str] = mapped_column(String(128))
    access_mode: Mapped[str] = mapped_column(String(16), default="allow")
    mask_type: Mapped[str | None] = mapped_column(String(32))
    status: Mapped[str] = mapped_column(String(16), default="active", index=True)
    create_time: Mapped[datetime] = mapped_column(DateTime, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(
        DateTime, server_default=func.now(), onupdate=func.now()
    )


class Nl2SqlSensitiveField(Base):
    __tablename__ = "nl2sql_sensitive_field"
    __table_args__ = (
        CheckConstraint("status IN ('active', 'inactive')", name="ck_nl2sql_sensitive_status"),
        {"comment": "NL2SQL 全局敏感字段规则"},
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    table_name: Mapped[str] = mapped_column(String(128))
    column_name: Mapped[str] = mapped_column(String(128))
    mask_type: Mapped[str] = mapped_column(String(32), default="partial")
    status: Mapped[str] = mapped_column(String(16), default="active", index=True)
    description: Mapped[str | None] = mapped_column(String(255))
    create_time: Mapped[datetime] = mapped_column(DateTime, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(
        DateTime, server_default=func.now(), onupdate=func.now()
    )


class Nl2SqlQueryHistory(Base):
    __tablename__ = "nl2sql_query_history"
    __table_args__ = (
        CheckConstraint(
            "status IN ('success', 'failed', 'blocked', 'timeout')",
            name="ck_nl2sql_history_status",
        ),
        {"comment": "NL2SQL 查询归档，不保存完整结果行"},
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    query_id: Mapped[str] = mapped_column(String(64), unique=True, index=True)
    user_id: Mapped[int] = mapped_column(BigInteger, index=True)
    session_id: Mapped[str | None] = mapped_column(String(64), index=True)
    caller_agent: Mapped[str | None] = mapped_column(String(64), index=True)
    question: Mapped[str] = mapped_column(Text)
    generated_sql: Mapped[str | None] = mapped_column(Text)
    access_tables: Mapped[list | None] = mapped_column(JSON)
    status: Mapped[str] = mapped_column(String(16), index=True)
    error_code: Mapped[str | None] = mapped_column(String(64))
    error_message: Mapped[str | None] = mapped_column(String(512))
    row_count: Mapped[int] = mapped_column(BigInteger, default=0)
    truncated: Mapped[bool] = mapped_column(Boolean, default=False)
    elapsed_ms: Mapped[float | None] = mapped_column()
    trace_id: Mapped[str | None] = mapped_column(String(64), index=True)
    create_time: Mapped[datetime] = mapped_column(
        DateTime, server_default=func.now(), index=True
    )