From e899ab3b1d4931e6799185088727dae4a4603307 Mon Sep 17 00:00:00 2001 From: junjun Date: Thu, 17 Jul 2025 15:49:22 +0800 Subject: [PATCH] feat: ds add oid --- backend/alembic/versions/028_ds_oid.py | 28 +++++++++ backend/apps/datasource/api/datasource.py | 8 +-- backend/apps/datasource/crud/datasource.py | 22 +++---- backend/apps/datasource/models/datasource.py | 2 + backend/apps/mcp/mcp.py | 63 ++++++++++++++------ 5 files changed, 92 insertions(+), 31 deletions(-) create mode 100644 backend/alembic/versions/028_ds_oid.py diff --git a/backend/alembic/versions/028_ds_oid.py b/backend/alembic/versions/028_ds_oid.py new file mode 100644 index 00000000..83afde51 --- /dev/null +++ b/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 ### diff --git a/backend/apps/datasource/api/datasource.py b/backend/apps/datasource/api/datasource.py index 2c3d36e8..43a01862 100644 --- a/backend/apps/datasource/api/datasource.py +++ b/backend/apps/datasource/api/datasource.py @@ -19,8 +19,8 @@ path = "/opt/sqlbot/data/excel" @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) diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 5647f7c8..3f6a32d3 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -22,10 +22,10 @@ from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, Cor 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,14 +45,15 @@ 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' @@ -60,10 +61,11 @@ def check_name(session: SessionDep, ds: CoreDatasource): 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() diff --git a/backend/apps/datasource/models/datasource.py b/backend/apps/datasource/models/datasource.py index ab73eaa3..4377344a 100644 --- a/backend/apps/datasource/models/datasource.py +++ b/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] = [] diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index 3530f493..42f945e5 100644 --- a/backend/apps/mcp/mcp.py +++ b/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 = APIRouter(tags=["mcp"], prefix="/mcp") # )) -@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) + 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") - llm_service = LLMService(session, user, chat) + 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")