Merge branch 'main' of https://github.com/dataease/SQLBot

Author: dataeaseShu

backend/apps/chat/task/llm.py

@@ -834,6 +834,7 @@ def run_task(self, in_chat: bool = True):
 
             # 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'

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(

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)

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()

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,