feat: data permission

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 ###

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 ###

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

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'

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

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 ""

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] = []
 
 

backend/apps/template/filter/__init__.py

@@ -0,0 +1,2 @@
+# Author: Junjun
+# Date: 2025/7/14

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

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