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

Author: dataeaseShu

backend/alembic/versions/028_ds_oid.py

@@ -0,0 +1,28 @@
+"""028_ds_oid
+
+Revision ID: e96b16d3daab
+Revises: b049c9f8ca5b
+Create Date: 2025-07-17 14:40:48.522033
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'e96b16d3daab'
+down_revision = 'b049c9f8ca5b'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('core_datasource', sa.Column('oid', sa.BigInteger(), nullable=True))
+    op.execute('update core_datasource set oid = 1')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('core_datasource', 'oid')
+    # ### end Alembic commands ###

backend/alembic/versions/029_modify_chat.py

@@ -0,0 +1,54 @@
+"""029_modify_chat
+
+Revision ID: 77d4c39ec22f
+Revises: e96b16d3daab
+Create Date: 2025-07-17 17:05:13.392973
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '77d4c39ec22f'
+down_revision = 'e96b16d3daab'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('chat', sa.Column('oid', sa.BigInteger(), nullable=True))
+    op.execute('update chat set oid = 1')
+    op.alter_column('chat', 'create_time',
+               existing_type=postgresql.TIMESTAMP(timezone=True),
+               type_=sa.DateTime(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'create_time',
+               existing_type=postgresql.TIMESTAMP(timezone=True),
+               type_=sa.DateTime(),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'finish_time',
+               existing_type=postgresql.TIMESTAMP(timezone=True),
+               type_=sa.DateTime(),
+               existing_nullable=True)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('chat_record', 'finish_time',
+               existing_type=sa.DateTime(),
+               type_=postgresql.TIMESTAMP(timezone=True),
+               existing_nullable=True)
+    op.alter_column('chat_record', 'create_time',
+               existing_type=sa.DateTime(),
+               type_=postgresql.TIMESTAMP(timezone=True),
+               existing_nullable=True)
+    op.alter_column('chat', 'create_time',
+               existing_type=sa.DateTime(),
+               type_=postgresql.TIMESTAMP(timezone=True),
+               existing_nullable=True)
+    op.drop_column('chat', 'oid')
+    # ### end Alembic commands ###

backend/apps/chat/models/chat_model.py

@@ -17,7 +17,8 @@
 class Chat(SQLModel, table=True):
     __tablename__ = "chat"
     id: Optional[int] = Field(sa_column=Column(BigInteger, Identity(always=True), primary_key=True))
-    create_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
+    oid: Optional[int] = Field(sa_column=Column(BigInteger, nullable=True, default=1))
+    create_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True))
     create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
     brief: str = Field(max_length=64, nullable=True)
     chat_type: str = Field(max_length=20, default="chat")  # chat, datasource
@@ -31,8 +32,8 @@ class ChatRecord(SQLModel, table=True):
     chat_id: int = Field(sa_column=Column(BigInteger, nullable=False))
     ai_modal_id: Optional[int] = Field(sa_column=Column(BigInteger))
     first_chat: bool = Field(sa_column=Column(Boolean, nullable=True, default=False))
-    create_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
-    finish_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
+    create_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True))
+    finish_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True))
     create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
     datasource: int = Field(sa_column=Column(BigInteger, nullable=True))
     engine_type: str = Field(max_length=64)
@@ -66,6 +67,7 @@ class ChatRecord(SQLModel, table=True):
     analysis_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
     predict_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
 
+
 class CreateChat(BaseModel):
     id: int = None
     question: str = None

backend/apps/datasource/api/datasource.py

@@ -19,8 +19,8 @@
 
 
 @router.get("/list")
-async def datasource_list(session: SessionDep):
-    return get_datasource_list(session=session)
+async def datasource_list(session: SessionDep, user: CurrentUser):
+    return get_datasource_list(session=session, user=user)
 
 
 @router.post("/get/{id}")
@@ -44,8 +44,8 @@ async def choose_tables(session: SessionDep, id: int, tables: List[CoreTable]):
 
 
 @router.post("/update", response_model=CoreDatasource)
-async def update(session: SessionDep, ds: CoreDatasource):
-    return update_ds(session, ds)
+async def update(session: SessionDep,user: CurrentUser, ds: CoreDatasource):
+    return update_ds(session,user, ds)
 
 
 @router.post("/delete/{id}", response_model=CoreDatasource)

backend/apps/datasource/crud/datasource.py

@@ -22,10 +22,10 @@
     DatasourceConf, TableAndFields
 
 
-def get_datasource_list(session: SessionDep):
-    statement = select(CoreDatasource).order_by(CoreDatasource.create_time.desc())
-    datasource_list = session.exec(statement).fetchall()
-    return datasource_list
+def get_datasource_list(session: SessionDep, user: CurrentUser):
+    oid = user.oid if user.oid is not None else 1
+    return session.query(CoreDatasource).filter(CoreDatasource.oid == oid).order_by(
+        CoreDatasource.create_time.desc()).all()
 
 
 def get_ds(session: SessionDep, id: int):
@@ -45,25 +45,27 @@ def check_status(session: SessionDep, ds: CoreDatasource):
         return False
 
 
-def check_name(session: SessionDep, ds: CoreDatasource):
+def check_name(session: SessionDep, user: CurrentUser, ds: CoreDatasource):
     if ds.id is not None:
         ds_list = session.query(CoreDatasource).filter(
-            and_(CoreDatasource.name == ds.name, CoreDatasource.id != ds.id)).all()
+            and_(CoreDatasource.name == ds.name, CoreDatasource.id != ds.id, CoreDatasource.oid == user.oid)).all()
         if ds_list is not None and len(ds_list) > 0:
             raise 'Name exist'
     else:
-        ds_list = session.query(CoreDatasource).filter(CoreDatasource.name == ds.name).all()
+        ds_list = session.query(CoreDatasource).filter(
+            and_(CoreDatasource.name == ds.name, CoreDatasource.oid == user.oid)).all()
         if ds_list is not None and len(ds_list) > 0:
             raise 'Name exist'
 
 
 def create_ds(session: SessionDep, user: CurrentUser, create_ds: CreateDatasource):
     ds = CoreDatasource()
     deepcopy_ignore_extra(create_ds, ds)
-    check_name(session, ds)
+    check_name(session, user, ds)
     ds.create_time = datetime.datetime.now()
     # status = check_status(session, ds)
     ds.create_by = user.id
+    ds.oid = user.oid if user.oid is not None else 1
     ds.status = "Success"
     ds.type_name = db_type_relation()[ds.type]
     record = CoreDatasource(**ds.model_dump())
@@ -85,9 +87,9 @@ def chooseTables(session: SessionDep, id: int, tables: List[CoreTable]):
     updateNum(session, ds)
 
 
-def update_ds(session: SessionDep, ds: CoreDatasource):
+def update_ds(session: SessionDep, user: CurrentUser, ds: CoreDatasource):
     ds.id = int(ds.id)
-    check_name(session, ds)
+    check_name(session, user, ds)
     status = check_status(session, ds)
     ds.status = "Success" if status is True else "Fail"
     record = session.exec(select(CoreDatasource).where(CoreDatasource.id == ds.id)).first()

backend/apps/datasource/models/datasource.py

@@ -18,6 +18,7 @@ class CoreDatasource(SQLModel, table=True):
     create_by: int = Field(sa_column=Column(BigInteger()))
     status: str = Field(max_length=64, nullable=True)
     num: str = Field(max_length=256, nullable=True)
+    oid: int = Field(sa_column=Column(BigInteger()))
 
 
 class CoreTable(SQLModel, table=True):
@@ -53,6 +54,7 @@ class CreateDatasource(BaseModel):
     create_by: int = 0
     status: str = ''
     num: str = ''
+    oid: int = 1
     tables: List[CoreTable] = []
 
 

backend/apps/mcp/mcp.py

@@ -3,20 +3,29 @@
 
 from datetime import timedelta
 
-from fastapi import APIRouter, HTTPException
+import jwt
+from fastapi import HTTPException, status, APIRouter
 from fastapi.responses import StreamingResponse
+# from fastapi.security import OAuth2PasswordBearer
+from jwt.exceptions import InvalidTokenError
+from pydantic import ValidationError
 
 from apps.chat.api.chat import create_chat
 from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart
 from apps.chat.task.llm import LLMService, run_task
-from apps.datasource.crud.datasource import get_datasource_list
-from apps.system.crud.user import authenticate
-from apps.system.models.system_model import AiModelDetail
+from apps.system.crud.user import authenticate, get_user_info
+from apps.system.schemas.system_schema import BaseUserDTO
+from apps.system.schemas.system_schema import UserInfoDTO
+from common.core import security
 from common.core.config import settings
-from common.core.deps import SessionDep, get_current_user
-from common.core.schemas import Token
+from common.core.deps import SessionDep
+from common.core.schemas import TokenPayload, XOAuth2PasswordBearer, Token
 from common.core.security import create_access_token
 
+reusable_oauth2 = XOAuth2PasswordBearer(
+    tokenUrl=f"{settings.API_V1_STR}/login/access-token"
+)
+
 router = APIRouter(tags=["mcp"], prefix="/mcp")
 
 
@@ -35,21 +44,24 @@
 #     ))
 
 
-@router.get("/ds_list", operation_id="get_datasource_list")
-async def datasource_list(session: SessionDep):
-    return get_datasource_list(session=session)
-
-
-@router.get("/model_list", operation_id="get_model_list")
-async def get_model_list(session: SessionDep):
-    return session.query(AiModelDetail).all()
+# @router.get("/ds_list", operation_id="get_datasource_list")
+# async def datasource_list(session: SessionDep):
+#     return get_datasource_list(session=session)
+#
+#
+# @router.get("/model_list", operation_id="get_model_list")
+# async def get_model_list(session: SessionDep):
+#     return session.query(AiModelDetail).all()
 
 
 @router.post("/mcp_start", operation_id="mcp_start")
 async def mcp_start(session: SessionDep, chat: ChatStart):
-    user = authenticate(session=session, account=chat.username, password=chat.password)
+    user: BaseUserDTO = authenticate(session=session, account=chat.username, password=chat.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()
     t = Token(access_token=create_access_token(
@@ -61,9 +73,26 @@ async def mcp_start(session: SessionDep, chat: ChatStart):
 
 @router.post("/mcp_question", operation_id="mcp_question")
 async def mcp_question(session: SessionDep, chat: ChatMcp):
-    user = await get_current_user(session, chat.token)
-
-    llm_service = LLMService(session, user, chat)
+    try:
+        payload = jwt.decode(
+            chat.token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
+        )
+        token_data = TokenPayload(**payload)
+    except (InvalidTokenError, ValidationError):
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Could not validate credentials",
+        )
+    session_user = await get_user_info(session=session, user_id=token_data.id)
+    session_user = UserInfoDTO.model_validate(session_user)
+    if not session_user:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    if session_user.status != 1:
+        raise HTTPException(status_code=400, detail="Inactive user")
+
+    # ask
+    llm_service = LLMService(session, session_user, chat)
     llm_service.init_record()
 
     return StreamingResponse(run_task(llm_service, False), media_type="text/event-stream")