perf: X-Pack controlled by license

Author: fit2cloud-chenyw

backend/apps/chat/task/llm.py

@@ -13,13 +13,8 @@
 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 @@
     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 @@ def build_table_filter(self, sql: str, filters: list):
         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 @@ def run_task(self, in_chat: bool = True):
             # 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:

backend/apps/datasource/crud/datasource.py

@@ -3,14 +3,11 @@
 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.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

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)]

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']

backend/main.py

@@ -3,7 +3,7 @@
 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 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")