diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 1062c36d..fe0dd7bd 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -28,6 +28,7 @@ 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.db.db import exec_sql @@ -791,26 +792,28 @@ class LLMService: SQLBotLogUtil.info(full_sql_text) # todo row permission - sql_json_str = extract_nested_json(full_sql_text) - data = orjson.loads(sql_json_str) + 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'] + 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) diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 9a7e8a6f..75edb427 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -5,10 +5,12 @@ from typing import List 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) diff --git a/backend/apps/datasource/crud/row_permission.py b/backend/apps/datasource/crud/row_permission.py index b9b84558..84f79cca 100644 --- a/backend/apps/datasource/crud/row_permission.py +++ b/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 ( diff --git a/backend/apps/system/api/user.py b/backend/apps/system/api/user.py index d5a7a1db..cb4e9fdb 100644 --- a/backend/apps/system/api/user.py +++ b/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): diff --git a/backend/apps/system/crud/user.py b/backend/apps/system/crud/user.py index ffad38a8..fa0166e8 100644 --- a/backend/apps/system/crud/user.py +++ b/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,7 +58,7 @@ 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) @@ -66,6 +66,6 @@ async def single_delete(session: SessionDep, id: int): 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") diff --git a/backend/common/core/sqlbot_cache.py b/backend/common/core/sqlbot_cache.py index e13bc607..e3fe8482 100644 --- a/backend/common/core/sqlbot_cache.py +++ b/backend/common/core/sqlbot_cache.py @@ -124,6 +124,8 @@ def cache( 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)") diff --git a/frontend/src/components/layout/Workspace.vue b/frontend/src/components/layout/Workspace.vue index 6c4d34e5..da669173 100644 --- a/frontend/src/components/layout/Workspace.vue +++ b/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) diff --git a/frontend/src/router/index.ts b/frontend/src/router/index.ts index 8a75f65d..fb34662f 100644 --- a/frontend/src/router/index.ts +++ b/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, diff --git a/frontend/src/router/watch.ts b/frontend/src/router/watch.ts index 14bf2ad9..b480d2ca 100644 --- a/frontend/src/router/watch.ts +++ b/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() diff --git a/frontend/src/style.less b/frontend/src/style.less index c0d125ef..685e6f51 100644 --- a/frontend/src/style.less +++ b/frontend/src/style.less @@ -164,6 +164,10 @@ body { .tip { margin-top: 24px; } + + .ed-message-box__message { + width: 95%; + } } .form-content_error { diff --git a/frontend/src/views/chat/ChatCreator.vue b/frontend/src/views/chat/ChatCreator.vue index 68fac999..a46dae2b 100644 --- a/frontend/src/views/chat/ChatCreator.vue +++ b/frontend/src/views/chat/ChatCreator.vue @@ -1,8 +1,13 @@ -