zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-29 04:08:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: add context, part7 (#17402)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-27 11:23:06 +08:00
committed by GitHub
parent 9d4847beaf
commit 1571abd98a
4 changed files with 71 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
internal/dao/mcp.go
View File

@@ -17,6 +17,7 @@
package dao
import (
"context"
"errors"
"fmt"
"strings"
@@ -45,9 +46,9 @@ func NewMCPServerDAO() *MCPServerDAO {
}
// GetByID returns an MCP server by ID.
func (dao *MCPServerDAO) GetByID(id string) (*entity.MCPServer, error) {
func (dao *MCPServerDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.MCPServer, error) {
var server entity.MCPServer
if err := DB.Where("id = ?", id).First(&server).Error; err != nil {
if err := db.WithContext(ctx).Where("id = ?", id).First(&server).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, nil
}
@@ -57,9 +58,9 @@ func (dao *MCPServerDAO) GetByID(id string) (*entity.MCPServer, error) {
}
// ExistsByNameAndTenant returns whether an MCP server name already exists for a tenant.
func (dao *MCPServerDAO) ExistsByNameAndTenant(name, tenantID string) (bool, error) {
func (dao *MCPServerDAO) ExistsByNameAndTenant(ctx context.Context, db *gorm.DB, name, tenantID string) (bool, error) {
var count int64
if err := DB.Model(&entity.MCPServer{}).
if err := db.WithContext(ctx).Model(&entity.MCPServer{}).
Where("name = ? AND tenant_id = ?", name, tenantID).
Count(&count).Error; err != nil {
return false, err
@@ -68,16 +69,16 @@ func (dao *MCPServerDAO) ExistsByNameAndTenant(name, tenantID string) (bool, err
}
// CreateMCPServer creates an MCP server.
func (dao *MCPServerDAO) CreateMCPServer(server *entity.MCPServer) error {
return DB.Create(server).Error
func (dao *MCPServerDAO) CreateMCPServer(ctx context.Context, db *gorm.DB, server *entity.MCPServer) error {
return db.WithContext(ctx).Create(server).Error
}
// ListMCPServers returns MCP servers for a tenant with optional filtering.
func (dao *MCPServerDAO) ListMCPServers(tenantID string, ids []string, keywords string, orderby string, desc bool) ([]*entity.MCPServer, int64, error) {
func (dao *MCPServerDAO) ListMCPServers(ctx context.Context, db *gorm.DB, tenantID string, ids []string, keywords string, orderby string, desc bool) ([]*entity.MCPServer, int64, error) {
var servers []*entity.MCPServer
var total int64
query := DB.Model(&entity.MCPServer{}).Where("tenant_id = ?", tenantID)
query := db.WithContext(ctx).Model(&entity.MCPServer{}).Where("tenant_id = ?", tenantID)
if len(ids) > 0 {
query = query.Where("id IN ?", ids)
@@ -111,17 +112,17 @@ func (dao *MCPServerDAO) ListMCPServers(tenantID string, ids []string, keywords
}
// GetByIDAndTenant returns an MCP server owned by a tenant.
func (dao *MCPServerDAO) GetByIDAndTenant(id, tenantID string) (*entity.MCPServer, error) {
func (dao *MCPServerDAO) GetByIDAndTenant(ctx context.Context, db *gorm.DB, id, tenantID string) (*entity.MCPServer, error) {
var server entity.MCPServer
if err := DB.Where("id = ? AND tenant_id = ?", id, tenantID).First(&server).Error; err != nil {
if err := db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).First(&server).Error; err != nil {
return nil, err
}
return &server, nil
}
// DeleteMCPServer deletes an MCP server owned by a tenant.
func (dao *MCPServerDAO) DeleteMCPServer(id, tenantID string) (bool, error) {
result := DB.Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&entity.MCPServer{})
func (dao *MCPServerDAO) DeleteMCPServer(ctx context.Context, db *gorm.DB, id, tenantID string) (bool, error) {
result := db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&entity.MCPServer{})
if result.Error != nil {
return false, result.Error
}
@@ -129,8 +130,8 @@ func (dao *MCPServerDAO) DeleteMCPServer(id, tenantID string) (bool, error) {
}
// UpdateMCPServer updates an MCP server owned by a tenant.
func (dao *MCPServerDAO) UpdateMCPServer(id, tenantID string, updates map[string]interface{}) (bool, error) {
result := DB.Model(&entity.MCPServer{}).
func (dao *MCPServerDAO) UpdateMCPServer(ctx context.Context, db *gorm.DB, id, tenantID string, updates map[string]interface{}) (bool, error) {
result := db.WithContext(ctx).Model(&entity.MCPServer{}).
Where("id = ? AND tenant_id = ?", id, tenantID).
Updates(updates)
if result.Error != nil {

25
internal/handler/mcp.go
View File

@@ -80,8 +80,9 @@ func (h *MCPHandler) CreateMCPServer(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error())
return
}
ctx := c.Request.Context()
result, code, err := h.mcpService.CreateMCPServer(user.ID, req)
result, code, err := h.mcpService.CreateMCPServer(ctx, user.ID, req)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -108,13 +109,13 @@ func (h *MCPHandler) ListMCPServers(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error())
return
}
ctx := c.Request.Context()
orderby := c.DefaultQuery("orderby", "create_time")
desc := strings.ToLower(c.DefaultQuery("desc", "true")) != "false"
keywords := c.Query("keywords")
mcpIDs := getMCPIDsFromQuery(c)
result, code, err := h.mcpService.ListMCPServers(user.ID, mcpIDs, keywords, page, pageSize, orderby, desc)
result, code, err := h.mcpService.ListMCPServers(ctx, user.ID, mcpIDs, keywords, page, pageSize, orderby, desc)
if err != nil {
if code == common.CodeServerError {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, code, nil, err.Error())
@@ -133,10 +134,10 @@ func (h *MCPHandler) GetMCPServer(c *gin.Context) {
common.ErrorWithCode(c, errorCode, errorMessage)
return
}
ctx := c.Request.Context()
mcpID := c.Param("mcp_id")
if c.Query("mode") == "download" {
result, code, err := h.mcpService.ExportMCPServer(user.ID, mcpID)
result, code, err := h.mcpService.ExportMCPServer(ctx, user.ID, mcpID)
if err != nil {
mcpDetailError(c, code, err)
return
@@ -145,7 +146,7 @@ func (h *MCPHandler) GetMCPServer(c *gin.Context) {
return
}
result, code, err := h.mcpService.GetMCPServer(user.ID, mcpID)
result, code, err := h.mcpService.GetMCPServer(ctx, user.ID, mcpID)
if err != nil {
mcpDetailError(c, code, err)
return
@@ -175,8 +176,8 @@ func (h *MCPHandler) UpdateMCPServer(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error())
return
}
result, code, err := h.mcpService.UpdateMCPServer(user.ID, mcpID, req)
ctx := c.Request.Context()
result, code, err := h.mcpService.UpdateMCPServer(ctx, user.ID, mcpID, req)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -192,8 +193,8 @@ func (h *MCPHandler) DeleteMCPServer(c *gin.Context) {
common.ErrorWithCode(c, errorCode, errorMessage)
return
}
result, code, err := h.mcpService.DeleteMCPServer(user.ID, c.Param("mcp_id"))
ctx := c.Request.Context()
result, code, err := h.mcpService.DeleteMCPServer(ctx, user.ID, c.Param("mcp_id"))
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -320,8 +321,8 @@ func (h *MCPHandler) ImportMCPServers(c *gin.Context) {
// 10 s fallback when timeout <= 0.
_ = json.Unmarshal(rawTimeout, &timeout)
}
results, err := h.mcpService.ImportServers(user.ID, servers, timeout)
ctx := c.Request.Context()
results, err := h.mcpService.ImportServers(ctx, user.ID, servers, timeout)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeBadRequest, nil, err.Error())
return

82
internal/service/mcp.go
View File

@@ -115,39 +115,39 @@ type ListMCPServersResponse struct {
const maxMCPFetchTimeoutSec = 60
// CreateMCPServer creates an MCP server owned by a tenant.
func (s *MCPService) CreateMCPServer(tenantID string, req CreateMCPServerRequest) (*CreateMCPServerResponse, common.ErrorCode, error) {
func (s *MCPService) CreateMCPServer(ctx context.Context, tenantID string, req CreateMCPServerRequest) (*CreateMCPServerResponse, common.ErrorCode, error) {
if req.Timeout < 0 || req.Timeout > maxMCPFetchTimeoutSec {
return nil, common.CodeDataError, errors.New("Invalid timeout.")
return nil, common.CodeDataError, errors.New("invalid timeout")
}
if !isValidMCPServerType(req.ServerType) {
return nil, common.CodeDataError, errors.New("Unsupported MCP server type.")
return nil, common.CodeDataError, errors.New("unsupported MCP server type")
}
if req.Name == "" || len([]byte(req.Name)) > mcpServerNameLimit {
return nil, common.CodeDataError, fmt.Errorf("Invalid MCP name or length is %d which is large than 255.", len([]byte(req.Name)))
}
exists, err := s.mcpServerDAO.ExistsByNameAndTenant(req.Name, tenantID)
exists, err := s.mcpServerDAO.ExistsByNameAndTenant(ctx, dao.DB, req.Name, tenantID)
if err != nil {
return nil, common.CodeServerError, err
}
if exists {
return nil, common.CodeDataError, errors.New("Duplicated MCP server name.")
return nil, common.CodeDataError, errors.New("duplicated MCP server name")
}
if req.URL == "" {
return nil, common.CodeDataError, errors.New("Invalid url.")
return nil, common.CodeDataError, errors.New("invalid url")
}
if _, err := s.tenantDAO.GetByID(tenantID); err != nil {
return nil, common.CodeDataError, errors.New("Tenant not found.")
if _, err = s.tenantDAO.GetByID(tenantID); err != nil {
return nil, common.CodeDataError, errors.New("tenant not found")
}
headers := safeJSONMap(req.Headers)
variables := safeJSONMap(req.Variables)
delete(variables, "tools")
tools, err := fetchMCPTools(context.Background(), req.URL, req.ServerType, headers, variables, req.Timeout)
tools, err := fetchMCPTools(ctx, req.URL, req.ServerType, headers, variables, req.Timeout)
if err != nil {
return nil, common.CodeDataError, err
}
@@ -164,8 +164,8 @@ func (s *MCPService) CreateMCPServer(tenantID string, req CreateMCPServerRequest
Headers: headers,
}
if err := s.mcpServerDAO.CreateMCPServer(server); err != nil {
return nil, common.CodeDataError, errors.New("Failed to create MCP server.")
if err = s.mcpServerDAO.CreateMCPServer(ctx, dao.DB, server); err != nil {
return nil, common.CodeDataError, errors.New("failed to create MCP server")
}
return &CreateMCPServerResponse{
@@ -231,8 +231,8 @@ func jsonMapStringValues(values entity.JSONMap) map[string]string {
return out
}
func (s *MCPService) GetMCPServer(tenantID, mcpID string) (*entity.MCPServer, common.ErrorCode, error) {
server, err := s.mcpServerDAO.GetByIDAndTenant(mcpID, tenantID)
func (s *MCPService) GetMCPServer(ctx context.Context, tenantID, mcpID string) (*entity.MCPServer, common.ErrorCode, error) {
server, err := s.mcpServerDAO.GetByIDAndTenant(ctx, dao.DB, mcpID, tenantID)
if err != nil {
if isMCPServerNotFound(err) {
return nil, common.CodeDataError, mcpServerNotFoundError(mcpID, tenantID)
@@ -245,8 +245,8 @@ func (s *MCPService) GetMCPServer(tenantID, mcpID string) (*entity.MCPServer, co
return server, common.CodeSuccess, nil
}
func (s *MCPService) ExportMCPServer(userID, mcpID string) (*ExportMCPServerResponse, common.ErrorCode, error) {
server, code, err := s.GetMCPServer(userID, mcpID)
func (s *MCPService) ExportMCPServer(ctx context.Context, userID, mcpID string) (*ExportMCPServerResponse, common.ErrorCode, error) {
server, code, err := s.GetMCPServer(ctx, userID, mcpID)
if err != nil {
return nil, code, err
}
@@ -281,8 +281,8 @@ func newExportMCPServerResponse(server *entity.MCPServer) *ExportMCPServerRespon
}
// UpdateMCPServer updates an MCP server owned by a tenant.
func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServerRequest) (*entity.MCPServer, common.ErrorCode, error) {
server, err := s.mcpServerDAO.GetByIDAndTenant(mcpID, tenantID)
func (s *MCPService) UpdateMCPServer(ctx context.Context, tenantID, mcpID string, req UpdateMCPServerRequest) (*entity.MCPServer, common.ErrorCode, error) {
server, err := s.mcpServerDAO.GetByIDAndTenant(ctx, dao.DB, mcpID, tenantID)
if err != nil {
if isMCPServerNotFound(err) {
return nil, common.CodeDataError, mcpServerNotFoundError(mcpID, tenantID)
@@ -302,7 +302,7 @@ func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServer
serverTypeProvided = true
}
if serverTypeProvided && !isValidMCPServerType(serverType) {
return nil, common.CodeDataError, errors.New("Unsupported MCP server type.")
return nil, common.CodeDataError, errors.New("unsupported MCP server type")
}
serverName := server.Name
@@ -324,12 +324,12 @@ func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServer
} else if ok {
serverURL = strings.TrimSpace(value)
if serverURL == "" {
return nil, common.CodeDataError, errors.New("Invalid url.")
return nil, common.CodeDataError, errors.New("invalid url")
}
serverURLProvided = true
}
if serverURL == "" {
return nil, common.CodeDataError, errors.New("Invalid url.")
return nil, common.CodeDataError, errors.New("invalid url")
}
headers := server.Headers
@@ -355,7 +355,7 @@ func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServer
if err != nil {
return nil, common.CodeDataError, err
}
tools, err := fetchMCPTools(context.Background(), serverURL, serverType, headers, variables, timeoutSeconds)
tools, err := fetchMCPTools(ctx, serverURL, serverType, headers, variables, timeoutSeconds)
if err != nil {
return nil, common.CodeDataError, err
}
@@ -387,11 +387,11 @@ func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServer
updates["description"] = description
}
if _, err := s.mcpServerDAO.UpdateMCPServer(mcpID, tenantID, updates); err != nil {
if _, err = s.mcpServerDAO.UpdateMCPServer(ctx, dao.DB, mcpID, tenantID, updates); err != nil {
return nil, common.CodeServerError, err
}
updatedServer, err := s.mcpServerDAO.GetByIDAndTenant(mcpID, tenantID)
updatedServer, err := s.mcpServerDAO.GetByIDAndTenant(ctx, dao.DB, mcpID, tenantID)
if err != nil {
if isMCPServerNotFound(err) {
return nil, common.CodeDataError, mcpServerNotFoundError(mcpID, tenantID)
@@ -419,8 +419,8 @@ func isMCPServerNotFound(err error) bool {
}
// ListMCPServers lists MCP servers owned by a tenant.
func (s *MCPService) ListMCPServers(tenantID string, ids []string, keywords string, page, pageSize int, orderby string, desc bool) (*ListMCPServersResponse, common.ErrorCode, error) {
servers, total, err := s.mcpServerDAO.ListMCPServers(tenantID, ids, keywords, orderby, desc)
func (s *MCPService) ListMCPServers(ctx context.Context, tenantID string, ids []string, keywords string, page, pageSize int, orderby string, desc bool) (*ListMCPServersResponse, common.ErrorCode, error) {
servers, total, err := s.mcpServerDAO.ListMCPServers(ctx, dao.DB, tenantID, ids, keywords, orderby, desc)
if err != nil {
var orderbyErr *dao.InvalidMCPServerOrderByError
if errors.As(err, &orderbyErr) {
@@ -458,8 +458,8 @@ func (s *MCPService) ListMCPServers(tenantID string, ids []string, keywords stri
}
// DeleteMCPServer deletes an MCP server owned by a tenant.
func (s *MCPService) DeleteMCPServer(tenantID, mcpID string) (bool, common.ErrorCode, error) {
server, err := s.mcpServerDAO.GetByID(mcpID)
func (s *MCPService) DeleteMCPServer(ctx context.Context, tenantID, mcpID string) (bool, common.ErrorCode, error) {
server, err := s.mcpServerDAO.GetByID(ctx, dao.DB, mcpID)
if err != nil {
return false, common.CodeServerError, fmt.Errorf("failed to get MCP server %s: %w", mcpID, err)
}
@@ -467,7 +467,7 @@ func (s *MCPService) DeleteMCPServer(tenantID, mcpID string) (bool, common.Error
return false, common.CodeDataError, mcpServerNotFoundError(mcpID, tenantID)
}
deleted, err := s.mcpServerDAO.DeleteMCPServer(mcpID, tenantID)
deleted, err := s.mcpServerDAO.DeleteMCPServer(ctx, dao.DB, mcpID, tenantID)
if err != nil {
return false, common.CodeServerError, err
}
@@ -479,7 +479,7 @@ func (s *MCPService) DeleteMCPServer(tenantID, mcpID string) (bool, common.Error
}
func mcpServerNotFoundError(mcpID, tenantID string) error {
return fmt.Errorf("Cannot find MCP server %s for user %s", mcpID, tenantID)
return fmt.Errorf("cannot find MCP server %s for user %s", mcpID, tenantID)
}
func isValidMCPServerType(serverType string) bool {
@@ -519,7 +519,7 @@ func safeJSONMap(raw json.RawMessage) entity.JSONMap {
var value map[string]interface{}
if err := json.Unmarshal(raw, &value); err == nil && value != nil {
return entity.JSONMap(value)
return value
}
var textValue string
@@ -530,7 +530,7 @@ func safeJSONMap(raw json.RawMessage) entity.JSONMap {
if err := json.Unmarshal([]byte(textValue), &value); err != nil || value == nil {
return entity.JSONMap{}
}
return entity.JSONMap(value)
return value
}
// ---------- import + test (this PR's additions) ----------
@@ -556,7 +556,7 @@ type ImportResult struct {
}
// ImportServers bulk-imports MCP servers from a {"mcpServers": {name: config}} map.
func (s *MCPService) ImportServers(tenantID string, servers map[string]map[string]interface{}, timeoutSeconds float64) ([]ImportResult, error) {
func (s *MCPService) ImportServers(ctx context.Context, tenantID string, servers map[string]map[string]interface{}, timeoutSeconds float64) ([]ImportResult, error) {
if timeoutSeconds <= 0 {
timeoutSeconds = defaultMCPFetchTimeoutSec
}
@@ -580,7 +580,7 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin
}
baseName := serverName
newName, err := s.nextAvailableMCPName(baseName, tenantID)
newName, err := s.nextAvailableMCPName(ctx, baseName, tenantID)
if err != nil {
return nil, err
}
@@ -627,8 +627,8 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin
}
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
tools, fetchErr := utility.FetchTools(ctx, utility.FetchOptions{
mcpCtx, cancel := context.WithTimeout(context.Background(), timeout)
tools, fetchErr := utility.FetchTools(mcpCtx, utility.FetchOptions{
URL: url,
ServerType: stype,
Headers: headers,
@@ -651,7 +651,7 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin
Variables: entity.JSONMap(variables),
Headers: entity.JSONMap(headerVals),
}
if err := s.mcpServerDAO.CreateMCPServer(server); err != nil {
if err = s.mcpServerDAO.CreateMCPServer(ctx, dao.DB, server); err != nil {
results = append(results, ImportResult{Server: serverName, Success: false, Message: "Failed to create MCP server."})
continue
}
@@ -665,11 +665,11 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin
return results, nil
}
func (s *MCPService) nextAvailableMCPName(base, tenantID string) (string, error) {
func (s *MCPService) nextAvailableMCPName(ctx context.Context, base, tenantID string) (string, error) {
name := base
counter := 0
for {
exists, err := s.mcpServerDAO.ExistsByNameAndTenant(name, tenantID)
exists, err := s.mcpServerDAO.ExistsByNameAndTenant(ctx, dao.DB, name, tenantID)
if err != nil {
return "", err
}
@@ -696,7 +696,7 @@ type TestServerRequest struct {
// TestServer opens a live MCP session and returns the tools the server advertises.
func (s *MCPService) TestServer(mcpID string, req *TestServerRequest) ([]map[string]interface{}, error) {
if req == nil || req.URL == "" {
return nil, fmt.Errorf("%w: Invalid MCP url.", ErrMCPInvalidURL)
return nil, fmt.Errorf("%w: Invalid MCP url", ErrMCPInvalidURL)
}
if !isValidMCPServerType(req.ServerType) {
return nil, ErrMCPInvalidType
@@ -730,9 +730,9 @@ func (s *MCPService) TestServer(mcpID string, req *TestServerRequest) ([]map[str
}
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
mcpCtx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
tools, err := utility.FetchTools(ctx, utility.FetchOptions{
tools, err := utility.FetchTools(mcpCtx, utility.FetchOptions{
URL: req.URL,
ServerType: req.ServerType,
Headers: headers,

3
internal/service/mcp_test.go
View File

@@ -67,7 +67,8 @@ func TestImportServersValidationErrors(t *testing.T) {
"missing-fields": {"foo": "bar"},
"bad-type": {"url": "http://example.com", "type": "stdio"},
}
results, err := s.ImportServers("tenant-1", configs, 1)
ctx := t.Context()
results, err := s.ImportServers(ctx, "tenant-1", configs, 1)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}