Go: add context to DAO (#17269)

### Summary DAO layer doesn't use context, this PR is to fix it. --------- Signed-off-by: Jin Hai <haijin.chn@gmail.com>
2026-07-24 17:36:47 +08:00 · 2026-07-23 12:15:58 +08:00
parent 7b29fc10ca
commit 5478f955ab
45 changed files with 723 additions and 578 deletions
--- a/internal/handler/agent.go
+++ b/internal/handler/agent.go
@@ -359,8 +359,8 @@ func (h *AgentHandler) ListTemplates(c *gin.Context) {
 		common.ErrorWithCode(c, errorCode, errorMessage)
 		return
 	}
-
-	templates, err := h.agentService.ListTemplates()
+	ctx := c.Request.Context()
+	templates, err := h.agentService.ListTemplates(ctx)
 	if err != nil {
 		jsonInternalError(c, err)
 		return
@@ -629,7 +629,8 @@ func (h *AgentHandler) ListAgentTemplates(c *gin.Context) {
 		common.ResponseWithCodeData(c, code, nil, msg)
 		return
 	}
-	rows, err := h.agentService.ListTemplates()
+	ctx := c.Request.Context()
+	rows, err := h.agentService.ListTemplates(ctx)
 	if err != nil {
 		common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error())
 		return
@@ -711,8 +712,8 @@ func (h *AgentHandler) ListAgentSessions(c *gin.Context) {
 	sessionID := c.Query("id")
 	expUserID := c.Query("user_id")
 	includeDSL := c.Query("dsl") == "true"
-
-	resp, code, err := h.agentService.ListAgentSessions(user.ID, user.ID, canvasID, service.ListAgentSessionsRequest{
+	ctx := c.Request.Context()
+	resp, code, err := h.agentService.ListAgentSessions(ctx, user.ID, user.ID, canvasID, service.ListAgentSessionsRequest{
 		Page:       page,
 		PageSize:   pageSize,
 		Keywords:   keywords,
@@ -749,7 +750,8 @@ func (h *AgentHandler) CreateAgentSession(c *gin.Context) {
 		common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "Invalid request: "+err.Error())
 		return
 	}
-	row, code, err := h.agentService.CreateAgentSession(&service.CreateAgentSessionRequest{
+	ctx := c.Request.Context()
+	row, code, err := h.agentService.CreateAgentSession(ctx, &service.CreateAgentSessionRequest{
 		UserID:  user.ID,
 		AgentID: canvasID,
 		Name:    body.Name,
@@ -772,7 +774,8 @@ func (h *AgentHandler) GetAgentSession(c *gin.Context) {
 	}
 	canvasID := c.Param("canvas_id")
 	sessionID := c.Param("session_id")
-	row, code, err := h.agentService.GetAgentSession(user.ID, canvasID, sessionID)
+	ctx := c.Request.Context()
+	row, code, err := h.agentService.GetAgentSession(ctx, user.ID, canvasID, sessionID)
 	if err != nil {
 		common.ErrorWithCode(c, code, err.Error())
 		return
@@ -791,10 +794,11 @@ func (h *AgentHandler) DeleteAgentSession(c *gin.Context) {
 		common.ResponseWithCodeData(c, code, nil, msg)
 		return
 	}
+	ctx := c.Request.Context()
 	canvasID := c.Param("canvas_id")
 	sessionID := c.Param("session_id")
 	if sessionID != "" {
-		ok, code, err := h.agentService.DeleteAgentSessionItem(user.ID, canvasID, sessionID)
+		ok, code, err := h.agentService.DeleteAgentSessionItem(ctx, user.ID, canvasID, sessionID)
 		if err != nil {
 			common.ErrorWithCode(c, code, err.Error())
 			return
@@ -812,7 +816,7 @@ func (h *AgentHandler) DeleteAgentSession(c *gin.Context) {
 			}
 		}
 	}
-	result, code, err := h.agentService.DeleteAgentSessions(user.ID, canvasID, ids, deleteAll)
+	result, code, err := h.agentService.DeleteAgentSessions(ctx, user.ID, canvasID, ids, deleteAll)
 	if err != nil {
 		common.ErrorWithCode(c, code, err.Error())
 		return
--- a/internal/handler/api_token.go
+++ b/internal/handler/api_token.go
@@ -50,9 +50,10 @@ func (h *SystemHandler) ListAPIKeys(c *gin.Context) {
 	}

 	tenantID := tenants[0].TenantID
+	ctx := c.Request.Context()

 	// Get keys for the tenant
-	keys, err := h.systemService.ListAPIKeys(tenantID)
+	keys, err := h.systemService.ListAPIKeys(ctx, tenantID)
 	if err != nil {
 		common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, "Failed to list keys")
 		return
@@ -92,8 +93,10 @@ func (h *SystemHandler) CreateKey(c *gin.Context) {
 		return
 	}

+	ctx := c.Request.Context()
+
 	// Create key
-	key, err := h.systemService.CreateAPIKey(tenantID, &req)
+	key, err := h.systemService.CreateAPIKey(ctx, tenantID, &req)
 	if err != nil {
 		common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, "Failed to create key")
 		return
@@ -133,8 +136,10 @@ func (h *SystemHandler) DeleteKey(c *gin.Context) {
 		return
 	}

+	ctx := c.Request.Context()
+
 	// Delete key
-	if err = h.systemService.DeleteAPIKey(tenantID, key); err != nil {
+	if err = h.systemService.DeleteAPIKey(ctx, tenantID, key); err != nil {
 		common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, "Failed to delete key")
 		return
 	}
--- a/internal/handler/chat.go
+++ b/internal/handler/chat.go
@@ -103,9 +103,10 @@ func (h *ChatHandler) ListChats(c *gin.Context) {
 	}

 	ownerIDs := getOwnerIDs(c)
+	ctx := c.Request.Context()

 	// List chats - default to valid status "1" (same as Python StatusEnum.VALID.value)
-	result, err := h.chatService.ListChats(userID, "1", keywords, page, pageSize, orderby, desc, ownerIDs)
+	result, err := h.chatService.ListChats(ctx, userID, "1", keywords, page, pageSize, orderby, desc, ownerIDs)
 	if err != nil {
 		common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error())
 		return
@@ -141,7 +142,8 @@ func (h *ChatHandler) Create(c *gin.Context) {
 		req = map[string]interface{}{}
 	}

-	result, code, err := h.chatService.Create(user.ID, req)
+	ctx := c.Request.Context()
+	result, code, err := h.chatService.Create(ctx, user.ID, req)
 	if err != nil {
 		common.ErrorWithCode(c, code, err.Error())
 		return
@@ -232,7 +234,8 @@ func (h *ChatHandler) DeleteChat(c *gin.Context) {
 		return
 	}

-	if err := h.chatService.DeleteChat(userID, chatID); err != nil {
+	ctx := c.Request.Context()
+	if err := h.chatService.DeleteChat(ctx, userID, chatID); err != nil {
 		if err.Error() == "no authorization" {
 			common.ResponseWithCodeData(c, common.CodeAuthenticationError, false, "No authorization.")
 			return
@@ -263,9 +266,10 @@ func (h *ChatHandler) BulkDeleteChats(c *gin.Context) {
 		return
 	}

+	ctx := c.Request.Context()
 	if len(req.IDs) == 0 && !req.DeleteAll {
 		if req.ChatID != "" {
-			if err := h.chatService.DeleteChat(userID, req.ChatID); err != nil {
+			if err := h.chatService.DeleteChat(ctx, userID, req.ChatID); err != nil {
 				if err.Error() == "no authorization" {
 					common.ResponseWithCodeData(c, common.CodeAuthenticationError, false, "No authorization.")
 					return
@@ -282,7 +286,7 @@ func (h *ChatHandler) BulkDeleteChats(c *gin.Context) {
 		return
 	}

-	result, err := h.chatService.BulkDeleteChats(userID, &req)
+	result, err := h.chatService.BulkDeleteChats(ctx, userID, &req)
 	if err != nil {
 		common.ErrorWithCode(c, common.CodeDataError, err.Error())
 		return
@@ -327,7 +331,8 @@ func (h *ChatHandler) GetChat(c *gin.Context) {
 	}

 	// Get chat detail with permission check
-	chat, err := h.chatService.GetChat(userID, chatID)
+	ctx := c.Request.Context()
+	chat, err := h.chatService.GetChat(ctx, userID, chatID)
 	if err != nil {
 		errMsg := err.Error()
 		// Check if it's an authorization error
@@ -404,14 +409,13 @@ func (h *ChatHandler) updateChatByMethod(c *gin.Context, patch bool) {
 		return
 	}

-	var (
-		result map[string]interface{}
-		err    error
-	)
+	var result map[string]interface{}
+	var err error
+	ctx := c.Request.Context()
 	if patch {
-		result, err = h.chatService.PatchChat(user.ID, chatID, req)
+		result, err = h.chatService.PatchChat(ctx, user.ID, chatID, req)
 	} else {
-		result, err = h.chatService.UpdateChat(user.ID, chatID, req)
+		result, err = h.chatService.UpdateChat(ctx, user.ID, chatID, req)
 	}
 	if err != nil {
 		if err.Error() == "no authorization" {
--- a/internal/handler/chat_channel.go
+++ b/internal/handler/chat_channel.go
@@ -15,6 +15,7 @@
 package handler

 import (
+	"context"
 	"strings"

 	"github.com/gin-gonic/gin"
@@ -25,11 +26,11 @@ import (
 )

 type ChatChannelService interface {
-	CreateChatChannel(tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error)
-	List(tenantID string) ([]*entity.ChatChannelListResponse, error)
-	GetChatChannel(userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error)
-	UpdateChatChannel(userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error)
-	DeleteChatChannel(userID, channelID string) (bool, common.ErrorCode, error)
+	CreateChatChannel(ctx context.Context, userID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error)
+	List(ctx context.Context, tenantID string) ([]*entity.ChatChannelListResponse, error)
+	GetChatChannel(ctx context.Context, userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error)
+	UpdateChatChannel(ctx context.Context, userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error)
+	DeleteChatChannel(ctx context.Context, userID, channelID string) (bool, common.ErrorCode, error)
 }

 type ChatChannelHandler struct {
@@ -66,7 +67,10 @@ func (h *ChatChannelHandler) CreateChatChannel(c *gin.Context) {
 		return
 	}

+	ctx := c.Request.Context()
+
 	row, err := h.chatChannelService.CreateChatChannel(
+		ctx,
 		user.ID,
 		req.Name,
 		req.Channel,
@@ -88,7 +92,8 @@ func (h *ChatChannelHandler) ListChatChannel(c *gin.Context) {
 		return
 	}

-	rows, err := h.chatChannelService.List(user.ID)
+	ctx := c.Request.Context()
+	rows, err := h.chatChannelService.List(ctx, user.ID)
 	if err != nil {
 		common.ResponseWithCodeData(c, common.CodeServerError, nil, err.Error())
 		return
@@ -116,7 +121,8 @@ func (h *ChatChannelHandler) GetChatChannel(c *gin.Context) {
 		return
 	}

-	channel, code, err := h.chatChannelService.GetChatChannel(userID, channelID)
+	ctx := c.Request.Context()
+	channel, code, err := h.chatChannelService.GetChatChannel(ctx, userID, channelID)
 	if code != common.CodeSuccess || err != nil {
 		writeChatChannelError(c, code, chatChannelErrMsg(code, err))
 		return
@@ -151,7 +157,8 @@ func (h *ChatChannelHandler) UpdateChatChannel(c *gin.Context) {
 		return
 	}

-	result, code, err := h.chatChannelService.UpdateChatChannel(userID, channelID, unwrapChatChannelPayload(request))
+	ctx := c.Request.Context()
+	result, code, err := h.chatChannelService.UpdateChatChannel(ctx, userID, channelID, unwrapChatChannelPayload(request))
 	if code != common.CodeSuccess || err != nil {
 		writeChatChannelError(c, code, chatChannelErrMsg(code, err))
 		return
@@ -180,7 +187,8 @@ func (h *ChatChannelHandler) DeleteChatChannel(c *gin.Context) {
 		return
 	}

-	result, code, err := h.chatChannelService.DeleteChatChannel(userID, channelID)
+	ctx := c.Request.Context()
+	result, code, err := h.chatChannelService.DeleteChatChannel(ctx, userID, channelID)
 	if code != common.CodeSuccess || err != nil {
 		writeChatChannelError(c, code, chatChannelErrMsg(code, err))
 		return
--- a/internal/handler/chat_channel_test.go
+++ b/internal/handler/chat_channel_test.go
@@ -1,6 +1,7 @@
 package handler

 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"net/http"
@@ -22,35 +23,35 @@ type fakeChatChannelService struct {
 	deleteFn func(userID, channelID string) (bool, common.ErrorCode, error)
 }

-func (f fakeChatChannelService) CreateChatChannel(tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) {
+func (f fakeChatChannelService) CreateChatChannel(ctx context.Context, tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) {
 	if f.createFn == nil {
 		return nil, errors.New("unexpected CreateChatChannel call")
 	}
 	return f.createFn(tenantID, name, channelType, config, chatID)
 }

-func (f fakeChatChannelService) List(tenantID string) ([]*entity.ChatChannelListResponse, error) {
+func (f fakeChatChannelService) List(ctx context.Context, tenantID string) ([]*entity.ChatChannelListResponse, error) {
 	if f.listFn == nil {
 		return nil, errors.New("unexpected List call")
 	}
 	return f.listFn(tenantID)
 }

-func (f fakeChatChannelService) GetChatChannel(userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error) {
+func (f fakeChatChannelService) GetChatChannel(ctx context.Context, userID, channelID string) (*entity.ChatChannel, common.ErrorCode, error) {
 	if f.getFn == nil {
 		return nil, common.CodeServerError, errors.New("unexpected GetChatChannel call")
 	}
 	return f.getFn(userID, channelID)
 }

-func (f fakeChatChannelService) UpdateChatChannel(userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error) {
+func (f fakeChatChannelService) UpdateChatChannel(ctx context.Context, userID, channelID string, req map[string]interface{}) (*entity.ChatChannel, common.ErrorCode, error) {
 	if f.updateFn == nil {
 		return nil, common.CodeServerError, errors.New("unexpected UpdateChatChannel call")
 	}
 	return f.updateFn(userID, channelID, req)
 }

-func (f fakeChatChannelService) DeleteChatChannel(userID, channelID string) (bool, common.ErrorCode, error) {
+func (f fakeChatChannelService) DeleteChatChannel(ctx context.Context, userID, channelID string) (bool, common.ErrorCode, error) {
 	if f.deleteFn == nil {
 		return false, common.CodeServerError, errors.New("unexpected DeleteChatChannel call")
 	}
--- a/internal/handler/chat_session.go
+++ b/internal/handler/chat_session.go
@@ -68,7 +68,8 @@ func (h *ChatSessionHandler) ListChatSessions(c *gin.Context) {
 	}

 	// Call service to list chat sessions
-	result, err := h.chatSessionService.ListChatSessions(userID, chatID)
+	ctx := c.Request.Context()
+	result, err := h.chatSessionService.ListChatSessions(ctx, userID, chatID)
 	if err != nil {
 		// Check if it's an authorization error
 		if err.Error() == "Only owner of dialog authorized for this operation" {
@@ -247,7 +248,8 @@ func (h *ChatSessionHandler) GetSession(c *gin.Context) {
 	userID := user.ID
 	chatID, sessionID := c.Param("chat_id"), c.Param("session_id")

-	result, code, err := h.chatSessionService.GetSession(userID, chatID, sessionID)
+	ctx := c.Request.Context()
+	result, code, err := h.chatSessionService.GetSession(ctx, userID, chatID, sessionID)
 	if err != nil {
 		common.ErrorWithCode(c, code, err.Error())
 		return
@@ -288,7 +290,8 @@ func (h *ChatSessionHandler) CreateSession(c *gin.Context) {
 		req = map[string]interface{}{}
 	}

-	result, code, err := h.chatSessionService.CreateSession(userID, chatID, req)
+	ctx := c.Request.Context()
+	result, code, err := h.chatSessionService.CreateSession(ctx, userID, chatID, req)
 	if err != nil {
 		if code == common.CodeAuthenticationError {
 			common.ResponseWithCodeData(c, code, false, err.Error())
@@ -334,7 +337,8 @@ func (h *ChatSessionHandler) DeleteSessions(c *gin.Context) {
 		req = map[string]interface{}{}
 	}

-	result, message, code, err := h.chatSessionService.DeleteSessions(userID, chatID, req)
+	ctx := c.Request.Context()
+	result, message, code, err := h.chatSessionService.DeleteSessions(ctx, userID, chatID, req)
 	if err != nil {
 		if code == common.CodeAuthenticationError {
 			common.ResponseWithCodeData(c, code, false, err.Error())
@@ -372,7 +376,8 @@ func (h *ChatSessionHandler) UpdateSession(c *gin.Context) {
 		return
 	}

-	result, code, err := h.chatSessionService.UpdateSession(userID, chatID, sessionID, req)
+	ctx := c.Request.Context()
+	result, code, err := h.chatSessionService.UpdateSession(ctx, userID, chatID, sessionID, req)
 	if err != nil {
 		common.ErrorWithCode(c, code, err.Error())
 		return
@@ -389,7 +394,8 @@ func (h *ChatSessionHandler) DeleteSessionMessage(c *gin.Context) {
 	userID := user.ID
 	chatID, sessionID, msgID := c.Param("chat_id"), c.Param("session_id"), c.Param("msg_id")

-	result, code, err := h.chatSessionService.DeleteSessionMessage(userID, chatID, sessionID, msgID)
+	ctx := c.Request.Context()
+	result, code, err := h.chatSessionService.DeleteSessionMessage(ctx, userID, chatID, sessionID, msgID)
 	if err != nil {
 		if code == common.CodeAuthenticationError {
 			common.ResponseWithCodeData(c, code, false, err.Error())
--- a/internal/handler/mcp_server.go
+++ b/internal/handler/mcp_server.go
@@ -17,6 +17,7 @@
 package handler

 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -123,8 +124,8 @@ func MCPListDatasets(ds *dataset.DatasetService, userID string, page, pageSize i

 // MCPListChats wraps ChatService.ListChats for the MCP tool handler,
 // converting the typed response into a generic []map[string]interface{}.
-func MCPListChats(cs *service.ChatService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) {
-	resp, err := cs.ListChats(userID, "1", "", page, pageSize, orderby, desc, nil)
+func MCPListChats(ctx context.Context, chatService *service.ChatService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) {
+	resp, err := chatService.ListChats(ctx, userID, "1", "", page, pageSize, orderby, desc, nil)
 	if err != nil {
 		return nil, 0, err
 	}
--- a/internal/handler/stats.go
+++ b/internal/handler/stats.go
@@ -57,8 +57,10 @@ func (h *StatsHandler) GetStats(c *gin.Context) {
 		agentSource := "agent"
 		source = &agentSource
 	}
+	ctx := c.Request.Context()

-	stats, err := h.statsService.GetStats(user.ID, fromDate, toDate, source)
+	// Get stats
+	stats, err := h.statsService.GetStats(ctx, user.ID, fromDate, toDate, source)
 	if err != nil {
 		code := common.CodeExceptionError
 		if errors.Is(err, service.ErrTenantNotFound) {