mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-14 00:48:26 +08:00
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)
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user