feat: API permission control (#613)

Author: fit2cloud-chenyw

backend/apps/chat/api/chat.py

@@ -14,6 +14,7 @@
     format_json_data, format_json_list_data, get_chart_config, list_recent_questions
 from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion, AxisObj, QuickCommand
 from apps.chat.task.llm import LLMService
+from apps.system.schemas.permission import SqlbotPermission, require_permissions
 from common.core.deps import CurrentAssistant, SessionDep, CurrentUser, Trans
 from common.utils.command_utils import parse_quick_command
 from common.utils.data_format import DataFormat
@@ -87,6 +88,7 @@ async def delete(session: SessionDep, chart_id: int):
 
 
 @router.post("/start")
+@require_permissions(permission=SqlbotPermission(type='ds', keyExpression="create_chat_obj.datasource"))
 async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat):
     try:
         return create_chat(session, current_user, create_chat_obj)
@@ -138,6 +140,7 @@ def _err(_e: Exception):
 
 
 @router.get("/recent_questions/{datasource_id}")
+@require_permissions(permission=SqlbotPermission(type='ds', keyExpression="datasource_id"))
 async def recommend_questions(session: SessionDep, current_user: CurrentUser, datasource_id: int):
     return list_recent_questions(session=session, current_user=current_user, datasource_id=datasource_id)
 
@@ -156,6 +159,7 @@ def find_base_question(record_id: int, session: SessionDep):
 
 
 @router.post("/question")
+@require_permissions(permission=SqlbotPermission(type='chat', keyExpression="request_question.chat_id"))
 async def question_answer(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion,
                           current_assistant: CurrentAssistant):
     try:

backend/apps/datasource/api/datasource.py

@@ -13,6 +13,7 @@
 from apps.db.db import get_schema
 from apps.db.engine import get_engine_conn
 from apps.swagger.i18n import PLACEHOLDER_PREFIX
+from apps.system.schemas.permission import SqlbotPermission, require_permissions
 from common.core.config import settings
 from common.core.deps import SessionDep, CurrentUser, Trans
 from common.utils.utils import SQLBotLogUtil
@@ -80,7 +81,9 @@ def inner():
     await asyncio.to_thread(inner)
 
 
+
 @router.post("/update", response_model=CoreDatasource, summary=f"{PLACEHOLDER_PREFIX}ds_update")
+@require_permissions(permission=SqlbotPermission(type='ds', keyExpression="ds.id"))
 async def update(session: SessionDep, trans: Trans, user: CurrentUser, ds: CoreDatasource):
     def inner():
         return update_ds(session, trans, user, ds)
@@ -89,11 +92,13 @@ def inner():
 
 
 @router.post("/delete/{id}", response_model=None, summary=f"{PLACEHOLDER_PREFIX}ds_delete")
+@require_permissions(permission=SqlbotPermission(type='ds', keyExpression="id"))
 async def delete(session: SessionDep, id: int = Path(..., description=f"{PLACEHOLDER_PREFIX}ds_id")):
     return delete_ds(session, id)
 
 
 @router.post("/getTables/{id}", response_model=List[TableSchemaResponse], summary=f"{PLACEHOLDER_PREFIX}ds_get_tables")
+@require_permissions(permission=SqlbotPermission(type='ds', keyExpression="id"))
 async def get_tables(session: SessionDep, id: int = Path(..., description=f"{PLACEHOLDER_PREFIX}ds_id")):
     return getTables(session, id)
 

backend/apps/system/api/aimodel.py

@@ -8,6 +8,7 @@
 from sqlmodel import func, select, update
 
 from apps.system.models.system_model import AiModelDetail
+from apps.system.schemas.permission import SqlbotPermission, require_permissions
 from common.core.deps import SessionDep, Trans
 from common.utils.crypto import sqlbot_decrypt
 from common.utils.time import get_timestamp
@@ -51,6 +52,7 @@ async def check_default(session: SessionDep, trans: Trans):
         raise Exception(trans('i18n_llm.miss_default'))
 
 @router.put("/default/{id}")
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def set_default(session: SessionDep, id: int):
     db_model = session.get(AiModelDetail, id)
     if not db_model:
@@ -70,6 +72,7 @@ async def set_default(session: SessionDep, id: int):
         raise e
 
 @router.get("", response_model=list[AiModelGridItem])
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def query(
         session: SessionDep,
         keyword: Union[str, None] = Query(default=None, max_length=255)
@@ -113,6 +116,7 @@ async def get_model_by_id(
     return AiModelEditor(**data)
 
 @router.post("")
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def add_model(
         session: SessionDep,
         creator: AiModelCreator
@@ -129,6 +133,7 @@ async def add_model(
     session.commit()
 
 @router.put("")
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def update_model(
         session: SessionDep,
         editor: AiModelEditor
@@ -144,6 +149,7 @@ async def update_model(
     session.commit()
 
 @router.delete("/{id}")
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def delete_model(
         session: SessionDep,
         trans: Trans,

backend/apps/system/api/parameter.py

@@ -4,6 +4,7 @@
 
 
 from apps.system.crud.parameter_manage import get_groups, get_parameter_args, save_parameter_args
+from apps.system.schemas.permission import SqlbotPermission, require_permissions
 from common.core.deps import SessionDep
 
 router = APIRouter(tags=["system/parameter"], prefix="/system/parameter")
@@ -13,9 +14,11 @@ async def get_login_args(session: SessionDep) -> list[SysArgModel]:
     return await get_groups(session, "login")
 
 @router.get("")
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def get_args(session: SessionDep) -> list[SysArgModel]:
     return await get_parameter_args(session)
 
 @router.post("", )
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def save_args(session: SessionDep, request: Request):
     return await save_parameter_args(session = session, request = request)

backend/apps/system/api/user.py

@@ -6,6 +6,7 @@
 from apps.system.models.system_model import UserWsModel, WorkspaceModel
 from apps.system.models.user import UserModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
+from apps.system.schemas.permission import SqlbotPermission, require_permissions
 from apps.system.schemas.system_schema import PwdEditor, UserCreator, UserEditor, UserGrid, UserLanguage, UserStatus, UserWs
 from common.core.deps import CurrentUser, SessionDep, Trans
 from common.core.pagination import Paginator
@@ -20,11 +21,14 @@
 async def user_info(current_user: CurrentUser):
     return current_user
 
+ 
 @router.get("/defaultPwd")
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 async def default_pwd() -> str:
     return settings.DEFAULT_PWD
 
 @router.get("/pager/{pageNum}/{pageSize}", response_model=PaginatedResponse[UserGrid])
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 async def pager(
     session: SessionDep,
     pageNum: int,
@@ -123,6 +127,7 @@ async def ws_change(session: SessionDep, current_user: CurrentUser, trans:Trans,
     session.commit()
 
 @router.get("/{id}", response_model=UserEditor)
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 async def query(session: SessionDep, trans: Trans, id: int) -> UserEditor:
     db_user: UserModel = get_db_user(session = session, user_id = id)
     u_ws_options = await user_ws_options(session, id, trans)
@@ -131,7 +136,9 @@ async def query(session: SessionDep, trans: Trans, id: int) -> UserEditor:
         result.oid_list = [item.id for item in u_ws_options]
     return result
 
+
 @router.post("")
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 async def create(session: SessionDep, creator: UserCreator, trans: Trans):
     if check_account_exists(session=session, account=creator.account):
         raise Exception(trans('i18n_exist', msg = f"{trans('i18n_user.account')} [{creator.account}]"))
@@ -158,8 +165,10 @@ async def create(session: SessionDep, creator: UserCreator, trans: Trans):
         user_model.oid = creator.oid_list[0]   
     session.add(user_model)
     session.commit()
+
 
 @router.put("")
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 @clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="editor.id")
 async def update(session: SessionDep, editor: UserEditor, trans: Trans):
     user_model: UserModel = get_db_user(session = session, user_id = editor.id)
@@ -193,12 +202,14 @@ async def update(session: SessionDep, editor: UserEditor, trans: Trans):
         user_model.oid = origin_oid if origin_oid in editor.oid_list else  editor.oid_list[0]
     session.add(user_model)
     session.commit()
-    
+
 @router.delete("/{id}")
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def delete(session: SessionDep, id: int):
     await single_delete(session, id)
 
-@router.delete("")    
+@router.delete("")
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 async def batch_del(session: SessionDep, id_list: list[int]):
     for id in id_list:
         await single_delete(session, id)
@@ -213,8 +224,10 @@ async def langChange(session: SessionDep, current_user: CurrentUser, trans: Tran
     db_user.language = lang
     session.add(db_user)
     session.commit()
-    
+
+   
 @router.patch("/pwd/{id}")
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 @clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")
 async def pwdReset(session: SessionDep, current_user: CurrentUser, trans: Trans, id: int):
     if not current_user.isAdmin:
@@ -236,8 +249,10 @@ async def pwdUpdate(session: SessionDep, current_user: CurrentUser, trans: Trans
     db_user.password = md5pwd(new_pwd)
     session.add(db_user)
     session.commit()
+
 
 @router.patch("/status")
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 @clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="statusDto.id")
 async def langChange(session: SessionDep, current_user: CurrentUser, trans: Trans, statusDto: UserStatus):
     if not current_user.isAdmin:

backend/apps/system/api/workspace.py

@@ -5,6 +5,7 @@
 from apps.system.crud.workspace import reset_single_user_oid, reset_user_oid
 from apps.system.models.system_model import UserWsModel, WorkspaceBase, WorkspaceEditor, WorkspaceModel
 from apps.system.models.user import UserModel
+from apps.system.schemas.permission import SqlbotPermission, require_permissions
 from apps.system.schemas.system_schema import UserWsBase, UserWsDTO, UserWsEditor, UserWsOption, WorkspaceUser
 from common.core.deps import CurrentUser, SessionDep, Trans
 from common.core.pagination import Paginator
@@ -14,6 +15,7 @@
 router = APIRouter(tags=["system/workspace"], prefix="/system/workspace")
 
 @router.get("/uws/option/pager/{pageNum}/{pageSize}", response_model=PaginatedResponse[UserWsOption])
+@require_permissions(permission=SqlbotPermission(role=['ws_admin']))
 async def option_pager(
     session: SessionDep,
     current_user: CurrentUser,
@@ -48,6 +50,7 @@ async def option_pager(
     )
 
 @router.get("/uws/option", response_model=UserWsOption | None)
+@require_permissions(permission=SqlbotPermission(role=['ws_admin'])) 
 async def option_user(
     session: SessionDep, 
     current_user: CurrentUser,
@@ -74,7 +77,9 @@ async def option_user(
         )
     return session.exec(stmt).first()
 
+
 @router.get("/uws/pager/{pageNum}/{pageSize}", response_model=PaginatedResponse[WorkspaceUser])
+@require_permissions(permission=SqlbotPermission(role=['ws_admin'])) 
 async def pager(
     session: SessionDep,
     current_user: CurrentUser,
@@ -114,7 +119,8 @@ async def pager(
     )
 
 
-@router.post("/uws")     
+@router.post("/uws")
+@require_permissions(permission=SqlbotPermission(role=['ws_admin']))     
 async def create(session: SessionDep, current_user: CurrentUser, trans: Trans, creator: UserWsDTO):
     if not current_user.isAdmin and current_user.weight == 0:
         raise Exception(trans('i18n_permission.no_permission', url = '', msg = ''))
@@ -136,7 +142,8 @@ async def create(session: SessionDep, current_user: CurrentUser, trans: Trans, c
     session.add_all(db_model_list)
     session.commit()
 
-@router.put("/uws")     
+@router.put("/uws")
+@require_permissions(permission=SqlbotPermission(role=['admin']))     
 async def edit(session: SessionDep, trans: Trans, editor: UserWsEditor):
     if not editor.oid or not editor.uid:
         raise Exception(trans('i18n_miss_args', key = '[oid, uid]'))
@@ -152,7 +159,8 @@ async def edit(session: SessionDep, trans: Trans, editor: UserWsEditor):
     await clean_user_cache(editor.uid)
     session.commit()
 
-@router.delete("/uws")     
+@router.delete("/uws")
+@require_permissions(permission=SqlbotPermission(role=['ws_admin']))     
 async def delete(session: SessionDep, current_user: CurrentUser, trans: Trans, dto: UserWsBase):
     if not current_user.isAdmin and current_user.weight == 0:
         raise Exception(trans('i18n_permission.no_permission', url = '', msg = ''))
@@ -170,6 +178,7 @@ async def delete(session: SessionDep, current_user: CurrentUser, trans: Trans, d
     session.commit()
 
 @router.get("", response_model=list[WorkspaceModel])
+@require_permissions(permission=SqlbotPermission(role=['admin'])) 
 async def query(session: SessionDep, trans: Trans):
     list_result = session.exec(select(WorkspaceModel)).all()
     for ws in list_result:
@@ -179,13 +188,15 @@ async def query(session: SessionDep, trans: Trans):
     return list_result
 
 @router.post("")
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 async def add(session: SessionDep, creator: WorkspaceBase):
     db_model = WorkspaceModel.model_validate(creator)
     db_model.create_time = get_timestamp()
     session.add(db_model)
     session.commit()
 
 @router.put("")
+@require_permissions(permission=SqlbotPermission(role=['admin']))
 async def update(session: SessionDep, editor: WorkspaceEditor):
     id = editor.id
     db_model = session.get(WorkspaceModel, id)
@@ -195,7 +206,8 @@ async def update(session: SessionDep, editor: WorkspaceEditor):
     session.add(db_model)
     session.commit()
 
-@router.get("/{id}", response_model=WorkspaceModel)    
+@router.get("/{id}", response_model=WorkspaceModel)
+@require_permissions(permission=SqlbotPermission(role=['admin']))    
 async def get_one(session: SessionDep, trans: Trans, id: int):
     db_model = session.get(WorkspaceModel, id)
     if not db_model:
@@ -204,7 +216,8 @@ async def get_one(session: SessionDep, trans: Trans, id: int):
         db_model.name = trans(db_model.name)
     return db_model
 
-@router.delete("/{id}")  
+@router.delete("/{id}")
+@require_permissions(permission=SqlbotPermission(role=['admin']))  
 async def single_delete(session: SessionDep, current_user: CurrentUser, id: int):
     if not current_user.isAdmin:
         raise HTTPException("only admin can delete workspace")