feat: data permission

Author: XiaJunjie2020

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_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 @@ def guess_user_question(self, old_questions: str = "[]"):
         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 @@
 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 @@
     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 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
 
         # 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 @@ def generate_sql(self):
                                                            [{'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 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.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)]