From 0afccc9267b606bf88c051109bdfe9a322c5b6a4 Mon Sep 17 00:00:00 2001
From: junjun <junjie.xia@fit2cloud.com>
Date: Tue, 10 Jun 2025 14:36:26 +0800
Subject: [PATCH] feat: mcp

---
 backend/apps/datasource/api/datasource.py |  2 +-
 backend/common/utils/whitelist.py         |  4 +++-
 backend/main.py                           | 24 +++++++++++++++++++++--
 backend/pyproject.toml                    |  1 +
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/backend/apps/datasource/api/datasource.py b/backend/apps/datasource/api/datasource.py
index b0921673..08f2e6f7 100644
--- a/backend/apps/datasource/api/datasource.py
+++ b/backend/apps/datasource/api/datasource.py
@@ -18,7 +18,7 @@ router = APIRouter(tags=["datasource"], prefix="/datasource")
 path = "/opt/sqlbot/data/excel"
 
 
-@router.get("/list")
+@router.get("/list", operation_id="get_datasource_list")
 async def datasource_list(session: SessionDep):
     return get_datasource_list(session=session)
 
diff --git a/backend/common/utils/whitelist.py b/backend/common/utils/whitelist.py
index c02dfe44..78010184 100644
--- a/backend/common/utils/whitelist.py
+++ b/backend/common/utils/whitelist.py
@@ -20,7 +20,9 @@ wlist = [
     "*.woff2",
     "*.ttf",
     "*.eot",
-    "*.otf"
+    "*.otf",
+    "/mcp",
+    "/mcp/message"
 ]
 
 class WhitelistChecker:
diff --git a/backend/main.py b/backend/main.py
index 2cb6b8d3..bef058cb 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -14,15 +14,20 @@ from common.core.config import settings
 from common.core.response_middleware import ResponseMiddleware, exception_handler
 from alembic.config import Config
 from alembic import command
+from fastapi_mcp import FastApiMCP
+
 
 def run_migrations():
     alembic_cfg = Config("alembic.ini")
     command.upgrade(alembic_cfg, "head")
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     run_migrations()
     yield
 
+
 def custom_generate_unique_id(route: APIRoute) -> str:
     tag = route.tags[0] if route.tags and len(route.tags) > 0 else ""
     return f"{tag}-{route.name}"
@@ -38,6 +43,17 @@ app = FastAPI(
     lifespan=lifespan
 )
 
+mcp_app = FastAPI()
+
+mcp = FastApiMCP(
+    app,
+    name="SQLBot MCP Server",
+    description="SQLBot MCP Server",
+    include_operations=["get_datasource_list"]
+)
+
+mcp.mount(mcp_app)
+
 # Set all CORS enabled origins
 if settings.all_cors_origins:
     app.add_middleware(
@@ -66,9 +82,13 @@ else:
     @app.get("/", include_in_schema=False)
     async def read_index():
         return FileResponse(path=os.path.join(frontend_dist, "index.html"))
-    
+
+
     app.mount("/", StaticFiles(directory=frontend_dist), name="static")
 
+mcp.setup_server()
+
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
\ No newline at end of file
+
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
diff --git a/backend/pyproject.toml b/backend/pyproject.toml
index 535c219b..41aaea6a 100644
--- a/backend/pyproject.toml
+++ b/backend/pyproject.toml
@@ -34,6 +34,7 @@ dependencies = [
     "psycopg2-binary (>=2.9.10,<3.0.0)",
     "oracledb (>=3.1.1,<4.0.0)",
     "pyyaml (>=6.0.2,<7.0.0)",
+    "fastapi-mcp (>=0.3.4,<0.4.0)",
 ]
 [[tool.uv.index]]
 url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"