From ca2f853bf1272d2105372aabd93ea8fdaf0a2e3b Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Mon, 4 Aug 2025 12:37:37 +0800 Subject: [PATCH] perf: X-Pack controlled by license --- backend/apps/chat/task/llm.py | 38 ++-------- backend/apps/datasource/crud/datasource.py | 75 ++++--------------- backend/apps/datasource/crud/permission.py | 66 ++++++++++++++++ .../apps/datasource/crud/row_permission.py | 7 +- backend/main.py | 4 +- 5 files changed, 90 insertions(+), 100 deletions(-) create mode 100644 backend/apps/datasource/crud/permission.py diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index b07b0e6c..d65688b9 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -13,13 +13,8 @@ import sqlparse 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, or_ from sqlalchemy import select -from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import sessionmaker -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 sqlmodel import create_engine, Session from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_default_config @@ -30,9 +25,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.crud.datasource import is_normal_user -from apps.datasource.crud.row_permission import transFilterTree -from apps.datasource.models.datasource import CoreDatasource, CoreTable +from apps.datasource.crud.permission import get_row_permission_filters, is_normal_user +from apps.datasource.models.datasource import CoreDatasource from apps.db.db import exec_sql from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds from apps.system.schemas.system_schema import AssistantOutDsSchema @@ -606,28 +600,9 @@ class LLMService: return full_filter_text 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)), - or_(DsRules.user_list.op('@>')(cast([f'{self.current_user.id}'], JSONB)), - DsRules.user_list.op('@>')(cast([self.current_user.id], JSONB)))) - ).first() - if obj is not None: - res.append(transRecord2DTO(self.session, permission)) - where_str = transFilterTree(self.session, res, self.ds) - filters.append({"table": table.table_name, "filter": where_str}) - + filters = get_row_permission_filters(session=self.session, current_user=self.current_user, ds=self.ds, tables=tables) + if not filters: + return None return self.build_table_filter(sql=sql, filters=filters) def generate_assistant_filter(self, sql, tables: List): @@ -864,7 +839,8 @@ class LLMService: # todo row permission if (not self.current_assistant and is_normal_user(self.current_user)) or use_dynamic_ds: sql, tables = self.check_sql(res=full_sql_text) - + sql_result = None + dynamic_sql_result = None if self.current_assistant: dynamic_sql_result = self.generate_assistant_dynamic_sql(sql, tables) if dynamic_sql_result: diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 28b00829..e4309ffd 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -3,14 +3,11 @@ import json from typing import List, Optional from fastapi import HTTPException -from sqlalchemy import and_, text, cast, or_, func -from sqlalchemy.dialects.postgresql import JSONB -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 sqlalchemy import and_, text, func + from sqlmodel import select -from apps.datasource.crud.row_permission import transFilterTree + from apps.datasource.utils.utils import aes_decrypt from apps.db.constant import DB from apps.db.db import get_engine, get_tables, get_fields, exec_sql @@ -23,7 +20,7 @@ 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, \ DatasourceConf, TableAndFields - +from apps.datasource.crud.permission import get_column_permission_fields, get_row_permission_filters, is_normal_user def get_datasource_list(session: SessionDep, user: CurrentUser, oid: Optional[int] = None) -> List[CoreDatasource]: current_oid = user.oid if user.oid is not None else 1 @@ -252,35 +249,14 @@ def preview(session: SessionDep, current_user: CurrentUser, id: int, data: Table f_list = [f for f in data.fields if f.checked] if is_normal_user(current_user): # column is checked, and, column permission for data.fields - column_permissions = session.query(DsPermission).filter( - and_(DsPermission.table_id == data.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)), - or_(DsRules.user_list.op('@>')(cast([f'{current_user.id}'], JSONB)), - DsRules.user_list.op('@>')(cast([current_user.id], JSONB)))) - ).first() - if obj is not None: - permission_list = json.loads(permission.permissions) - f_list = filter_list(f_list, permission_list) + f_list = get_column_permission_fields(session=session, current_user=current_user, table=data.table, fields=f_list) or f_list # row permission tree - row_permissions = session.query(DsPermission).filter( - and_(DsPermission.table_id == data.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 = session.query(DsRules).filter( - and_(DsRules.permission_list.op('@>')(cast([permission.id], JSONB)), - or_(DsRules.user_list.op('@>')(cast([f'{current_user.id}'], JSONB)), - DsRules.user_list.op('@>')(cast([current_user.id], JSONB)))) - ).first() - if obj is not None: - res.append(transRecord2DTO(session, permission)) - where_str = transFilterTree(session, res, ds) + where_str = '' + filter_mapping = get_row_permission_filters(session=session, current_user=current_user, ds=ds, tables=None, single_table=data.table) + if filter_mapping: + mapping_dict = filter_mapping[0] + where_str = mapping_dict.get('filter') where = (' where ' + where_str) if where_str is not None and where_str != '' else '' fields = [f.field_name for f in f_list] @@ -349,25 +325,13 @@ def get_table_obj_by_ds(session: SessionDep, current_user: CurrentUser, ds: Core fields = session.query(CoreField).filter(and_(CoreField.table_id == table.id, CoreField.checked == True)).all() # do column permissions, filter fields - if is_normal_user(current_user): - 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)), - or_(DsRules.user_list.op('@>')(cast([f'{current_user.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) - + fields = get_column_permission_fields(session=session, current_user=current_user, table=table, fields=fields) or fields _list.append(TableAndFields(schema=schema, table=table, fields=fields)) return _list + + def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource) -> str: schema_str = "" table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds) @@ -397,16 +361,3 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat 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)] - - -def is_normal_user(current_user: CurrentUser): - return current_user.id != 1 diff --git a/backend/apps/datasource/crud/permission.py b/backend/apps/datasource/crud/permission.py new file mode 100644 index 00000000..e3f3c3d8 --- /dev/null +++ b/backend/apps/datasource/crud/permission.py @@ -0,0 +1,66 @@ + +import json +from typing import List, Optional + +from sqlalchemy import and_, cast, or_ +from apps.datasource.crud.row_permission import transFilterTree +from apps.datasource.models.datasource import CoreDatasource, CoreField, CoreTable +from common.core.deps import CurrentUser, SessionDep +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 sqlalchemy.dialects.postgresql import JSONB + +def get_row_permission_filters(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource, tables: Optional[list] = None, single_table: Optional[CoreTable] = None): + if single_table: + table_list = [session.get(CoreTable, single_table.id)] + else: + table_list =session.query(CoreTable).filter( + and_(CoreTable.ds_id == ds.id, CoreTable.table_name.in_(tables)) + ).all() + + filters = [] + for table in table_list: + row_permissions = 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 = session.query(DsRules).filter( + and_(DsRules.permission_list.op('@>')(cast([permission.id], JSONB)), + or_(DsRules.user_list.op('@>')(cast([f'{current_user.id}'], JSONB)), + DsRules.user_list.op('@>')(cast([current_user.id], JSONB)))) + ).first() + if obj is not None: + res.append(transRecord2DTO(session, permission)) + where_str = transFilterTree(session, res, ds) + filters.append({"table": table.table_name, "filter": where_str}) + return filters + +def get_column_permission_fields(session: SessionDep, current_user: CurrentUser, table: CoreTable, fields: list[CoreField]): + if is_normal_user(current_user): + 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)), + or_(DsRules.user_list.op('@>')(cast([f'{current_user.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) + return fields +def is_normal_user(current_user: CurrentUser): + return current_user.id != 1 + +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)] \ No newline at end of file diff --git a/backend/apps/datasource/crud/row_permission.py b/backend/apps/datasource/crud/row_permission.py index 84f79cca..16f97971 100644 --- a/backend/apps/datasource/crud/row_permission.py +++ b/backend/apps/datasource/crud/row_permission.py @@ -2,15 +2,12 @@ # Date: 2025/6/25 from typing import List, Dict - -from sqlbot_xpack.permissions.models.ds_permission import PermissionTree, PermissionDTO - from apps.datasource.models.datasource import CoreField, CoreDatasource from apps.db.constant import DB from common.core.deps import SessionDep -def transFilterTree(session: SessionDep, tree_list: List[PermissionDTO], ds: CoreDatasource) -> str | None: +def transFilterTree(session: SessionDep, tree_list: List[any], ds: CoreDatasource) -> str | None: if tree_list is None: return None res: List[str] = [] @@ -24,7 +21,7 @@ def transFilterTree(session: SessionDep, tree_list: List[PermissionDTO], ds: Cor return " AND ".join(res) -def transTreeToWhere(session: SessionDep, tree: PermissionTree, ds: CoreDatasource) -> str | None: +def transTreeToWhere(session: SessionDep, tree: any, ds: CoreDatasource) -> str | None: if tree is None: return None logic = tree['logic'] diff --git a/backend/main.py b/backend/main.py index d578213a..e72c915d 100644 --- a/backend/main.py +++ b/backend/main.py @@ -3,7 +3,7 @@ from fastapi import FastAPI from fastapi.routing import APIRoute from starlette.middleware.cors import CORSMiddleware from starlette.exceptions import HTTPException as StarletteHTTPException -from apps.api import api_router + from apps.system.crud.assistant import init_dynamic_cors from apps.system.middleware.auth import TokenMiddleware from common.core.config import settings @@ -13,9 +13,9 @@ from alembic import command from fastapi_mcp import FastApiMCP from fastapi.staticfiles import StaticFiles import sqlbot_xpack - from common.utils.utils import SQLBotLogUtil from common.core.sqlbot_cache import init_sqlbot_cache +from apps.api import api_router def run_migrations(): alembic_cfg = Config("alembic.ini")