From 9c41ec8dbdf2d543e6cfd71a7cdbed7b6e732f8a Mon Sep 17 00:00:00 2001 From: adws <2834758851@qq.com> Date: Thu, 22 Oct 2026 16:23:44 +0800 Subject: [PATCH] refactor(admin): move ingestion task listing to EE (#21312) --- internal/admin/handler.go | 94 ---------- internal/admin/handler_ee.go | 30 ++++ internal/admin/router.go | 6 - internal/admin/router_ee.go | 3 + internal/admin/service.go | 13 -- internal/admin/service_ee.go | 79 +++----- internal/cli/admin_command.go | 159 ---------------- internal/cli/admin_command_ee.go | 48 +++++ internal/cli/admin_parser.go | 169 +----------------- internal/cli/parser.go | 2 +- internal/service/ingestion_task_service.go | 82 ++------- .../service/ingestion_task_service_test.go | 60 ------- 12 files changed, 121 insertions(+), 624 deletions(-) diff --git a/internal/admin/handler.go b/internal/admin/handler.go index db06da86c2..f8ffc2696d 100644 --- a/internal/admin/handler.go +++ b/internal/admin/handler.go @@ -25,7 +25,6 @@ import ( "ragflow/internal/common" "ragflow/internal/engine" "ragflow/internal/engine/kvrocks" - "ragflow/internal/handler" "ragflow/internal/server" "ragflow/internal/service" "ragflow/internal/storage" @@ -1016,99 +1015,6 @@ func (h *Handler) ShowMessageQueue(c *gin.Context) { common.SuccessWithData(c, result, "show message queue successfully") } -type RemoveIngestionTaskRequest struct { - Tasks []string `json:"tasks"` - Email *string `json:"email"` - Status *string `json:"status"` -} - -func (h *Handler) RemoveIngestionTasks(c *gin.Context) { - var req RemoveIngestionTaskRequest - if err := c.ShouldBindJSON(&req); err != nil { - common.ErrorWithCode(c, common.CodeBadRequest, "Task ID is required") - return - } - ctx := c.Request.Context() - if req.Email == nil && req.Status == nil { - tasks, err := h.service.RemoveIngestionTasks(ctx, req.Tasks) - if err != nil { - common.ErrorWithCode(c, handler.IngestionTaskErrorCode(err), err.Error()) - return - } - - common.SuccessWithData(c, tasks, "Remove tasks successfully") - } else { - tasks, err := h.service.RemoveIngestionTasksByCondition(ctx, req.Tasks, req.Email, req.Status) - if err != nil { - common.ErrorWithCode(c, handler.IngestionTaskErrorCode(err), err.Error()) - return - } - common.SuccessWithData(c, tasks, "Remove tasks successfully") - } -} - -type StopIngestionTaskRequest struct { - Tasks []string `json:"tasks"` - Email *string `json:"email"` - Status *string `json:"status"` -} - -func (h *Handler) StopIngestionTasks(c *gin.Context) { - var req StopIngestionTaskRequest - if err := c.ShouldBindJSON(&req); err != nil { - common.ErrorWithCode(c, common.CodeBadRequest, "Task ID is required") - return - } - ctx := c.Request.Context() - if req.Email == nil && req.Status == nil { - tasks, err := h.service.StopIngestionTasks(ctx, req.Tasks) - if err != nil { - common.ErrorWithCode(c, handler.IngestionTaskErrorCode(err), err.Error()) - return - } - var result []map[string]string - for _, task := range tasks { - result = append(result, map[string]string{ - "task_id": task.ID, - "status": task.Status, - }) - } - - common.SuccessWithData(c, result, "Stop tasks successfully") - } else { - tasks, err := h.service.StopIngestionTasksByCondition(ctx, req.Tasks, req.Email, req.Status) - if err != nil { - common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) - return - } - common.SuccessWithData(c, tasks, "Stop tasks successfully") - } -} - -type ListIngestionTasksRequest struct { - Email *string `json:"email"` - Status *string `json:"status"` -} - -// ListIngestionTasks handle list ingestion tasks -func (h *Handler) ListIngestionTasks(c *gin.Context) { - var err error - var tasks []map[string]interface{} - var req ListIngestionTasksRequest - ctx := c.Request.Context() - - if err = c.ShouldBindJSON(&req); err != nil { - tasks, err = h.service.ListIngestionTasks(ctx) - } else { - tasks, err = h.service.ListIngestionTasksByCondition(ctx, req.Email, req.Status) - } - - if err != nil { - common.ErrorWithCode(c, common.CodeServerError, err.Error()) - } - common.SuccessWithData(c, tasks, "Get all tasks") -} - func (h *Handler) ListIngestors(c *gin.Context) { serverList := GlobalServerStore.ListInfos() serverList = AppendServiceList(serverList) diff --git a/internal/admin/handler_ee.go b/internal/admin/handler_ee.go index 1ebe6d2c1d..d489efe6b1 100644 --- a/internal/admin/handler_ee.go +++ b/internal/admin/handler_ee.go @@ -1581,6 +1581,36 @@ func (h *Handler) ShowIngestionTasksSummary(c *gin.Context) { common.SuccessWithData(c, ingestionTasksSummary, "") } +// ListIngestionTasks requires the enterprise edition. +func (h *Handler) ListIngestionTasks(c *gin.Context) { + tasks, err := h.service.ListIngestionTasks(c.Request.Context()) + if err != nil { + common.ErrorWithCode(c, common.CodeServerError, err.Error()) + return + } + common.SuccessWithData(c, tasks, "") +} + +// RemoveIngestionTasks requires the enterprise edition. +func (h *Handler) RemoveIngestionTasks(c *gin.Context) { + result, err := h.service.RemoveIngestionTasks(c.Request.Context()) + if err != nil { + common.ErrorWithCode(c, common.CodeServerError, err.Error()) + return + } + common.SuccessWithData(c, result, "") +} + +// StopIngestionTasks requires the enterprise edition. +func (h *Handler) StopIngestionTasks(c *gin.Context) { + result, err := h.service.StopIngestionTasks(c.Request.Context()) + if err != nil { + common.ErrorWithCode(c, common.CodeServerError, err.Error()) + return + } + common.SuccessWithData(c, result, "") +} + // ShowDataSummary handle show data summary func (h *Handler) ShowDataSummary(c *gin.Context) { ctx := c.Request.Context() diff --git a/internal/admin/router.go b/internal/admin/router.go index b54b1c76ba..0a6e62d541 100644 --- a/internal/admin/router.go +++ b/internal/admin/router.go @@ -52,7 +52,6 @@ func (r *Router) Setup(engine *gin.Engine) { //admin.POST("/ingestion/tasks", r.handler.StartIngestionTask) //admin.DELETE("/ingestion", r.handler.CancelIngestionTask) // cancel ingestion - //admin.GET("/ingestion/tasks", r.handler.ListIngestionTasks) // Protected routes protected := admin.Group("") @@ -123,11 +122,6 @@ func (r *Router) Setup(engine *gin.Engine) { protected.GET("/all-models", r.handler.ListAllModels) protected.GET("/all-models/:model_name", r.handler.ShowModel) - // Ingestion tasks - protected.DELETE("/ingestion/tasks", r.handler.RemoveIngestionTasks) - protected.PUT("/ingestion/tasks", r.handler.StopIngestionTasks) - protected.GET("/ingestion/tasks", r.handler.ListIngestionTasks) - RegisterEERouter(protected, r) } } diff --git a/internal/admin/router_ee.go b/internal/admin/router_ee.go index af3d0c90a8..ab0388ad5a 100644 --- a/internal/admin/router_ee.go +++ b/internal/admin/router_ee.go @@ -111,6 +111,9 @@ func RegisterEERouter(protected *gin.RouterGroup, r *Router) { protected.GET("/users/plan/summary", r.handler.ShowUsersPlanSummary) protected.GET("/users/plan", r.handler.ShowUsersPlan) protected.GET("/users/quota/summary", r.handler.ShowUsersQuotaSummary) + protected.DELETE("/ingestion/tasks", r.handler.RemoveIngestionTasks) + protected.PUT("/ingestion/tasks", r.handler.StopIngestionTasks) + protected.GET("/ingestion/tasks", r.handler.ListIngestionTasks) protected.GET("/ingestion/tasks/summary", r.handler.ShowIngestionTasksSummary) protected.GET("/data/summary", r.handler.ShowDataSummary) protected.GET("/data/orphan", r.handler.ShowDataOrphan) diff --git a/internal/admin/service.go b/internal/admin/service.go index 2b451718de..98eae5e44c 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -120,19 +120,6 @@ func (s *Service) Logout(ctx context.Context, user interface{}) error { return nil } -// ListIngestionTasks list all ingestion tasks for admin user -func (s *Service) ListIngestionTasks(ctx context.Context) ([]map[string]interface{}, error) { - return s.ingestionTaskSvc.ListAllForAdmin(ctx) -} - -func (s *Service) RemoveIngestionTasks(ctx context.Context, tasks []string) ([]map[string]string, error) { - return s.ingestionTaskSvc.RemoveMany(ctx, tasks, nil) -} - -func (s *Service) StopIngestionTasks(ctx context.Context, tasks []string) ([]*entity.IngestionTask, error) { - return s.ingestionTaskSvc.RequestStopMany(ctx, tasks, nil) -} - // GetUserByToken get user by access token func (s *Service) GetUserByToken(ctx context.Context, token string) (*entity.User, error) { user, err := s.userDAO.GetByAccessToken(ctx, dao.DB, token) diff --git a/internal/admin/service_ee.go b/internal/admin/service_ee.go index 09e69d0294..a0ecdcbda0 100644 --- a/internal/admin/service_ee.go +++ b/internal/admin/service_ee.go @@ -1128,69 +1128,30 @@ func (s *Service) ListUserAPIKeys(ctx context.Context, username string) ([]map[s return result, nil } -func (s *Service) ListIngestionTasksByCondition(ctx context.Context, email, status *string) ([]map[string]interface{}, error) { - - if email == nil && status == nil { - return nil, fmt.Errorf("email or status are required") - } - - element := map[string]interface{}{ - "command": "list_ingestion_tasks_by_condition", - "error": "'List ingestion tasks by condition' is not supported", - } - - if email != nil { - element["email"] = *email - } - if status != nil { - element["status"] = *status - } - - return []map[string]interface{}{element}, nil +// ListIngestionTasks reports that admin task listing is not supported. +func (s *Service) ListIngestionTasks(ctx context.Context) ([]map[string]interface{}, error) { + return []map[string]interface{}{ + { + "command": "list_ingestion_tasks", + "error": "'list ingestion tasks' is not supported", + }, + }, nil } -func (s *Service) StopIngestionTasksByCondition(ctx context.Context, tasks []string, email, status *string) ([]map[string]interface{}, error) { - - if email == nil && status == nil { - return nil, fmt.Errorf("email or status are required") - } - - element := map[string]interface{}{ - "command": "stop_ingestion_tasks_by_condition", - "tasks": tasks, - "error": "'Stop ingestion tasks by condition' is not supported", - } - - if email != nil { - element["email"] = *email - } - if status != nil { - element["status"] = *status - } - - return []map[string]interface{}{element}, nil +// RemoveIngestionTasks reports that admin task removal is not supported. +func (s *Service) RemoveIngestionTasks(ctx context.Context) (map[string]interface{}, error) { + return map[string]interface{}{ + "command": "remove_ingestion_tasks", + "error": "'remove ingestion tasks' is not supported", + }, nil } -func (s *Service) RemoveIngestionTasksByCondition(ctx context.Context, tasks []string, email, status *string) ([]map[string]interface{}, error) { - - if email == nil && status == nil { - return nil, fmt.Errorf("email or status are required") - } - - element := map[string]interface{}{ - "command": "remove_ingestion_tasks_by_condition", - "tasks": tasks, - "error": "'Remove ingestion tasks by condition' is not supported", - } - - if email != nil { - element["email"] = *email - } - if status != nil { - element["status"] = *status - } - - return []map[string]interface{}{element}, nil +// StopIngestionTasks reports that admin task stopping is not supported. +func (s *Service) StopIngestionTasks(ctx context.Context) (map[string]interface{}, error) { + return map[string]interface{}{ + "command": "stop_ingestion_tasks", + "error": "'stop ingestion tasks' is not supported", + }, nil } func CheckClientStatus(serverName string, status *common.BaseMessage) error { diff --git a/internal/cli/admin_command.go b/internal/cli/admin_command.go index 8919a6fbff..c184f4b162 100644 --- a/internal/cli/admin_command.go +++ b/internal/cli/admin_command.go @@ -1085,61 +1085,6 @@ func (c *CLI) ListAdminIngestors(commandCount int, cmd *Command) (ResponseIf, er return HandleCommonResponse(resp, "list ingestors") } -func (c *CLI) ListAdminIngestionTasks(commandCount int, cmd *Command) (ResponseIf, error) { - if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil { - return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login") - } - - resp, err := c.AdminServerClient.Request(commandCount, "GET", "/admin/ingestion/tasks", "admin", nil, nil) - if err != nil { - return nil, fmt.Errorf("failed to list admin tasks: %w", err) - } - - return HandleCommonResponse(resp, "list admin tasks") -} - -func (c *CLI) AdminStopIngestionCommand(commandCount int, cmd *Command) (ResponseIf, error) { - if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil { - return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login") - } - - tasks, ok := cmd.Params["tasks"].([]string) - if !ok { - return nil, fmt.Errorf("uri not provided") - } - payload := map[string]interface{}{ - "tasks": tasks, - } - - resp, err := c.AdminServerClient.Request(commandCount, "PUT", "/admin/ingestion/tasks", "admin", nil, payload) - if err != nil { - return nil, fmt.Errorf("failed to stop ingestion: %w", err) - } - - return HandleCommonResponse(resp, "stop ingestion") -} - -func (c *CLI) AdminRemoveIngestionCommand(commandCount int, cmd *Command) (ResponseIf, error) { - if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil { - return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login") - } - - tasks, ok := cmd.Params["tasks"].([]string) - if !ok { - return nil, fmt.Errorf("uri not provided") - } - payload := map[string]interface{}{ - "tasks": tasks, - } - - resp, err := c.AdminServerClient.Request(commandCount, "DELETE", "/admin/ingestion/tasks", "admin", nil, payload) - if err != nil { - return nil, fmt.Errorf("failed to remove ingestion: %w", err) - } - - return HandleCommonResponse(resp, "remove ingestion") -} - func (c *CLI) AdminShutdownIngestor(commandCount int, cmd *Command) (ResponseIf, error) { if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil { return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login") @@ -2195,50 +2140,6 @@ func (c *CLI) AdminPurgeUsersCommand(commandCount int, cmd *Command) (ResponseIf return HandleCommonDataResponse(resp, "purge users data") } -func (c *CLI) AdminListUserIngestionTasksCommand(commandCount int, cmd *Command) (ResponseIf, error) { - - if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil { - return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login") - } - - userName, ok := cmd.Params["user_name"].(string) - if !ok { - return nil, fmt.Errorf("plan_name not provided") - } - - payload := map[string]interface{}{ - "email": userName, - } - - status, ok := cmd.Params["status"].(string) - if ok { - payload["status"] = status - } - - apiURL := "/admin/ingestion/tasks" - - resp, err := c.AdminServerClient.Request(commandCount, "GET", apiURL, "admin", nil, payload) - if err != nil { - return nil, fmt.Errorf("failed to list user %s ingestion tasks: %w", userName, err) - } - - if resp.StatusCode != 200 { - return nil, fmt.Errorf("failed to list user %s ingestion tasks: HTTP %d, body: %s", userName, resp.StatusCode, string(resp.Body)) - } - - var result OrderedCommonResponse - if err = json.Unmarshal(resp.Body, &result); err != nil { - return nil, fmt.Errorf("list user %s ingestion tasks failed: invalid JSON (%w)", userName, err) - } - - if result.Code != 0 { - return nil, fmt.Errorf("%s", result.Message) - } - - result.Duration = resp.Duration - return &result, nil -} - func (c *CLI) AdminListUserChatsCommand(commandCount int, cmd *Command) (ResponseIf, error) { if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil { @@ -2504,66 +2405,6 @@ func (c *CLI) AdminListUserLogsCommand(commandCount int, cmd *Command) (Response return HandleCommonResponse(resp, fmt.Sprintf("list user %s logs", userName)) } -func (c *CLI) AdminStopUserIngestionTasksCommand(commandCount int, cmd *Command) (ResponseIf, error) { - - if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil { - return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login") - } - - userName, ok := cmd.Params["user_name"].(string) - if !ok { - return nil, fmt.Errorf("user_name not provided") - } - - payload := map[string]interface{}{ - "email": userName, - } - - status, ok := cmd.Params["status"].(string) - if ok { - payload["status"] = status - } - - apiURL := "/admin/ingestion/tasks" - - resp, err := c.AdminServerClient.Request(commandCount, "PUT", apiURL, "admin", nil, payload) - if err != nil { - return nil, fmt.Errorf("failed to stop user %s ingestion tasks: %w", userName, err) - } - - return HandleCommonResponse(resp, fmt.Sprintf("stop user %s ingestion tasks", userName)) -} - -func (c *CLI) AdminRemoveUserIngestionTasksCommand(commandCount int, cmd *Command) (ResponseIf, error) { - - if c.Config.CLIMode != AdminMode || c.AdminServerClient.LoginToken == nil { - return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login") - } - - userName, ok := cmd.Params["user_name"].(string) - if !ok { - return nil, fmt.Errorf("user_name not provided") - } - - payload := map[string]interface{}{ - "email": userName, - } - - status, ok := cmd.Params["status"].(string) - if ok { - payload["status"] = status - } - - apiURL := "/admin/ingestion/tasks" - - resp, err := c.AdminServerClient.Request(commandCount, "DELETE", apiURL, "admin", nil, payload) - if err != nil { - return nil, fmt.Errorf("failed to remove user %s ingestion tasks: %w", userName, err) - } - - return HandleCommonResponse(resp, fmt.Sprintf("remove user %s ingestion tasks", userName)) -} - // AdminAddProviderCommand add provider func (c *CLI) AdminAddProviderCommand(commandCount int, cmd *Command) (ResponseIf, error) { diff --git a/internal/cli/admin_command_ee.go b/internal/cli/admin_command_ee.go index 7fba1a0e76..21afe7a3ad 100644 --- a/internal/cli/admin_command_ee.go +++ b/internal/cli/admin_command_ee.go @@ -35,3 +35,51 @@ func (c *CLI) AdminListUserDatasetsCommand(commandCount int, cmd *Command) (Resp func (p *Parser) parseAdminListUserDatasetsEE(userName string) (*Command, error) { return nil, errors.New("listing user datasets is only available in the enterprise edition") } + +func (c *CLI) ListAdminIngestionTasks(commandCount int, cmd *Command) (ResponseIf, error) { + return nil, errors.New("listing admin ingestion tasks is not supported in the open-source edition") +} + +func (p *Parser) parseAdminListIngestionTasks() (*Command, error) { + return nil, errors.New("listing admin ingestion tasks is not supported in the open-source edition") +} + +func (c *CLI) AdminListUserIngestionTasksCommand(commandCount int, cmd *Command) (ResponseIf, error) { + return nil, errors.New("listing user ingestion tasks in admin mode is not supported in the open-source edition") +} + +func (p *Parser) parseAdminListUserIngestionTasks(userName string) (*Command, error) { + return nil, errors.New("listing user ingestion tasks in admin mode is not supported in the open-source edition") +} + +func (c *CLI) AdminStopIngestionCommand(commandCount int, cmd *Command) (ResponseIf, error) { + return nil, errors.New("stopping admin ingestion tasks is not supported in the open-source edition") +} + +func (c *CLI) AdminRemoveIngestionCommand(commandCount int, cmd *Command) (ResponseIf, error) { + return nil, errors.New("removing admin ingestion tasks is not supported in the open-source edition") +} + +func (c *CLI) AdminStopUserIngestionTasksCommand(commandCount int, cmd *Command) (ResponseIf, error) { + return nil, errors.New("stopping user ingestion tasks in admin mode is not supported in the open-source edition") +} + +func (c *CLI) AdminRemoveUserIngestionTasksCommand(commandCount int, cmd *Command) (ResponseIf, error) { + return nil, errors.New("removing user ingestion tasks in admin mode is not supported in the open-source edition") +} + +func (p *Parser) parseAdminStopIngestionTasks() (*Command, error) { + return nil, errors.New("stopping admin ingestion tasks is not supported in the open-source edition") +} + +func (p *Parser) parseAdminRemoveIngestionTasks() (*Command, error) { + return nil, errors.New("removing admin ingestion tasks is not supported in the open-source edition") +} + +func (p *Parser) parseAdminStopUserIngestionTasks(userName string) (*Command, error) { + return nil, errors.New("stopping user ingestion tasks in admin mode is not supported in the open-source edition") +} + +func (p *Parser) parseAdminRemoveUserIngestionTasks(userName string) (*Command, error) { + return nil, errors.New("removing user ingestion tasks in admin mode is not supported in the open-source edition") +} diff --git a/internal/cli/admin_parser.go b/internal/cli/admin_parser.go index 2b79a71422..49d199e0c9 100644 --- a/internal/cli/admin_parser.go +++ b/internal/cli/admin_parser.go @@ -255,24 +255,6 @@ func (p *Parser) parseAdminListIngestors() (*Command, error) { return cmd, nil } -func (p *Parser) parseAdminListIngestionTasks() (*Command, error) { - p.nextToken() // consume Ingestion - - if p.curToken.Type != TokenTasks { - return nil, fmt.Errorf("expected TASKS") - } - p.nextToken() // consume TASKS - - cmd := NewCommand("list_admin_ingestion_tasks") - - // Semicolon is optional - if p.curToken.Type == TokenSemicolon { - p.nextToken() - } - - return cmd, nil -} - // LIST PROVIDER 'provider_name' MODELS; func (p *Parser) parseAdminListProviderCommands() (*Command, error) { p.nextToken() // consume PROVIDER @@ -1242,69 +1224,17 @@ func (p *Parser) parseAdminCheckProviderInstance(providerName string) (*Command, return cmd, nil } -// STOP INGESTION TASKS 'task_id1 task_id2'; -func (p *Parser) parseAdminStopIngestionTasks() (*Command, error) { +// STOP commands dispatch ingestion-specific parsing to the EE implementation. +func (p *Parser) parseAdminStopCommands() (*Command, error) { p.nextToken() // consume STOP - - var cmd *Command - switch p.curToken.Type { case TokenIngestion: - p.nextToken() - if p.curToken.Type != TokenTasks { - return nil, fmt.Errorf("expected TASKS") - } - p.nextToken() // consume TASK - - taskString, err := p.parseQuotedString() - if err != nil { - return nil, err - } - - tasks := strings.Split(taskString, " ") - p.nextToken() // consume TASK - - cmd = NewCommand("admin_stop_ingestion_tasks") - cmd.Params["tasks"] = tasks + return p.parseAdminStopIngestionTasks() case TokenUser: return p.parseAdminStopUserCommand() default: return nil, fmt.Errorf("expected USER or INGESTION") } - - // Semicolon is optional - if p.curToken.Type == TokenSemicolon { - p.nextToken() - } - - return cmd, nil -} - -func (p *Parser) parseAdminRemoveIngestionTasks() (*Command, error) { - p.nextToken() // consume INGESTION - - if p.curToken.Type != TokenTasks { - return nil, fmt.Errorf("expected TASKS") - } - p.nextToken() // consume TASKS - - taskString, err := p.parseQuotedString() - if err != nil { - return nil, err - } - - tasks := strings.Split(taskString, " ") - p.nextToken() // consume TASKS - - cmd := NewCommand("admin_remove_ingestion_tasks") - cmd.Params["tasks"] = tasks - - // Semicolon is optional - if p.curToken.Type == TokenSemicolon { - p.nextToken() - } - - return cmd, nil } // region CREATE commands @@ -2649,30 +2579,6 @@ func (p *Parser) parseStartIngestion() (*Command, error) { return cmd, nil } -func (p *Parser) parseStopIngestion() (*Command, error) { - p.nextToken() // consume Stop - - if p.curToken.Type != TokenIngestion { - return nil, fmt.Errorf("expect INGESTION") - } - p.nextToken() // consume Ingest - - taskID, err := p.parseQuotedString() - if err != nil { - return nil, err - } - - cmd := NewCommand("admin_stop_ingestion_command") - cmd.Params["task_id"] = taskID - p.nextToken() - - // Semicolon is optional - if p.curToken.Type == TokenSemicolon { - p.nextToken() - } - return cmd, nil -} - func (p *Parser) parseAdminIngestCommand() (*Command, error) { p.nextToken() // consume Ingest @@ -3300,31 +3206,6 @@ func (p *Parser) parseAdminListUserCommand() (*Command, error) { return cmd, nil } -// LIST USER 'user@example.com' INGESTION TASKS 'status'; -func (p *Parser) parseAdminListUserIngestionTasks(userName string) (*Command, error) { - p.nextToken() // consume INGESTION - - if p.curToken.Type != TokenTasks { - return nil, fmt.Errorf("expected TASKS after INGESTION") - } - p.nextToken() - - cmd := NewCommand("admin_list_user_ingestion_tasks") - cmd.Params["user_name"] = userName - - if p.curToken.Type == TokenQuotedString { - var status string - var err error - status, err = p.parseQuotedString() - if err != nil { - return nil, err - } - cmd.Params["status"] = status - p.nextToken() - } - return cmd, nil -} - // LIST USER 'user_name' PROVIDER 'provider_name' INSTANCES; // LIST USER 'user_name' PROVIDER 'provider_name' INSTANCE 'instance_name' MODELS; func (p *Parser) parseAdminListUserProviderInstanceModels(userName string) (*Command, error) { @@ -3421,71 +3302,31 @@ func (p *Parser) parseAdminListUserLogs(userName string) (*Command, error) { // STOP USER 'user@example.com' INGESTION TASKS 'created'; func (p *Parser) parseAdminStopUserCommand() (*Command, error) { p.nextToken() // consume USER - userName, err := p.parseQuotedString() if err != nil { return nil, err } p.nextToken() - - var cmd *Command switch p.curToken.Type { case TokenIngestion: - p.nextToken() - if p.curToken.Type != TokenTasks { - return nil, fmt.Errorf("expected TASKS after INGESTION") - } - p.nextToken() - cmd = NewCommand("admin_stop_user_ingestion_tasks_command") - if p.curToken.Type == TokenQuotedString { - var status string - status, err = p.parseQuotedString() - if err != nil { - return nil, err - } - cmd.Params["status"] = status - p.nextToken() - } + return p.parseAdminStopUserIngestionTasks(userName) default: return nil, fmt.Errorf("expected INGESTION after USER") } - - cmd.Params["user_name"] = userName - return cmd, nil } // REMOVE USER 'user@example.com' INGESTION TASKS 'created'; func (p *Parser) parseAdminRemoveUserCommand() (*Command, error) { p.nextToken() // consume USER - userName, err := p.parseQuotedString() if err != nil { return nil, err } p.nextToken() - - var cmd *Command switch p.curToken.Type { case TokenIngestion: - p.nextToken() - if p.curToken.Type != TokenTasks { - return nil, fmt.Errorf("expected TASKS after INGESTION") - } - p.nextToken() - cmd = NewCommand("admin_remove_user_ingestion_tasks_command") - if p.curToken.Type == TokenQuotedString { - var status string - status, err = p.parseQuotedString() - if err != nil { - return nil, err - } - cmd.Params["status"] = status - p.nextToken() - } + return p.parseAdminRemoveUserIngestionTasks(userName) default: return nil, fmt.Errorf("expected INGESTION after USER") } - - cmd.Params["user_name"] = userName - return cmd, nil } diff --git a/internal/cli/parser.go b/internal/cli/parser.go index 426634c637..5dbee1483c 100644 --- a/internal/cli/parser.go +++ b/internal/cli/parser.go @@ -125,7 +125,7 @@ func (p *Parser) parseAdminCommand() (*Command, error) { case TokenRemove: return p.parseAdminRemoveCommands() case TokenStop: - return p.parseAdminStopIngestionTasks() + return p.parseAdminStopCommands() case TokenAdd: return p.parseAdminAddCommand() case TokenDelete: diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go index 1a42483ced..8f87c600f3 100644 --- a/internal/service/ingestion_task_service.go +++ b/internal/service/ingestion_task_service.go @@ -63,28 +63,24 @@ func (e *TaskStatusConflictError) Error() string { } type IngestionTaskService struct { - documentDAO *dao.DocumentDAO - userDAO *dao.UserDAO - ingestionTaskDAO *dao.IngestionTaskDAO - ingestionTaskLogDAO *dao.IngestionTaskLogDAO - pipelineLogDAO *dao.PipelineOperationLogDAO - kbDAO *dao.KnowledgebaseDAO - userCanvasDAO *dao.UserCanvasDAO - taskPublisher TaskPublisher - logSettings IngestionLogSettings + documentDAO *dao.DocumentDAO + ingestionTaskDAO *dao.IngestionTaskDAO + pipelineLogDAO *dao.PipelineOperationLogDAO + kbDAO *dao.KnowledgebaseDAO + userCanvasDAO *dao.UserCanvasDAO + taskPublisher TaskPublisher + logSettings IngestionLogSettings } func NewIngestionTaskService() *IngestionTaskService { return &IngestionTaskService{ - documentDAO: dao.NewDocumentDAO(), - userDAO: dao.NewUserDAO(), - ingestionTaskDAO: dao.NewIngestionTaskDAO(), - ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), - pipelineLogDAO: dao.NewPipelineOperationLogDAO(), - kbDAO: dao.NewKnowledgebaseDAO(), - userCanvasDAO: dao.NewUserCanvasDAO(), - taskPublisher: NewMessageQueueTaskPublisher(), - logSettings: DefaultIngestionLogSettings(), + documentDAO: dao.NewDocumentDAO(), + ingestionTaskDAO: dao.NewIngestionTaskDAO(), + pipelineLogDAO: dao.NewPipelineOperationLogDAO(), + kbDAO: dao.NewKnowledgebaseDAO(), + userCanvasDAO: dao.NewUserCanvasDAO(), + taskPublisher: NewMessageQueueTaskPublisher(), + logSettings: DefaultIngestionLogSettings(), } } @@ -181,56 +177,6 @@ func (s *IngestionTaskService) RemoveMany(ctx context.Context, tasks []string, o return deletedTasks, nil } -func (s *IngestionTaskService) ListAllForAdmin(ctx context.Context) ([]map[string]interface{}, error) { - ingestionTasks, err := s.ingestionTaskDAO.GetAllTasks(ctx, dao.DB, 0, 0) - if err != nil { - return nil, err - } - - showTasks := make([]map[string]interface{}, 0, len(ingestionTasks)) - for _, task := range ingestionTasks { - var user *entity.User - user, err = s.userDAO.GetByTenantID(ctx, dao.DB, task.UserID) - if err != nil { - return nil, err - } - - showTask := map[string]interface{}{ - "id": task.ID, - "user_id": task.UserID, - "user": user.Email, - "document_id": task.DocumentID, - "status": task.Status, - } - - if task.PipelineLogID != nil && *task.PipelineLogID != "" { - if run, runErr := s.pipelineLogDAO.GetByID(ctx, dao.DB, *task.PipelineLogID); runErr == nil && run.RunCount != nil && *run.RunCount > 0 { - showTask["run_count"] = *run.RunCount - } - } - - showTask["component_total"] = task.ComponentTotal - if task.ComponentTotal > 0 { - progress := (*dao.TaskProgress)(nil) - if task.PipelineLogID != nil && *task.PipelineLogID != "" { - progress, err = s.ingestionTaskLogDAO.AggregateProgressByPipelineLogID(ctx, dao.DB, *task.PipelineLogID, task.ComponentTotal) - } else { - err = errors.New("task has no pipeline log identity") - } - if err == nil { - showTask["component_done"] = progress.Done - } else { - showTask["component_done"] = 0 - } - } else { - showTask["component_done"] = 0 - } - - showTasks = append(showTasks, showTask) - } - return showTasks, nil -} - // TransitionTaskToRunning performs only the task-state transition required // before worker identity validation. It intentionally does not reset the // document or advance a pipeline log. diff --git a/internal/service/ingestion_task_service_test.go b/internal/service/ingestion_task_service_test.go index bed0337a3c..2fee206a6d 100644 --- a/internal/service/ingestion_task_service_test.go +++ b/internal/service/ingestion_task_service_test.go @@ -426,66 +426,6 @@ func TestIngestionTaskServiceRemoveManyRemovesOwnedTasks(t *testing.T) { } } -func TestIngestionTaskServiceListAllForAdminIncludesRunAndUserEmail(t *testing.T) { - db := setupServiceTestDB(t) - pushServiceDB(t, db) - status := "1" - if err := dao.DB.Create(&entity.User{ - ID: "user-1", - Email: "user-1@test.com", - Nickname: "user-1", - IsAuthenticated: "1", - IsActive: "1", - IsAnonymous: "0", - Status: &status, - }).Error; err != nil { - t.Fatalf("insert user: %v", err) - } - insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") - runCount := 3 - if err := dao.DB.Create(&entity.PipelineOperationLog{ - ID: "run-1", - DocumentID: "doc-1", - TenantID: "tenant-1", - KbID: "kb-1", - TaskType: string(entity.PipelineTaskTypeParse), - OperationStatus: string(entity.TaskStatusRunning), - RunCount: &runCount, - }).Error; err != nil { - t.Fatalf("insert pipeline operation log: %v", err) - } - if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("pipeline_log_id", "run-1").Error; err != nil { - t.Fatalf("bind pipeline operation log: %v", err) - } - if err := dao.DB.Exec(`INSERT INTO ingestion_task_log - (task_id, pipeline_log_id, checkpoint, event_type, component, phase, message) - VALUES (?, ?, '{}', ?, ?, ?, ?)`, "task-1", "run-1", dao.EventTypeLifecycle, "Parser", 1, "Parser Done").Error; err != nil { - t.Fatalf("insert task log: %v", err) - } - - ctx := t.Context() - svc := NewIngestionTaskService() - tasks, err := svc.ListAllForAdmin(ctx) - if err != nil { - t.Fatalf("ListAllForAdmin failed: %v", err) - } - if len(tasks) != 1 { - t.Fatalf("expected 1 task, got %d", len(tasks)) - } - if tasks[0]["user"] != "user-1@test.com" { - t.Fatalf("user = %v, want user-1@test.com", tasks[0]["user"]) - } - if tasks[0]["run_count"] != 3 { - t.Fatalf("run_count = %v, want 3", tasks[0]["run_count"]) - } - if tasks[0]["component_total"] != 0 { - t.Fatalf("component_total = %v, want 0", tasks[0]["component_total"]) - } - if tasks[0]["component_done"] != 0 { - t.Fatalf("component_done = %v, want 0", tasks[0]["component_done"]) - } -} - func TestIngestionTaskServiceStartRunningTransitionsCreatedTask(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db)