From ed878930fb7e32d9fb6d45a592d90692e3566b76 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Thu, 28 May 2026 16:44:35 +0800 Subject: [PATCH] feat[Go]: implement delete/ rebuild/ listlog api for connector (#15300) ### What problem does this PR solve? implement delete, rebuild api for connector ### Type of change - [x] New Feature (non-breaking change which adds functionality) --- internal/dao/connector.go | 164 ++++++++++++++++++++++++++++ internal/entity/connector.go | 67 ++++++++++++ internal/handler/connector.go | 129 +++++++++++++++++++++- internal/handler/connector_test.go | 169 ++++++++++++++++++++++++++++- internal/router/router.go | 5 + internal/service/connector.go | 140 +++++++++++++++++++++++- 6 files changed, 668 insertions(+), 6 deletions(-) diff --git a/internal/dao/connector.go b/internal/dao/connector.go index 260e1596a9..312d56a37b 100644 --- a/internal/dao/connector.go +++ b/internal/dao/connector.go @@ -18,6 +18,9 @@ package dao import ( "ragflow/internal/entity" + "time" + + "gorm.io/gorm" ) // ConnectorDAO connector data access object @@ -103,3 +106,164 @@ func (dao *ConnectorDAO) UpdateByID(id string, updates map[string]interface{}) e func (dao *ConnectorDAO) DeleteByID(id string) error { return DB.Where("id = ?", id).Delete(&entity.Connector{}).Error } + +// CancelRunningOrScheduledLogs marks active sync logs as canceled for a connector. +func (dao *ConnectorDAO) CancelRunningOrScheduledLogs(connectorID string) error { + return DB.Model(&entity.SyncLogs{}). + Where("connector_id = ? AND status IN ?", connectorID, []string{string(entity.TaskStatusSchedule), string(entity.TaskStatusRunning)}). + Update("status", string(entity.TaskStatusCancel)).Error +} + +// ListDocumentsByKBAndSourceType lists connector documents in a dataset. +func (dao *ConnectorDAO) ListDocumentsByKBAndSourceType(kbID, sourceType string) ([]*entity.Document, error) { + var documents []*entity.Document + err := DB.Where("kb_id = ? AND source_type = ?", kbID, sourceType).Find(&documents).Error + return documents, err +} + +// RebuildConnector replaces old connector documents with scheduled sync tasks. +func (dao *ConnectorDAO) RebuildConnector(connector *entity.Connector, kbID string, documents []*entity.Document) error { + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Where("connector_id = ? AND kb_id = ?", connector.ID, kbID).Delete(&entity.SyncLogs{}).Error; err != nil { + return err + } + + if len(documents) > 0 { + docIDs := make([]string, 0, len(documents)) + var tokenNum int64 + var chunkNum int64 + for _, document := range documents { + docIDs = append(docIDs, document.ID) + tokenNum += document.TokenNum + chunkNum += document.ChunkNum + } + + var mappings []entity.File2Document + if err := tx.Where("document_id IN ?", docIDs).Find(&mappings).Error; err != nil { + return err + } + fileIDs := make([]string, 0, len(mappings)) + seenFileIDs := make(map[string]struct{}, len(mappings)) + for _, mapping := range mappings { + if mapping.FileID == nil || *mapping.FileID == "" { + continue + } + if _, ok := seenFileIDs[*mapping.FileID]; ok { + continue + } + seenFileIDs[*mapping.FileID] = struct{}{} + fileIDs = append(fileIDs, *mapping.FileID) + } + + if err := tx.Where("doc_id IN ?", docIDs).Delete(&entity.Task{}).Error; err != nil { + return err + } + if err := tx.Where("document_id IN ?", docIDs).Delete(&entity.File2Document{}).Error; err != nil { + return err + } + if len(fileIDs) > 0 { + if err := tx.Unscoped(). + Where("id IN ? AND source_type = ?", fileIDs, string(entity.FileSourceKnowledgebase)). + Delete(&entity.File{}).Error; err != nil { + return err + } + } + if err := tx.Where("id IN ?", docIDs).Delete(&entity.Document{}).Error; err != nil { + return err + } + if err := tx.Model(&entity.Knowledgebase{}). + Where("id = ?", kbID). + Updates(map[string]interface{}{ + "doc_num": gorm.Expr("doc_num - ?", len(docIDs)), + "token_num": gorm.Expr("token_num - ?", tokenNum), + "chunk_num": gorm.Expr("chunk_num - ?", chunkNum), + }).Error; err != nil { + return err + } + } + + if err := tx.Model(&entity.Connector{}). + Where("id = ?", connector.ID). + Update("status", string(entity.TaskStatusSchedule)).Error; err != nil { + return err + } + + if err := createRebuildSyncLog(tx, connector.ID, kbID, connectorTaskTypeSync, true); err != nil { + return err + } + if syncDeletedFiles, _ := connector.Config["sync_deleted_files"].(bool); syncDeletedFiles { + if err := createRebuildSyncLog(tx, connector.ID, kbID, connectorTaskTypePrune, false); err != nil { + return err + } + } + return nil + }) +} + +const ( + connectorTaskTypeSync = "sync" + connectorTaskTypePrune = "prune" +) + +func createRebuildSyncLog(tx *gorm.DB, connectorID, kbID, taskType string, reindex bool) error { + fromBeginning := "0" + if reindex { + fromBeginning = "1" + } + now := time.Now().Local() + return tx.Create(&entity.SyncLogs{ + ID: generateUUID(), + ConnectorID: connectorID, + KbID: kbID, + TaskType: taskType, + Status: string(entity.TaskStatusSchedule), + FromBeginning: &fromBeginning, + TimeStarted: &now, + ErrorMsg: "", + TotalDocsIndexed: 0, + }).Error +} + +// ListLogsByConnectorID lists sync logs for one connector with pagination. +func (dao *ConnectorDAO) ListLogsByConnectorID(connectorID string, offset, limit int) ([]*entity.ConnectorSyncLog, int64, error) { + baseQuery := DB.Model(&entity.SyncLogs{}). + Joins("JOIN connector ON sync_logs.connector_id = connector.id"). + Joins("JOIN connector2kb ON sync_logs.connector_id = connector2kb.connector_id AND sync_logs.kb_id = connector2kb.kb_id"). + Joins("JOIN knowledgebase ON sync_logs.kb_id = knowledgebase.id"). + Where("sync_logs.connector_id = ?", connectorID) + + var total int64 + if err := baseQuery.Distinct("sync_logs.id").Count(&total).Error; err != nil { + return nil, 0, err + } + + var logs []*entity.ConnectorSyncLog + err := baseQuery. + Select( + "sync_logs.id", + "sync_logs.connector_id", + "sync_logs.task_type", + "sync_logs.kb_id", + "sync_logs.update_date", + "sync_logs.new_docs_indexed", + "sync_logs.total_docs_indexed", + "sync_logs.docs_removed_from_index", + "sync_logs.error_msg", + "sync_logs.error_count", + "sync_logs.time_started", + "connector.refresh_freq AS refresh_freq", + "connector.prune_freq AS prune_freq", + "knowledgebase.name AS kb_name", + "sync_logs.status", + ). + Distinct(). + Order("sync_logs.update_time DESC"). + Offset(offset). + Limit(limit). + Scan(&logs).Error + if err != nil { + return nil, 0, err + } + + return logs, total, nil +} diff --git a/internal/entity/connector.go b/internal/entity/connector.go index f33c427b9f..0658b42f60 100644 --- a/internal/entity/connector.go +++ b/internal/entity/connector.go @@ -107,6 +107,7 @@ func (Connector2Kb) TableName() string { type SyncLogs struct { ID string `gorm:"column:id;primaryKey;size:32" json:"id"` ConnectorID string `gorm:"column:connector_id;size:32;index" json:"connector_id"` + TaskType string `gorm:"column:task_type;size:32;not null;default:sync;index" json:"task_type"` Status string `gorm:"column:status;size:128;not null;index" json:"status"` FromBeginning *string `gorm:"column:from_beginning;size:1" json:"from_beginning,omitempty"` NewDocsIndexed int64 `gorm:"column:new_docs_indexed;default:0" json:"new_docs_indexed"` @@ -126,3 +127,69 @@ type SyncLogs struct { func (SyncLogs) TableName() string { return "sync_logs" } + +// ConnectorSyncLog is the API projection used by the connector logs endpoint. +type ConnectorSyncLog struct { + ID string `gorm:"column:id" json:"id"` + ConnectorID string `gorm:"column:connector_id" json:"connector_id"` + TaskType string `gorm:"column:task_type" json:"task_type"` + KbID string `gorm:"column:kb_id" json:"kb_id"` + UpdateDate *time.Time `gorm:"column:update_date" json:"update_date,omitempty"` + NewDocsIndexed int64 `gorm:"column:new_docs_indexed" json:"new_docs_indexed"` + TotalDocsIndexed int64 `gorm:"column:total_docs_indexed" json:"total_docs_indexed"` + DocsRemovedFromIndex int64 `gorm:"column:docs_removed_from_index" json:"docs_removed_from_index"` + ErrorMsg string `gorm:"column:error_msg" json:"error_msg"` + ErrorCount int64 `gorm:"column:error_count" json:"error_count"` + TimeStarted *time.Time `gorm:"column:time_started" json:"time_started,omitempty"` + RefreshFreq int64 `gorm:"column:refresh_freq" json:"refresh_freq"` + PruneFreq int64 `gorm:"column:prune_freq" json:"prune_freq"` + KbName string `gorm:"column:kb_name" json:"kb_name"` + Status string `gorm:"column:status" json:"status"` +} + +// MarshalJSON formats datetime fields to match the Python API encoder. +func (c ConnectorSyncLog) MarshalJSON() ([]byte, error) { + type connectorSyncLogJSON struct { + ID string `json:"id"` + ConnectorID string `json:"connector_id"` + TaskType string `json:"task_type"` + KbID string `json:"kb_id"` + UpdateDate *string `json:"update_date,omitempty"` + NewDocsIndexed int64 `json:"new_docs_indexed"` + TotalDocsIndexed int64 `json:"total_docs_indexed"` + DocsRemovedFromIndex int64 `json:"docs_removed_from_index"` + ErrorMsg string `json:"error_msg"` + ErrorCount int64 `json:"error_count"` + TimeStarted *string `json:"time_started,omitempty"` + RefreshFreq int64 `json:"refresh_freq"` + PruneFreq int64 `json:"prune_freq"` + KbName string `json:"kb_name"` + Status string `json:"status"` + } + + return json.Marshal(connectorSyncLogJSON{ + ID: c.ID, + ConnectorID: c.ConnectorID, + TaskType: c.TaskType, + KbID: c.KbID, + UpdateDate: formatConnectorLogTime(c.UpdateDate), + NewDocsIndexed: c.NewDocsIndexed, + TotalDocsIndexed: c.TotalDocsIndexed, + DocsRemovedFromIndex: c.DocsRemovedFromIndex, + ErrorMsg: c.ErrorMsg, + ErrorCount: c.ErrorCount, + TimeStarted: formatConnectorLogTime(c.TimeStarted), + RefreshFreq: c.RefreshFreq, + PruneFreq: c.PruneFreq, + KbName: c.KbName, + Status: c.Status, + }) +} + +func formatConnectorLogTime(value *time.Time) *string { + if value == nil { + return nil + } + formatted := value.Format("2006-01-02 15:04:05") + return &formatted +} diff --git a/internal/handler/connector.go b/internal/handler/connector.go index 0ba9eb41cc..f92cf451a9 100644 --- a/internal/handler/connector.go +++ b/internal/handler/connector.go @@ -20,6 +20,7 @@ import ( "net/http" "ragflow/internal/common" "ragflow/internal/entity" + "strconv" "strings" "github.com/gin-gonic/gin" @@ -30,7 +31,10 @@ import ( type connectorService interface { ListConnectors(userID string) (*service.ListConnectorsResponse, error) CreateConnector(userID string, req *service.CreateConnectorRequest) (*entity.Connector, error) - GetConnector(connectorID string, userID string) (*entity.Connector, common.ErrorCode, error) + GetConnector(connectorID, userID string) (*entity.Connector, common.ErrorCode, error) + ListLog(connectorID, userID string, page, pageSize int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) + DeleteConnector(connectorID, userID string) (bool, common.ErrorCode, error) + RebuildConnector(connectorID, userID, kbID string) (bool, common.ErrorCode, error) } // ConnectorHandler connector handler @@ -108,6 +112,54 @@ func (h *ConnectorHandler) GetConnector(c *gin.Context) { }) } +// ListLogs list connector sync logs. +// @Summary List Connector Logs +// @Description List sync logs for a connector the current user can access +// @Tags connector +// @Accept json +// @Produce json +// @Success 200 {object} map[string]interface{} +// @Router /api/v1/connectors/{connector_id}/logs [get] +func (h *ConnectorHandler) ListLogs(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + page := 1 + if rawPage := strings.TrimSpace(c.DefaultQuery("page", "1")); rawPage != "" { + parsedPage, err := strconv.Atoi(rawPage) + if err != nil { + jsonError(c, common.CodeArgumentError, "page must be an integer") + return + } + page = parsedPage + } + + pageSize := 15 + if rawPageSize := strings.TrimSpace(c.DefaultQuery("page_size", "15")); rawPageSize != "" { + parsedPageSize, err := strconv.Atoi(rawPageSize) + if err != nil { + jsonError(c, common.CodeArgumentError, "page_size must be an integer") + return + } + pageSize = parsedPageSize + } + + logs, total, code, err := h.connectorService.ListLog(c.Param("connector_id"), user.ID, page, pageSize) + if err != nil { + jsonError(c, code, err.Error()) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeSuccess, + "data": gin.H{"total": total, "logs": logs}, + "message": "success", + }) +} + // CreateConnector create connector // @Summary create Connectors // @Description create a connectors for the current user @@ -174,3 +226,78 @@ func (h *ConnectorHandler) CreateConnector(c *gin.Context) { "message": "success", }) } + +// DeleteConnector delete connector +// @Description Detele Connector +// @Tags connector +// @Accept json +// @Produce json +func (h *ConnectorHandler) DeleteConnector(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + ok, code, err := h.connectorService.DeleteConnector(c.Param("connector_id"), user.ID) + if err != nil { + jsonError(c, code, err.Error()) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeSuccess, + "data": ok, + "message": "success", + }) +} + +// RebuildConnector rebuild connector +// @Summary Rebuild Connector +// @Description Trigger a rebuild for an accessible connector and knowledge base +// @Tags connector +// @Accept json +// @Produce json +// @Success 200 {object} map[string]interface{} +// @Router /connector/:connector_id/rebuild [post] +func (h *ConnectorHandler) RebuildConnector(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + // Parse request body to get kb_id + var req struct { + KbID string `json:"kb_id" binding:"required"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeDataError, + "data": nil, + "message": "required argument is missing: kb_id", + }) + return + } + + if strings.TrimSpace(req.KbID) == "" { + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeDataError, + "data": nil, + "message": "kb_id cannot be empty", + }) + return + } + + ok, code, err := h.connectorService.RebuildConnector(c.Param("connector_id"), user.ID, req.KbID) + if err != nil { + jsonError(c, code, err.Error()) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeSuccess, + "data": ok, + "message": "success", + }) +} diff --git a/internal/handler/connector_test.go b/internal/handler/connector_test.go index abc768a18f..c1e17edd23 100644 --- a/internal/handler/connector_test.go +++ b/internal/handler/connector_test.go @@ -6,6 +6,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "github.com/gin-gonic/gin" @@ -17,6 +18,8 @@ import ( type fakeConnectorService struct { connector *entity.Connector + logs []*entity.ConnectorSyncLog + total int64 code common.ErrorCode err error } @@ -36,11 +39,25 @@ func (s fakeConnectorService) GetConnector(string, string) (*entity.Connector, c return s.connector, common.CodeSuccess, nil } -func (s fakeConnectorService) UpdateConnector(string, string, *service.UpdateConnectorRequest) (*entity.Connector, common.ErrorCode, error) { +func (s fakeConnectorService) ListLog(string, string, int, int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) { if s.err != nil { - return nil, s.code, s.err + return nil, 0, s.code, s.err } - return s.connector, common.CodeSuccess, nil + return s.logs, s.total, common.CodeSuccess, nil +} + +func (s fakeConnectorService) DeleteConnector(string, string) (bool, common.ErrorCode, error) { + if s.err != nil { + return false, s.code, s.err + } + return true, common.CodeSuccess, nil +} + +func (s fakeConnectorService) RebuildConnector(string, string, string) (bool, common.ErrorCode, error) { + if s.err != nil { + return false, s.code, s.err + } + return true, common.CodeSuccess, nil } func TestConnectorHandlerGetConnector(t *testing.T) { @@ -106,3 +123,149 @@ func TestConnectorHandlerGetConnector(t *testing.T) { }) } } + +func TestConnectorHandlerDeleteConnector(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + service fakeConnectorService + wantCode common.ErrorCode + wantData interface{} + wantMsg string + }{ + { + name: "success", + service: fakeConnectorService{}, + wantCode: common.CodeSuccess, + wantData: true, + }, + { + name: "unauthorized", + service: fakeConnectorService{code: common.CodeAuthenticationError, err: fmt.Errorf("No authorization.")}, + wantCode: common.CodeAuthenticationError, + wantData: nil, + wantMsg: "No authorization.", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &ConnectorHandler{connectorService: tt.service} + router := gin.New() + router.DELETE("/api/v1/connectors/:connector_id", func(c *gin.Context) { + c.Set("user", &entity.User{ID: "tenant-1"}) + h.DeleteConnector(c) + }) + + resp := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodDelete, "/api/v1/connectors/connector-1", nil) + router.ServeHTTP(resp, req) + if resp.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", resp.Code, resp.Body.String()) + } + + var body map[string]interface{} + if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v", err) + } + if body["code"] != float64(tt.wantCode) { + t.Fatalf("code=%v body=%v", body["code"], body) + } + if body["data"] != tt.wantData { + t.Fatalf("data=%v body=%v", body["data"], body) + } + if tt.wantMsg != "" && body["message"] != tt.wantMsg { + t.Fatalf("message=%v", body["message"]) + } + }) + } +} + +func TestConnectorHandlerListLogs(t *testing.T) { + gin.SetMode(gin.TestMode) + + startedAt := time.Date(2026, 5, 28, 8, 30, 0, 0, time.Local) + updatedAt := time.Date(2026, 5, 28, 9, 0, 0, 0, time.Local) + + tests := []struct { + name string + service fakeConnectorService + wantCode common.ErrorCode + wantMsg string + wantTotal float64 + wantLogID string + }{ + { + name: "success", + service: fakeConnectorService{ + logs: []*entity.ConnectorSyncLog{{ + ID: "log-1", + ConnectorID: "connector-1", + TaskType: "sync", + KbID: "kb-1", + UpdateDate: &updatedAt, + NewDocsIndexed: 2, + TotalDocsIndexed: 10, + DocsRemovedFromIndex: 1, + ErrorMsg: "", + ErrorCount: 0, + TimeStarted: &startedAt, + RefreshFreq: 5, + PruneFreq: 5, + KbName: "Docs", + Status: "3", + }}, + total: 1, + }, + wantCode: common.CodeSuccess, + wantTotal: 1, + wantLogID: "log-1", + }, + { + name: "unauthorized", + service: fakeConnectorService{code: common.CodeAuthenticationError, err: fmt.Errorf("No authorization.")}, + wantCode: common.CodeAuthenticationError, + wantMsg: "No authorization.", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &ConnectorHandler{connectorService: tt.service} + router := gin.New() + router.GET("/api/v1/connectors/:connector_id/logs", func(c *gin.Context) { + c.Set("user", &entity.User{ID: "tenant-1"}) + h.ListLogs(c) + }) + + resp := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/connectors/connector-1/logs?page=2&page_size=5", nil) + router.ServeHTTP(resp, req) + if resp.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", resp.Code, resp.Body.String()) + } + + var body map[string]interface{} + if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v", err) + } + if body["code"] != float64(tt.wantCode) { + t.Fatalf("code=%v body=%v", body["code"], body) + } + if tt.wantMsg != "" && body["message"] != tt.wantMsg { + t.Fatalf("message=%v", body["message"]) + } + if tt.wantLogID != "" { + data := body["data"].(map[string]interface{}) + if data["total"] != tt.wantTotal { + t.Fatalf("total=%v body=%v", data["total"], body) + } + logs := data["logs"].([]interface{}) + if logs[0].(map[string]interface{})["id"] != tt.wantLogID { + t.Fatalf("logs=%v body=%v", logs, body) + } + } + }) + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 4b5d2f21a4..395ea34567 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -319,6 +319,9 @@ func (r *Router) Setup(engine *gin.Engine) { connector.GET("/", r.connectorHandler.ListConnectors) connector.POST("/", r.connectorHandler.CreateConnector) connector.GET("/:connector_id", r.connectorHandler.GetConnector) + connector.GET("/:connector_id/logs", r.connectorHandler.ListLogs) + connector.DELETE("/:connector_id", r.connectorHandler.DeleteConnector) + connector.POST("/:connector_id/rebuild", r.connectorHandler.RebuildConnector) } system := v1.Group("/system") @@ -421,6 +424,8 @@ func (r *Router) Setup(engine *gin.Engine) { connector := authorized.Group("/v1/connector") { connector.GET("/list", r.connectorHandler.ListConnectors) + connector.GET("/:connector_id", r.connectorHandler.GetConnector) + connector.POST("/:connector_id/rebuild", r.connectorHandler.RebuildConnector) } // File routes diff --git a/internal/service/connector.go b/internal/service/connector.go index 1de962832a..01ed41dc81 100644 --- a/internal/service/connector.go +++ b/internal/service/connector.go @@ -17,10 +17,12 @@ package service import ( + "context" "errors" "fmt" "ragflow/internal/common" "ragflow/internal/dao" + "ragflow/internal/engine" "ragflow/internal/entity" "gorm.io/gorm" @@ -39,6 +41,11 @@ type ConnectorService struct { userTenantDAO *dao.UserTenantDAO } +func (s *ConnectorService) Rebuild() { + //TODO implement me + panic("implement me") +} + // NewConnectorService create connector service func NewConnectorService() *ConnectorService { return &ConnectorService{ @@ -62,6 +69,29 @@ type CreateConnectorRequest struct { TimeoutSecs *int64 `json:"timeout_secs,omitempty"` } +// RebuildConnectorRequest rebuild connector request. +type RebuildConnectorRequest struct { + KbID string `json:"kb_id"` +} + +// canAccessConnector Test Authentication +func (s *ConnectorService) canAccessConnector(connector *entity.Connector, userID string) bool { + if connector.TenantID == userID { + return true + } + + _, err := s.userTenantDAO.FilterByUserIDAndTenantID(userID, connector.TenantID) + return err == nil +} + +// cancelConnectorTasks Stop connector tasks +func (s *ConnectorService) cancelConnectorTasks(connectorID string) error { + if err := s.connectorDAO.CancelRunningOrScheduledLogs(connectorID); err != nil { + return err + } + return s.connectorDAO.UpdateByID(connectorID, map[string]interface{}{"status": string(entity.TaskStatusCancel)}) +} + // CreateConnector creates a connector owned by the current user. // Equivalent to Python's create_connector endpoint. func (s *ConnectorService) CreateConnector(userID string, req *CreateConnectorRequest) (*entity.Connector, error) { @@ -101,7 +131,7 @@ func (s *ConnectorService) CreateConnector(userID string, req *CreateConnectorRe } // GetConnector returns one connector when the user can access its tenant. -func (s *ConnectorService) GetConnector(connectorID string, userID string) (*entity.Connector, common.ErrorCode, error) { +func (s *ConnectorService) GetConnector(connectorID, userID string) (*entity.Connector, common.ErrorCode, error) { if connectorID == "" { return nil, common.CodeDataError, fmt.Errorf("connector_id is required") } @@ -109,7 +139,7 @@ func (s *ConnectorService) GetConnector(connectorID string, userID string) (*ent connector, err := s.connectorDAO.GetByID(connectorID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + return nil, common.CodeDataError, fmt.Errorf("Can't find this Connector!") } return nil, common.CodeServerError, err } @@ -159,3 +189,109 @@ func (s *ConnectorService) ListConnectors(userID string) (*ListConnectorsRespons Connectors: connectors, }, nil } + +func (s *ConnectorService) DeleteConnector(connectorID, userID string) (bool, common.ErrorCode, error) { + if connectorID == "" { + return false, common.CodeDataError, fmt.Errorf("connector_id is required") + } + + connector, err := s.connectorDAO.GetByID(connectorID) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return false, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + } + return false, common.CodeServerError, err + } + + if !s.canAccessConnector(connector, userID) { + return false, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + } + + if err = s.cancelConnectorTasks(connector.ID); err != nil { + return false, common.CodeServerError, err + } + + if err = s.connectorDAO.DeleteByID(connector.ID); err != nil { + return false, common.CodeServerError, err + } + return true, common.CodeSuccess, nil +} + +// RebuildConnector schedules a rebuild for an accessible connector and knowledge base. +func (s *ConnectorService) RebuildConnector(connectorID, userID, kbID string) (bool, common.ErrorCode, error) { + if connectorID == "" { + return false, common.CodeDataError, fmt.Errorf("connector_id is required") + } + if kbID == "" { + return false, common.CodeArgumentError, fmt.Errorf("required argument is missing: kb_id") + } + + connector, err := s.connectorDAO.GetByID(connectorID) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return false, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + } + return false, common.CodeServerError, err + } + + if !s.canAccessConnector(connector, userID) { + return false, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + } + + sourceType := fmt.Sprintf("%s/%s", connector.Source, connector.ID) + documents, err := s.connectorDAO.ListDocumentsByKBAndSourceType(kbID, sourceType) + if err != nil { + return false, common.CodeServerError, err + } + + s.deleteConnectorDocumentChunks(connector.TenantID, kbID, documents) + + if err := s.connectorDAO.RebuildConnector(connector, kbID, documents); err != nil { + return false, common.CodeServerError, err + } + return true, common.CodeSuccess, nil +} + +func (s *ConnectorService) deleteConnectorDocumentChunks(tenantID, kbID string, documents []*entity.Document) { + docEngine := engine.Get() + if docEngine == nil { + return + } + + indexName := fmt.Sprintf("ragflow_%s", tenantID) + for _, document := range documents { + _, _ = docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": document.ID}, indexName, kbID) + } +} + +func (s *ConnectorService) ListLog(connectorID, userID string, page, pageSize int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) { + if connectorID == "" { + return nil, 0, common.CodeDataError, fmt.Errorf("connector_id is required") + } + + connector, err := s.connectorDAO.GetByID(connectorID) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, 0, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + } + return nil, 0, common.CodeServerError, err + } + + if !s.canAccessConnector(connector, userID) { + return nil, 0, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + } + + if page < 1 { + page = 1 + } + if pageSize <= 0 || pageSize > 100 { + pageSize = 15 + } + offset := (page - 1) * pageSize + + logs, total, err := s.connectorDAO.ListLogsByConnectorID(connectorID, offset, pageSize) + if err != nil { + return nil, 0, common.CodeServerError, fmt.Errorf("failed to fetch connector logs: %w", err) + } + return logs, total, common.CodeSuccess, nil +}