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

Author: dataeaseShu

backend/apps/datasource/api/datasource.py

@@ -64,7 +64,7 @@ async def get_tables_by_conf(session: SessionDep, ds: CoreDatasource):
 
 
 @router.post("/getFields/{id}/{table_name}")
-async def get_tables(session: SessionDep, id: int, table_name: str):
+async def get_fields(session: SessionDep, id: int, table_name: str):
     return getFields(session, id, table_name)
 
 

backend/apps/datasource/crud/datasource.py

@@ -263,7 +263,7 @@ def fieldEnum(session: SessionDep, id: int):
 
 
 def updateNum(session: SessionDep, ds: CoreDatasource):
-    all_tables = get_tables(ds)
+    all_tables = get_tables(ds) if ds.type != 'excel' else json.loads(aes_decrypt(ds.configuration)).get('sheets')
     selected_tables = get_tables_by_ds_id(session, ds.id)
     num = f'{len(selected_tables)}/{len(all_tables)}'
 

backend/apps/system/api/login.py

@@ -1,6 +1,7 @@
 from typing import Annotated
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.security import OAuth2PasswordRequestForm
+from apps.system.schemas.system_schema import BaseUserDTO
 from common.core.deps import SessionDep
 from ..crud.user import authenticate
 from common.core.security import create_access_token
@@ -14,9 +15,12 @@ def local_login(
     session: SessionDep,
     form_data: Annotated[OAuth2PasswordRequestForm, Depends()]
 ) -> Token:
-    user = authenticate(session=session, account=form_data.username, password=form_data.password)
+    user: BaseUserDTO = authenticate(session=session, account=form_data.username, password=form_data.password)
     if not user:
         raise HTTPException(status_code=400, detail="Incorrect account or password")
+    
+    if not user.oid or user.oid == 0:
+        raise HTTPException(status_code=400, detail="No associated workspace, 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,8 +1,8 @@
 from typing import Optional
-from fastapi import APIRouter, Query
-from sqlmodel import func, or_, select
-from apps.system.crud.user import get_db_user, user_ws_options
-from apps.system.models.system_model import UserWsModel, WorkspaceModel
+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.models.system_model import UserWsModel
 from apps.system.models.user import UserModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
 from apps.system.schemas.system_schema import PwdEditor, UserCreator, UserEditor, UserGrid, UserLanguage, UserWs
@@ -34,14 +34,21 @@ async def pager(
     stmt = (
         select(
             UserModel,
-            func.coalesce(func.string_agg(WorkspaceModel.name, ','), '').label("space_name")
+            func.coalesce(
+                func.array_remove(
+                    func.array_agg(UserWsModel.oid),
+                    None
+                ),
+                []
+            ).label("oid_list")
+            #func.coalesce(func.string_agg(WorkspaceModel.name, ','), '').label("space_name")
         )
         .join(UserWsModel, UserModel.id == UserWsModel.uid, isouter=True)
-        .join(WorkspaceModel, UserWsModel.oid == WorkspaceModel.id, isouter=True)
+        #.join(WorkspaceModel, UserWsModel.oid == WorkspaceModel.id, isouter=True)
+        .where(UserModel.id != 1)
         .group_by(UserModel.id)
         .order_by(UserModel.create_time)
     )
-    
     if status is not None:
         stmt = stmt.where(UserModel.status == status)
 
@@ -67,10 +74,21 @@ async def pager(
             )
         )
 
-    return await paginator.get_paginated_response(
+    user_page = await paginator.get_paginated_response(
         stmt=stmt,
         pagination=pagination,
         **filters)
+    
+    """ for item in user_page.items:
+        space_name: str = item['space_name']
+        if space_name and 'i18n_default_workspace' in space_name:
+            parts = list(map(
+                lambda x: trans(x) if x == "i18n_default_workspace" else x,
+                space_name.split(',')
+            ))
+            output_str = ','.join(parts)
+            item['space_name'] = output_str """
+    return user_page
 
 @router.get("/ws")
 async def ws_options(session: SessionDep, current_user: CurrentUser, trans: Trans) -> list[UserWs]:
@@ -81,46 +99,78 @@ async def ws_options(session: SessionDep, current_user: CurrentUser, trans: Tran
 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):
-        raise RuntimeError(f"oid [{oid}] is invalid!")
+        raise HTTPException(f"oid [{oid}] is invalid!")
     user_model: UserModel = get_db_user(session = session, user_id = current_user.id)
     user_model.oid = oid
     session.add(user_model)
     session.commit()
 
 @router.get("/{id}", response_model=UserEditor)
-async def query(session: SessionDep, id: int) -> UserEditor:
+async def query(session: SessionDep, trans: Trans, id: int) -> UserEditor:
     db_user: UserModel = get_db_user(session = session, user_id = id)
-    return db_user
+    u_ws_options = await user_ws_options(session, id, trans)
+    result = UserEditor.model_validate(db_user.model_dump())
+    if u_ws_options:
+        result.oid_list = [item.id for item in u_ws_options]
+    return result
 
 @router.post("")
 async def create(session: SessionDep, creator: UserCreator):
     data = creator.model_dump(exclude_unset=True)
     user_model = UserModel.model_validate(data)
     #user_model.create_time = get_timestamp()
     user_model.language = "zh-CN"
+    user_model.oid = 0
+    if creator.oid_list:
+        # need to validate oid_list
+        db_model_list = [
+            UserWsModel.model_validate({
+                "oid": oid,
+                "uid": user_model.id,
+                "weight": 0
+            })
+            for oid in creator.oid_list
+        ]
+        session.add_all(db_model_list)
+        user_model.oid = creator.oid_list[0]   
     session.add(user_model)
     session.commit()
 
 @router.put("")
 @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
+    del_stmt = sqlmodel_delete(UserWsModel).where(UserWsModel.uid == editor.id)
+    session.exec(del_stmt)
+    
     data = editor.model_dump(exclude_unset=True)
     user_model.sqlmodel_update(data)
+    
+    user_model.oid = 0
+    if editor.oid_list:
+        # need to validate oid_list
+        db_model_list = [
+            UserWsModel.model_validate({
+                "oid": oid,
+                "uid": user_model.id,
+                "weight": 0
+            })
+            for oid in editor.oid_list
+        ]
+        session.add_all(db_model_list)
+        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}")
-@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")
 async def delete(session: SessionDep, id: int):
-    user_model: UserModel = get_db_user(session = session, user_id = id)
-    session.delete(user_model)
-    session.commit()
+    await single_delete(session, id)
 
 @router.delete("")    
 async def batch_del(session: SessionDep, id_list: list[int]):
     for id in id_list:
-        delete(session, id)
+        await single_delete(session, id)
 
 @router.put("/language")
 @clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
@@ -137,7 +187,7 @@ async def langChange(session: SessionDep, current_user: CurrentUser, language: U
 @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 RuntimeError('only for admin')
+        raise HTTPException('only for admin')
     db_user: UserModel = get_db_user(session=session, user_id=id)
     db_user.password = default_md5_pwd()
     session.add(db_user)
@@ -148,7 +198,7 @@ async def pwdReset(session: SessionDep, current_user: CurrentUser, id: int):
 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):
-        raise RuntimeError("pwd error")
+        raise HTTPException("pwd error")
     db_user.password = md5pwd(editor.new_pwd)
     session.add(db_user)
     session.commit()

backend/apps/system/crud/user.py

@@ -1,11 +1,12 @@
 
 from typing import Optional
-from sqlmodel import Session, select
+from sqlmodel import Session, 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
-from common.core.sqlbot_cache import cache
+from common.core.deps import SessionDep
+from common.core.sqlbot_cache import cache, clear_cache
 from common.utils.locale import I18n
 from ..models.user import UserModel
 from common.core.security import verify_md5pwd
@@ -55,4 +56,12 @@ async def user_ws_options(session: Session, uid: int, trans: Optional[I18n] = No
     return [
         UserWs(id = id, name = trans(name) if name.startswith('i18n') else name) 
         for id, name in result.all()
-    ]
+    ]
+    
+@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()

backend/apps/system/middleware/auth.py

@@ -1,6 +1,6 @@
 
 from typing import Optional
-from fastapi import Request
+from fastapi import HTTPException, Request
 from fastapi.responses import JSONResponse
 import jwt
 from sqlmodel import Session
@@ -35,15 +35,15 @@ async def dispatch(self, request, call_next):
                 request.state.current_user = validator[1]
                 request.state.assistant = validator[2]
                 return await call_next(request)
-            return JSONResponse({"error": f"Unauthorized:[{validator[1]}]"}, status_code=401)
+            return JSONResponse({"msg": f"Unauthorized:[{validator[1]}]"}, status_code=401)
         #validate pass
         tokenkey = settings.TOKEN_KEY
         token = request.headers.get(tokenkey)
         validate_pass, data = await self.validateToken(token)
         if validate_pass:
             request.state.current_user = data
             return await call_next(request)
-        return JSONResponse({"error": f"Unauthorized:[{data}]"}, status_code=401)
+        return JSONResponse({"msg": f"Unauthorized:[{data}]"}, status_code=401)
 
     def is_options(self, request: Request):
         return request.method == "OPTIONS"
@@ -62,6 +62,12 @@ async def validateToken(self, token: Optional[str]):
             with Session(engine) as session:
                 session_user = await get_user_info(session = session, user_id = token_data.id)
                 session_user = UserInfoDTO.model_validate(session_user)
+                session_user = UserInfoDTO.model_validate(session_user)
+                """ if token_data.oid != session_user.oid:
+                    raise HTTPException(
+                        status_code=401,
+                        detail="Default space has been changed, please login again!"
+                    ) """
                 return True, session_user
         except Exception as e:
             return False, e

backend/apps/system/models/user.py

@@ -10,7 +10,7 @@
 
 class BaseUserPO(SQLModel):
     account: str = Field(max_length=255, unique=True)
-    oid: int = Field(nullable=False, sa_type=BigInteger())
+    oid: int = Field(nullable=False, sa_type=BigInteger(), default=0)
     name: str = Field(max_length=255, unique=True)
     password: str = Field(default_factory=default_md5_pwd, max_length=255)
     email: str = Field(max_length=255)

backend/apps/system/schemas/system_schema.py

@@ -30,14 +30,15 @@ class UserCreator(BaseUser):
     name: str
     email: str
     status: int = 1
+    oid_list: Optional[list[int]] = None 
 
 class UserEditor(UserCreator, BaseCreatorDTO):
    pass
 
 class UserGrid(UserEditor):
     create_time: int
     language: str = "zh-CN"
-    space_name: Optional[str] = None
+    #space_name: Optional[str] = None
     origin: str = ''
 
 

frontend/src/api/chat.ts

@@ -24,6 +24,7 @@ export interface ChatMessage {
   isTyping?: boolean
   first_chat?: boolean
   recommended_question?: string
+  index: number
 }
 
 export class ChatRecord {