Merge remote-tracking branch 'origin/main'

Author: ziyujiahao

backend/apps/chat/task/llm.py

@@ -28,6 +28,7 @@
     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.db.db import exec_sql
@@ -791,26 +792,28 @@ def run_task(self, in_chat: bool = True):
             SQLBotLogUtil.info(full_sql_text)
 
             # todo row permission
-            sql_json_str = extract_nested_json(full_sql_text)
-            data = orjson.loads(sql_json_str)
-
-            sql = ''
-            message = ''
-            error = False
-            if data['success']:
-                sql = data['sql']
+            if is_normal_user(self.current_user):
+                sql_json_str = extract_nested_json(full_sql_text)
+                data = orjson.loads(sql_json_str)
+
+                sql = ''
+                message = ''
+                error = False
+                if data['success']:
+                    sql = data['sql']
+                else:
+                    message = data['message']
+                    error = True
+                if error:
+                    raise Exception(message)
+                if sql.strip() == '':
+                    raise Exception("SQL query is empty")
+
+                sql_result = self.generate_filter(data.get('sql'), data.get('tables'))  # maybe no sql and tables
+                SQLBotLogUtil.info(sql_result)
+                sql = self.check_save_sql(res=sql_result)
             else:
-                message = data['message']
-                error = True
-            if error:
-                raise Exception(message)
-            if sql.strip() == '':
-                raise Exception("SQL query is empty")
-
-            sql_result = self.generate_filter(data.get('sql'), data.get('tables'))  # maybe no sql and tables
-            SQLBotLogUtil.info(sql_result)
-            sql = self.check_save_sql(res=sql_result)
-            # sql = llm_service.check_save_sql(res=full_sql_text)
+                sql = self.check_save_sql(res=full_sql_text)
 
             SQLBotLogUtil.info(sql)
             format_sql = sqlparse.format(sql, reindent=True)

backend/apps/datasource/crud/datasource.py

@@ -5,10 +5,12 @@
 from fastapi import HTTPException
 from sqlalchemy import and_, text, cast, or_, func
 from sqlalchemy.dialects.postgresql import JSONB
-from sqlbot_xpack.permissions.models.ds_permission import DsPermission
+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 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
@@ -238,43 +240,69 @@ def updateField(session: SessionDep, field: CoreField):
 
 
 def preview(session: SessionDep, current_user: CurrentUser, id: int, data: TableObj):
+    ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
+    check_status(session, ds, True)
+
     if data.fields is None or len(data.fields) == 0:
         return {"fields": [], "data": [], "sql": ''}
 
-    # column is checked, and, column permission for data.fields
+    where = None
     f_list = [f for f in data.fields if f.checked]
-    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)
+    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)
+
+        # 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))
+        wheres = transFilterTree(session, res, ds)
+        where = (' where ' + wheres) if wheres is not None and wheres != '' else ''
 
     fields = [f.field_name for f in f_list]
     if fields is None or len(fields) == 0:
         return {"fields": [], "data": [], "sql": ''}
 
-    ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
-    check_status(session, ds, True)
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     sql: str = ""
     if ds.type == "mysql":
-        sql = f"""SELECT `{"`, `".join(fields)}` FROM `{data.table.table_name}` LIMIT 100"""
+        sql = f"""SELECT `{"`, `".join(fields)}` FROM `{data.table.table_name}` 
+            {where} 
+            LIMIT 100"""
     elif ds.type == "sqlServer":
         sql = f"""SELECT [{"], [".join(fields)}] FROM [{conf.dbSchema}].[{data.table.table_name}]
+            {where} 
             ORDER BY [{data.fields[0].field_name}]
             OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"""
     elif ds.type == "pg" or ds.type == "excel":
-        sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}" LIMIT 100"""
+        sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}" 
+            {where} 
+            LIMIT 100"""
     elif ds.type == "oracle":
         sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}"
+            {where} 
             ORDER BY "{data.fields[0].field_name}"
             OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"""
     return exec_sql(ds, sql)
@@ -320,19 +348,20 @@ 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
-        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)
+        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)
 
         _list.append(TableAndFields(schema=schema, table=table, fields=fields))
     return _list
@@ -376,3 +405,7 @@ def filter_list(list_a, list_b):
             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 and (current_user.weight is not None and current_user.weight != 1)

backend/apps/datasource/crud/row_permission.py

@@ -3,8 +3,10 @@
 
 from typing import List, Dict
 
-from apps.datasource.models.datasource import CoreField, CoreDatasource
 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
 
 
@@ -48,7 +50,8 @@ def transTreeItem(session: SessionDep, item: Dict, ds: CoreDatasource) -> str |
     if field is None:
         return None
 
-    whereName = field.field_name
+    db = DB.get_db(ds.type)
+    whereName = db.prefix + field.field_name + db.suffix
     if item['filter_type'] == 'enum':
         if len(item['enum_value']) > 0:
             if ds['type'] == 'sqlServer' and (

backend/apps/system/api/user.py

@@ -95,7 +95,7 @@ async def ws_options(session: SessionDep, current_user: CurrentUser, trans: Tran
     return await user_ws_options(session, current_user.id, trans)
 
 @router.put("/ws/{oid}")
-#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
+@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
 async def ws_change(session: SessionDep, current_user: CurrentUser, oid: int):
     ws_list: list[UserWs] = await user_ws_options(session, current_user.id)
     if not any(x.id == oid for x in ws_list):
@@ -137,7 +137,7 @@ async def create(session: SessionDep, creator: UserCreator):
     session.commit()
 
 @router.put("")
-#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="editor.id")
+@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="editor.id")
 async def update(session: SessionDep, editor: UserEditor):
     user_model: UserModel = get_db_user(session = session, user_id = editor.id)
     origin_oid: int = user_model.oid
@@ -173,7 +173,7 @@ async def batch_del(session: SessionDep, id_list: list[int]):
         await single_delete(session, id)
 
 @router.put("/language")
-#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
+@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
 async def langChange(session: SessionDep, current_user: CurrentUser, language: UserLanguage):
     lang = language.language
     if lang not in ["zh-CN", "en"]:
@@ -184,7 +184,7 @@ async def langChange(session: SessionDep, current_user: CurrentUser, language: U
     session.commit()
 
 @router.patch("/pwd/{id}")
-#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")
+@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")
 async def pwdReset(session: SessionDep, current_user: CurrentUser, id: int):
     if not current_user.isAdmin:
         raise HTTPException('only for admin')
@@ -194,7 +194,7 @@ async def pwdReset(session: SessionDep, current_user: CurrentUser, id: int):
     session.commit()
 
 @router.put("/pwd")
-#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
+@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
 async def pwdUpdate(session: SessionDep, current_user: CurrentUser, editor: PwdEditor):
     db_user: UserModel = get_db_user(session=session, user_id=current_user.id)
     if not verify_md5pwd(editor.pwd, db_user.password):

backend/apps/system/crud/user.py

@@ -22,7 +22,7 @@ def get_user_by_account(*, session: Session, account: str) -> BaseUserDTO | None
         return None
     return BaseUserDTO.model_validate(db_user.model_dump())
 
-#@cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="user_id")
+@cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="user_id")
 async def get_user_info(*, session: Session, user_id: int) -> UserInfoDTO | None:
     db_user: UserModel = get_db_user(session = session, user_id = user_id)
     userInfo = UserInfoDTO.model_validate(db_user.model_dump())
@@ -58,14 +58,14 @@ async def user_ws_options(session: Session, uid: int, trans: Optional[I18n] = No
         for id, name in result.all()
     ]
 
-#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")
+@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")
 async def single_delete(session: SessionDep, id: int):
     user_model: UserModel = get_db_user(session = session, user_id = id)
     del_stmt = sqlmodel_delete(UserWsModel).where(UserWsModel.uid == id)
     session.exec(del_stmt)
     session.delete(user_model)
     session.commit()
 
-#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")    
+@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")    
 async def clean_user_cache(id: int):
     SQLBotLogUtil.info(f"User cache for [{id}] has been cleaned")

backend/common/core/sqlbot_cache.py

@@ -124,6 +124,8 @@ async def wrapper(*args, **kwargs):
                 result = await func(*args, **kwargs)
 
                 actual_expire = expire + random.randint(-jitter, jitter)
+                if await backend.get(cache_key):
+                    await backend.clear(cache_key)
                 await backend.set(cache_key, result, actual_expire)
 
                 SQLBotLogUtil.debug(f"Cache set: {cache_key} (expire: {actual_expire}s)")

frontend/src/components/layout/Workspace.vue

@@ -6,13 +6,11 @@ import icon_done_outlined from '@/assets/svg/icon_done_outlined.svg'
 import icon_searchOutline_outlined from '@/assets/svg/icon_search-outline_outlined.svg'
 import { userApi } from '@/api/auth'
 import { ElMessage } from 'element-plus-secondary'
-import { useRouter } from 'vue-router'
 import { useI18n } from 'vue-i18n'
 import { useUserStore } from '@/stores/user'
 
 const userStore = useUserStore()
 const { t } = useI18n()
-const router = useRouter()
 defineProps({
   collapse: { type: [Boolean], required: true },
 })
@@ -46,7 +44,6 @@ const handleDefaultWorkspaceChange = (item: any) => {
   currentWorkspace.value = { id: item.id, name: item.name }
   userApi.ws_change(item.id).then(() => {
     ElMessage.success(t('common.switch_success'))
-    router.push('/chat/index')
     setTimeout(() => {
       location.reload()
     }, 300)

frontend/src/router/index.ts

@@ -114,38 +114,6 @@ const router = createRouter({
       component: DashboardPreview,
       meta: { title: 'DashboardPreview', icon: 'dashboard' },
     },
-    /* {
-      path: "/setting",
-      component: Layout,
-      redirect: "/setting/index",
-      children: [
-        {
-          path: "index",
-          name: "setting",
-          component: setting,
-          meta: { title: "Settings", icon: "setting" },
-        },
-      ],
-    }, */
-    // {
-    //   path: '/system',
-    //   component: Layout,
-    //   redirect: '/system/model',
-    //   children: [
-    //     /*  {
-    //       path: "user",
-    //       name: "user",
-    //       component: User,
-    //       meta: { title: "User Management", icon: "icon_user" },
-    //     }, */
-    //     {
-    //       path: 'model',
-    //       name: 'model',
-    //       component: Model,
-    //       meta: { title: 'AI Model Configuration', icon: 'icon_ai' },
-    //     },
-    //   ],
-    // },
     {
       path: '/system',
       component: LayoutDsl,

frontend/src/router/watch.ts

@@ -26,10 +26,7 @@ export const watchRouter = (router: any) => {
     if (!userStore.getUid) {
       await userStore.info()
     }
-    /* if (!wsCache.get('user.uid')) {
-      await userStore.info()
-    } */
-    if (to.path === '/') {
+    if (to.path === '/' || accessCrossPermission(to)) {
       next('/chat')
       return
     }
@@ -41,6 +38,14 @@ export const watchRouter = (router: any) => {
     }
   })
 }
+
+const accessCrossPermission = (to: any) => {
+  if (!to?.path) return false
+  return (
+    (to.path.startsWith('/system') && !userStore.isAdmin) ||
+    (to.path.startsWith('/set') && !userStore.isSpaceAdmin)
+  )
+}
 const loadXpackStatic = () => {
   if (document.getElementById('sqlbot_xpack_static')) {
     return Promise.resolve()