From 080d4e08a832cd4a90e2cfab7935e595970945e1 Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 15 Jul 2025 17:32:03 +0800 Subject: [PATCH] feat: data permission --- backend/alembic/versions/025_ds_num.py | 29 ++++ .../versions/026_row_column_permission.py | 54 ++++++++ backend/apps/chat/models/chat_model.py | 8 ++ backend/apps/chat/task/llm.py | 93 ++++++++++++- backend/apps/datasource/crud/datasource.py | 52 +++++++- .../apps/datasource/crud/row_permission.py | 126 ++++++++++++++++++ backend/apps/datasource/models/datasource.py | 2 + backend/apps/template/filter/__init__.py | 2 + backend/apps/template/filter/generator.py | 6 + backend/pyproject.toml | 2 +- 10 files changed, 366 insertions(+), 8 deletions(-) create mode 100644 backend/alembic/versions/025_ds_num.py create mode 100644 backend/alembic/versions/026_row_column_permission.py create mode 100644 backend/apps/datasource/crud/row_permission.py create mode 100644 backend/apps/template/filter/__init__.py create mode 100644 backend/apps/template/filter/generator.py diff --git a/backend/alembic/versions/025_ds_num.py b/backend/alembic/versions/025_ds_num.py new file mode 100644 index 00000000..35fdfaee --- /dev/null +++ b/backend/alembic/versions/025_ds_num.py @@ -0,0 +1,29 @@ +"""025_ds_num + +Revision ID: 97dcdbedaaf3 +Revises: 4c6d18a18bd4 +Create Date: 2025-07-15 15:50:56.942959 + +""" +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = '97dcdbedaaf3' +down_revision = '806bc67ff45f' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('core_datasource', sa.Column('num', sqlmodel.sql.sqltypes.AutoString(length=256), nullable=True)) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column('core_datasource', 'num') + # ### end Alembic commands ### diff --git a/backend/alembic/versions/026_row_column_permission.py b/backend/alembic/versions/026_row_column_permission.py new file mode 100644 index 00000000..fac243ea --- /dev/null +++ b/backend/alembic/versions/026_row_column_permission.py @@ -0,0 +1,54 @@ +"""026_row_column_permission + +Revision ID: 4c6d18a18bd4 +Revises: 863105882eba +Create Date: 2025-06-25 17:32:09.183257 + +""" +import sqlalchemy as sa +import sqlmodel.sql.sqltypes +from alembic import op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = '4c6d18a18bd4' +down_revision = '97dcdbedaaf3' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('ds_permission', + sa.Column('id', sa.Integer(), sa.Identity(always=True), nullable=False), + sa.Column('enable', sa.Boolean(), nullable=False), + sa.Column('auth_target_type', sa.String(128), nullable=False), + sa.Column('auth_target_id', sa.BigInteger(), nullable=True), + sa.Column('type', sa.String(64), nullable=False), + sa.Column('ds_id', sa.BigInteger(), nullable=True), + sa.Column('table_id', sa.BigInteger(), nullable=True), + sa.Column('expression_tree', sa.Text(), nullable=True), + sa.Column('permissions', sa.Text(), nullable=True), + sa.Column('white_list_user', sa.Text(), nullable=True), + sa.Column('create_time', sa.DateTime(), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + op.create_table('ds_rules', + sa.Column('id', sa.Integer(), sa.Identity(always=True), nullable=False), + sa.Column('enable', sa.Boolean(), nullable=False), + sa.Column('name', sqlmodel.sql.sqltypes.AutoString(length=128), nullable=False), + sa.Column('description', sqlmodel.sql.sqltypes.AutoString(length=512), nullable=True), + sa.Column('permission_list', postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.Column('user_list', postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.Column('white_list_user', sa.Text(), nullable=True), + sa.Column('create_time', sa.DateTime(), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table('ds_rules') + op.drop_table('ds_permission') + # ### end Alembic commands ### diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 4aa02266..3f891cbf 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -11,6 +11,7 @@ from apps.template.generate_guess_question.generator import get_guess_question_t from apps.template.generate_predict.generator import get_predict_template from apps.template.generate_sql.generator import get_sql_template from apps.template.select_datasource.generator import get_datasource_template +from apps.template.filter.generator import get_permissions_template class Chat(SQLModel, table=True): @@ -98,6 +99,7 @@ class AiModelQuestion(BaseModel): fields: str = "" data: str = "" lang: str = "zh-CN" + filter: str = [] def sql_sys_question(self): return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question) @@ -137,6 +139,12 @@ class AiModelQuestion(BaseModel): return get_guess_question_template()['user'].format(question=self.question, schema=self.db_schema, old_questions=old_questions, lang=self.lang) + def filter_sys_question(self): + return get_permissions_template()['system'] + + def filter_user_question(self): + return get_permissions_template()['user'].format(sql=self.sql, filter=self.filter, lang=self.lang) + class ChatQuestion(AiModelQuestion): question: str = '' diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index e5ee344f..ffb4736b 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -1,3 +1,4 @@ +import json import logging import traceback import warnings @@ -10,8 +11,13 @@ import requests from langchain.chat_models.base import BaseChatModel from langchain_community.utilities import SQLDatabase from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, BaseMessageChunk +from sqlalchemy import and_, cast from sqlalchemy import select +from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import load_only +from sqlbot_xpack.permissions.api.permission import transRecord2DTO +from sqlbot_xpack.permissions.models.ds_permission import DsPermission, PermissionDTO +from sqlbot_xpack.permissions.models.ds_rules import DsRules from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_default_config from apps.chat.curd.chat import save_question, save_full_sql_message, save_full_sql_message_and_answer, save_sql, \ @@ -21,7 +27,8 @@ from apps.chat.curd.chat import save_question, save_full_sql_message, save_full_ get_old_questions, save_analysis_predict_record, list_base_records, rename_chat from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat from apps.datasource.crud.datasource import get_table_schema -from apps.datasource.models.datasource import CoreDatasource +from apps.datasource.crud.row_permission import transFilterTree +from apps.datasource.models.datasource import CoreDatasource, CoreTable from apps.db.db import exec_sql from apps.system.crud.assistant import get_assistant_ds from common.core.config import settings @@ -77,7 +84,7 @@ class LLMService: # get schema if ds: - chat_question.db_schema = get_table_schema(session=self.session, ds=ds) + chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds) chat_question.lang = current_user.language @@ -467,6 +474,78 @@ class LLMService: [{'type': msg.type, 'content': msg.content} for msg in self.sql_message]).decode()) + def generate_filter(self, sql: str, tables: List): + table_list = self.session.query(CoreTable).filter( + and_(CoreTable.ds_id == self.ds.id, CoreTable.table_name.in_(tables)) + ).all() + + filters = [] + for table in table_list: + row_permissions = self.session.query(DsPermission).filter( + and_(DsPermission.table_id == table.id, DsPermission.type == 'row')).all() + res: List[PermissionDTO] = [] + if row_permissions is not None: + for permission in row_permissions: + # check permission and user in same rules + obj = self.session.query(DsRules).filter( + and_(DsRules.permission_list.op('@>')(cast([permission.id], JSONB)), + DsRules.user_list.op('@>')(cast([self.current_user.id], JSONB))) + ).first() + if obj is not None: + res.append(transRecord2DTO(self.session, permission)) + wheres = transFilterTree(self.session, res, self.ds) + filters.append({"table": table.table_name, "filter": wheres}) + + filter = json.dumps(filters, ensure_ascii=False) + # filter = f"""[{{"table":"{tables[0]}","filter":"省份 = '广东省' or 销售额（万元） > 10000"}}]""" # todo get filters + self.chat_question.sql = sql + self.chat_question.filter = filter + msg: List[Union[BaseMessage, dict[str, Any]]] = [] + msg.append(SystemMessage(content=self.chat_question.filter_sys_question())) + msg.append(HumanMessage(content=self.chat_question.filter_user_question())) + + history_msg = [] + # if self.record.full_analysis_message and self.record.full_analysis_message.strip() != '': + # history_msg = orjson.loads(self.record.full_analysis_message) + + # self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, answer='', + # full_message=orjson.dumps(history_msg + + # [{'type': msg.type, + # 'content': msg.content} for msg + # in + # msg]).decode()) + full_thinking_text = '' + full_filter_text = '' + res = self.llm.stream(msg) + token_usage = {} + for chunk in res: + print(chunk) + reasoning_content_chunk = '' + if 'reasoning_content' in chunk.additional_kwargs: + reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '') + # else: + # reasoning_content_chunk = chunk.get('reasoning_content') + if reasoning_content_chunk is None: + reasoning_content_chunk = '' + full_thinking_text += reasoning_content_chunk + + full_filter_text += chunk.content + # yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk} + get_token_usage(chunk, token_usage) + + msg.append(AIMessage(full_filter_text)) + # self.record = save_full_analysis_message_and_answer(session=self.session, record_id=self.record.id, + # token_usage=token_usage, + # answer=orjson.dumps({'content': full_analysis_text, + # 'reasoning_content': full_thinking_text}).decode(), + # full_message=orjson.dumps(history_msg + + # [{'type': msg.type, + # 'content': msg.content} for msg + # in + # analysis_msg]).decode()) + print(full_filter_text) + return full_filter_text + def generate_chart(self): # append current question self.chart_message.append(HumanMessage(self.chat_question.chart_user_question())) @@ -663,7 +742,15 @@ def run_task(llm_service: LLMService, in_chat: bool = True): # filter sql print(full_sql_text) - sql = llm_service.check_save_sql(res=full_sql_text) + + # todo row permission + sql_json_str = extract_nested_json(full_sql_text) + data = orjson.loads(sql_json_str) + sql_result = llm_service.generate_filter(data['sql'], data['tables']) + print(sql_result) + sql = llm_service.check_save_sql(res=sql_result) + # sql = llm_service.check_save_sql(res=full_sql_text) + print(sql) if in_chat: yield orjson.dumps({'content': sql, 'type': 'sql'}).decode() + '\n\n' diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 4fc7b183..9430a79f 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -2,7 +2,10 @@ import datetime import json from typing import List -from sqlalchemy import and_, text +from sqlalchemy import and_, text, cast +from sqlalchemy.dialects.postgresql import JSONB +from sqlbot_xpack.permissions.models.ds_permission import DsPermission +from sqlbot_xpack.permissions.models.ds_rules import DsRules from sqlmodel import select from apps.datasource.utils.utils import aes_decrypt @@ -12,6 +15,7 @@ from apps.db.engine import get_engine_conn from apps.db.type import db_type_relation from common.core.deps import SessionDep, CurrentUser from common.utils.utils import deepcopy_ignore_extra +from .table import get_tables_by_ds_id from ..crud.field import delete_field_by_ds_id, update_field from ..crud.table import delete_table_by_ds_id, update_table from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, TableObj, \ @@ -71,12 +75,14 @@ def create_ds(session: SessionDep, user: CurrentUser, create_ds: CreateDatasourc # save tables and fields sync_table(session, ds, create_ds.tables) + updateNum(session, ds) return ds def chooseTables(session: SessionDep, id: int, tables: List[CoreTable]): ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first() sync_table(session, ds, tables) + updateNum(session, ds) def update_ds(session: SessionDep, ds: CoreDatasource): @@ -239,20 +245,49 @@ def preview(session: SessionDep, id: int, data: TableObj): return exec_sql(ds, sql) -def get_table_obj_by_ds(session: SessionDep, ds: CoreDatasource) -> List[TableAndFields]: +def updateNum(session: SessionDep, ds: CoreDatasource): + all_tables = get_tables(ds) + selected_tables = get_tables_by_ds_id(session, ds.id) + num = f'{len(selected_tables)}/{len(all_tables)}' + + record = session.exec(select(CoreDatasource).where(CoreDatasource.id == ds.id)).first() + update_data = ds.model_dump(exclude_unset=True) + for field, value in update_data.items(): + setattr(record, field, value) + record.num = num + session.add(record) + session.commit() + + +def get_table_obj_by_ds(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource) -> List[TableAndFields]: _list: List = [] tables = session.query(CoreTable).filter(CoreTable.ds_id == ds.id).all() conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config() schema = conf.dbSchema if conf.dbSchema is not None and conf.dbSchema != "" else conf.database for table in tables: fields = session.query(CoreField).filter(and_(CoreField.table_id == table.id, CoreField.checked == True)).all() + + # do column permissions, filter fields + column_permissions = session.query(DsPermission).filter( + and_(DsPermission.table_id == table.id, DsPermission.type == 'column')).all() + if column_permissions is not None: + for permission in column_permissions: + # check permission and user in same rules + obj = session.query(DsRules).filter( + and_(DsRules.permission_list.op('@>')(cast([permission.id], JSONB)), + DsRules.user_list.op('@>')(cast([current_user.id], JSONB))) + ).first() + if obj is not None: + permission_list = json.loads(permission.permissions) + fields = filter_list(fields, permission_list) + _list.append(TableAndFields(schema=schema, table=table, fields=fields)) return _list -def get_table_schema(session: SessionDep, ds: CoreDatasource) -> str: +def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource) -> str: schema_str = "" - table_objs = get_table_obj_by_ds(session=session, ds=ds) + table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds) if len(table_objs) == 0: return schema_str db_name = table_objs[0].schema @@ -279,3 +314,12 @@ def get_table_schema(session: SessionDep, ds: CoreDatasource) -> str: schema_str += ",\n".join(field_list) schema_str += '\n]\n' return schema_str + + +def filter_list(list_a, list_b): + id_to_invalid = {} + for b in list_b: + if not b['enable']: + id_to_invalid[b['field_id']] = True + + return [a for a in list_a if not id_to_invalid.get(a.id, False)] diff --git a/backend/apps/datasource/crud/row_permission.py b/backend/apps/datasource/crud/row_permission.py new file mode 100644 index 00000000..b9b84558 --- /dev/null +++ b/backend/apps/datasource/crud/row_permission.py @@ -0,0 +1,126 @@ +# Author: Junjun +# Date: 2025/6/25 + +from typing import List, Dict + +from apps.datasource.models.datasource import CoreField, CoreDatasource +from sqlbot_xpack.permissions.models.ds_permission import PermissionTree, PermissionDTO +from common.core.deps import SessionDep + + +def transFilterTree(session: SessionDep, tree_list: List[PermissionDTO], ds: CoreDatasource) -> str | None: + if tree_list is None: + return None + res: List[str] = [] + for dto in tree_list: + tree = dto.tree + if tree is None: + continue + tree_exp = transTreeToWhere(session, tree, ds) + if tree_exp is not None: + res.append(tree_exp) + return " AND ".join(res) + + +def transTreeToWhere(session: SessionDep, tree: PermissionTree, ds: CoreDatasource) -> str | None: + if tree is None: + return None + logic = tree['logic'] + + items = tree['items'] + list: List[str] = [] + if items is not None: + for item in items: + exp: str = None + if item['type'] == 'item': + exp = transTreeItem(session, item, ds) + elif item['type'] == 'tree': + exp = transTreeToWhere(session, item['sub_tree'], ds) + + if exp is not None: + list.append(exp) + return '(' + f' {logic} '.join(list) + ')' if len(list) > 0 else None + + +def transTreeItem(session: SessionDep, item: Dict, ds: CoreDatasource) -> str | None: + res: str = None + field = session.query(CoreField).filter(CoreField.id == int(item['field_id'])).first() + if field is None: + return None + + whereName = field.field_name + if item['filter_type'] == 'enum': + if len(item['enum_value']) > 0: + if ds['type'] == 'sqlServer' and ( + field.field_type == 'nchar' or field.field_type == 'NCHAR' or field.field_type == 'nvarchar' or field.field_type == 'NVARCHAR'): + res = "(" + whereName + " IN (N'" + "',N'".join(item['enum_value']) + "'))" + else: + res = "(" + whereName + " IN ('" + "','".join(item['enum_value']) + "'))" + else: + value = item['value'] + whereTerm = transFilterTerm(item['term']) + whereValue = '' + + if item['term'] == 'null': + whereValue = '' + elif item['term'] == 'not_null': + whereValue = '' + elif item['term'] == 'empty': + whereValue = "''" + elif item['term'] == 'not_empty': + whereValue = "''" + elif item['term'] == 'in' or item['term'] == 'not in': + if ds.type == 'sqlServer' and ( + field.field_type == 'nchar' or field.field_type == 'NCHAR' or field.field_type == 'nvarchar' or field.field_type == 'NVARCHAR'): + whereValue = "(N'" + "', N'".join(value.split(",")) + "')" + else: + whereValue = "('" + "', '".join(value.split(",")) + "')" + elif item['term'] == 'like' or item['term'] == 'not like': + if ds.type == 'sqlServer' and ( + field.field_type == 'nchar' or field.field_type == 'NCHAR' or field.field_type == 'nvarchar' or field.field_type == 'NVARCHAR'): + whereValue = f"N'%{value}%'" + else: + whereValue = f"'%{value}%'" + else: + if ds.type == 'sqlServer' and ( + field.field_type == 'nchar' or field.field_type == 'NCHAR' or field.field_type == 'nvarchar' or field.field_type == 'NVARCHAR'): + whereValue = f"N'{value}'" + else: + whereValue = f"'{value}'" + + res = whereName + whereTerm + whereValue + return res + + +def transFilterTerm(term: str) -> str: + if term == "eq": + return " = " + if term == "not_eq": + return " <> " + if term == "lt": + return " < " + if term == "le": + return " <= " + if term == "gt": + return " > " + if term == "ge": + return " >= " + if term == "in": + return " IN " + if term == "not in": + return " NOT IN " + if term == "like": + return " LIKE " + if term == "not like": + return " NOT LIKE " + if term == "null": + return " IS NULL " + if term == "not_null": + return " IS NOT NULL " + if term == "empty": + return " = " + if term == "not_empty": + return " <> " + if term == "between": + return " BETWEEN " + return "" diff --git a/backend/apps/datasource/models/datasource.py b/backend/apps/datasource/models/datasource.py index f84aa691..ab73eaa3 100644 --- a/backend/apps/datasource/models/datasource.py +++ b/backend/apps/datasource/models/datasource.py @@ -17,6 +17,7 @@ class CoreDatasource(SQLModel, table=True): create_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True)) create_by: int = Field(sa_column=Column(BigInteger())) status: str = Field(max_length=64, nullable=True) + num: str = Field(max_length=256, nullable=True) class CoreTable(SQLModel, table=True): @@ -51,6 +52,7 @@ class CreateDatasource(BaseModel): create_time: Optional[datetime] = None create_by: int = 0 status: str = '' + num: str = '' tables: List[CoreTable] = [] diff --git a/backend/apps/template/filter/__init__.py b/backend/apps/template/filter/__init__.py new file mode 100644 index 00000000..e9302324 --- /dev/null +++ b/backend/apps/template/filter/__init__.py @@ -0,0 +1,2 @@ +# Author: Junjun +# Date: 2025/7/14 diff --git a/backend/apps/template/filter/generator.py b/backend/apps/template/filter/generator.py new file mode 100644 index 00000000..521601a4 --- /dev/null +++ b/backend/apps/template/filter/generator.py @@ -0,0 +1,6 @@ +from apps.template.template import get_base_template + + +def get_permissions_template(): + template = get_base_template() + return template['template']['permissions'] diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 056d88fb..27a7051f 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [ "pyyaml (>=6.0.2,<7.0.0)", "fastapi-mcp (>=0.3.4,<0.4.0)", "tabulate>=0.9.0", - "sqlbot-xpack==0.0.3.2", + "sqlbot-xpack==0.0.3.3", "fastapi-cache2>=0.2.2", ] [[tool.uv.index]]