diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 679717cb..59333281 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -834,6 +834,7 @@ class LLMService: # execute sql result = self.execute_sql(sql=sql) + print(result) self.save_sql_data(data_obj=result) if in_chat: yield orjson.dumps({'content': 'execute-success', 'type': 'sql-data'}).decode() + '\n\n' diff --git a/backend/apps/system/api/login.py b/backend/apps/system/api/login.py index 0aca2934..e2fd0bb2 100644 --- a/backend/apps/system/api/login.py +++ b/backend/apps/system/api/login.py @@ -21,6 +21,8 @@ def local_login( if not user.oid or user.oid == 0: raise HTTPException(status_code=400, detail="No associated workspace, Please contact the administrator") + if user.status != 1: + raise HTTPException(status_code=400, detail="User is disabled, Please contact the administrator") access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) user_dict = user.to_dict() return Token(access_token=create_access_token( diff --git a/backend/apps/system/api/user.py b/backend/apps/system/api/user.py index cb4e9fdb..7584d1db 100644 --- a/backend/apps/system/api/user.py +++ b/backend/apps/system/api/user.py @@ -1,7 +1,7 @@ from typing import Optional from fastapi import APIRouter, HTTPException, Query from sqlmodel import func, or_, select, delete as sqlmodel_delete -from apps.system.crud.user import get_db_user, single_delete, user_ws_options +from apps.system.crud.user import check_account_exists, get_db_user, single_delete, user_ws_options from apps.system.models.system_model import UserWsModel from apps.system.models.user import UserModel from apps.system.schemas.auth import CacheName, CacheNamespace @@ -116,6 +116,9 @@ async def query(session: SessionDep, trans: Trans, id: int) -> UserEditor: @router.post("") async def create(session: SessionDep, creator: UserCreator): + count = check_account_exists(session=session, account=creator.account) + if count > 0: + raise Exception(f"Account [{creator.account}] already exists!") data = creator.model_dump(exclude_unset=True) user_model = UserModel.model_validate(data) #user_model.create_time = get_timestamp() @@ -140,6 +143,10 @@ async def create(session: SessionDep, creator: UserCreator): @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) + if not user_model: + raise Exception(f"User with id [{editor.id}] not found!") + if editor.account != user_model.account: + raise Exception(f"account cannot be changed!") origin_oid: int = user_model.oid del_stmt = sqlmodel_delete(UserWsModel).where(UserWsModel.uid == editor.id) session.exec(del_stmt) diff --git a/backend/apps/system/crud/user.py b/backend/apps/system/crud/user.py index fa0166e8..ea6955bc 100644 --- a/backend/apps/system/crud/user.py +++ b/backend/apps/system/crud/user.py @@ -1,6 +1,6 @@ from typing import Optional -from sqlmodel import Session, select, delete as sqlmodel_delete +from sqlmodel import Session, func, select, delete as sqlmodel_delete from apps.system.models.system_model import UserWsModel, WorkspaceModel from apps.system.schemas.auth import CacheName, CacheNamespace from apps.system.schemas.system_schema import BaseUserDTO, UserInfoDTO, UserWs @@ -69,3 +69,7 @@ async def single_delete(session: SessionDep, id: int): @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") + + +def check_account_exists(*, session: Session, account: str) -> bool: + return session.exec(select(func.count()).select_from(UserModel).where(UserModel.account == account)).one() \ No newline at end of file diff --git a/backend/apps/system/schemas/system_schema.py b/backend/apps/system/schemas/system_schema.py index e6aa54b1..cb0c757f 100644 --- a/backend/apps/system/schemas/system_schema.py +++ b/backend/apps/system/schemas/system_schema.py @@ -19,6 +19,7 @@ class BaseUser(BaseModel): class BaseUserDTO(BaseUser, BaseCreatorDTO): language: str password: str + status: int = 1 def to_dict(self): return { "id": self.id,