mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-20 14:41:05 +08:00
feat: migrate POST /api/v1/datasets/<dataset_id>/documents/stop to Go (#15597)
## Summary
Migrate the stop parse documents endpoint from Python to Go.
### Python endpoint
`POST /api/v1/datasets/<dataset_id>/documents/stop` —
`api/apps/restful_apis/document_api.py:1542-1641`
### Changes
| File | Change |
|------|--------|
| `internal/dao/task.go` | Add `GetByDocID` method |
| `internal/dao/task_test.go` | 3 DAO tests (new file) |
| `internal/service/document.go` | Add `StopParseDocuments` + refactor
shared helpers |
| `internal/service/document_test.go` | 8 service tests |
| `internal/handler/document.go` | Add handler + request struct +
interface |
| `internal/handler/document_test.go` | 5 handler tests |
| `internal/router/router.go` | Add `POST /:dataset_id/documents/stop`
route |
### How it works
1. Validates all document IDs belong to the dataset
2. For each document in RUNNING/CANCEL state (or with unfinished tasks):
- Sets Redis cancel signal `{task_id}-cancel` for each associated task
- Updates `document.run` to CANCEL ("2")
3. Returns `{"success_count": N, "errors": [...]}`
### Test strategy
- **DAO/Service**: SQLite in-memory DB, zero mocks. Redis is nil-safe by
design.
- **Handler**: `fakeDocumentService` implementing `documentServiceIface`
interface.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
This commit is contained in:
@@ -44,6 +44,7 @@ type documentServiceIface interface {
|
||||
DeleteDocument(id string) error
|
||||
DeleteDocuments(ids []string, deleteAll bool, datasetID, userID string) (int, error)
|
||||
ParseDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
|
||||
StopParseDocuments(datasetID string, docIDs []string) (map[string]interface{}, error)
|
||||
ListDocuments(page, pageSize int) ([]*service.DocumentResponse, int64, error)
|
||||
ListDocumentsByDatasetID(kbID string, page, pageSize int) ([]*entity.DocumentListItem, int64, error)
|
||||
GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*service.DocumentResponse, int64, error)
|
||||
@@ -842,3 +843,46 @@ func (h *DocumentHandler) ParseDocuments(c *gin.Context) {
|
||||
"data": parseResult,
|
||||
})
|
||||
}
|
||||
|
||||
type StopParseDocumentRequest struct {
|
||||
DocumentIDs []string `json:"document_ids" binding:"required"`
|
||||
}
|
||||
|
||||
func (h *DocumentHandler) StopParseDocuments(c *gin.Context) {
|
||||
datasetID := c.Param("dataset_id")
|
||||
|
||||
var req StopParseDocumentRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": common.CodeBadRequest,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if len(req.DocumentIDs) == 0 {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": common.CodeBadRequest,
|
||||
"message": "`document_ids` is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
userID := c.GetString("user_id")
|
||||
|
||||
if !h.datasetService.Accessible(datasetID, userID) {
|
||||
jsonError(c, common.CodeAuthenticationError, "You don't own the dataset.")
|
||||
return
|
||||
}
|
||||
|
||||
result, err := h.documentService.StopParseDocuments(datasetID, req.DocumentIDs)
|
||||
if err != nil {
|
||||
jsonError(c, common.CodeExceptionError, err.Error())
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": 0,
|
||||
"message": "success",
|
||||
"data": result,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -24,17 +24,22 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/service"
|
||||
)
|
||||
|
||||
// fakeDocumentService implements documentServiceIface for handler tests.
|
||||
type fakeDocumentService struct {
|
||||
deleted int
|
||||
err error
|
||||
deleted int
|
||||
err error
|
||||
stopResult map[string]interface{}
|
||||
stopErr error
|
||||
}
|
||||
|
||||
func (f *fakeDocumentService) CreateDocument(req *service.CreateDocumentRequest) (*entity.Document, error) {
|
||||
@@ -55,6 +60,9 @@ func (f *fakeDocumentService) DeleteDocuments(ids []string, deleteAll bool, data
|
||||
func (f *fakeDocumentService) ParseDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeDocumentService) StopParseDocuments(datasetID string, docIDs []string) (map[string]interface{}, error) {
|
||||
return f.stopResult, f.stopErr
|
||||
}
|
||||
func (f *fakeDocumentService) ListDocuments(page, pageSize int) ([]*service.DocumentResponse, int64, error) {
|
||||
return nil, 0, nil
|
||||
}
|
||||
@@ -244,3 +252,193 @@ func TestDeleteDocumentsHandler_MissingDatasetID(t *testing.T) {
|
||||
t.Fatal("expected error for missing dataset_id")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopParseDocumentsHandler_EmptyDocIDs(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
fake := &fakeDocumentService{}
|
||||
h := &DocumentHandler{
|
||||
documentService: fake,
|
||||
datasetService: service.NewDatasetService(),
|
||||
}
|
||||
|
||||
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `{"document_ids": []}`)
|
||||
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
||||
|
||||
h.StopParseDocuments(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
code, _ := resp["code"].(float64)
|
||||
if code == float64(common.CodeSuccess) {
|
||||
t.Fatal("expected error for empty document_ids")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopParseDocumentsHandler_BadJSON(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
fake := &fakeDocumentService{}
|
||||
h := &DocumentHandler{
|
||||
documentService: fake,
|
||||
datasetService: service.NewDatasetService(),
|
||||
}
|
||||
|
||||
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `not json`)
|
||||
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
||||
|
||||
h.StopParseDocuments(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
code, _ := resp["code"].(float64)
|
||||
if code == float64(common.CodeSuccess) {
|
||||
t.Fatal("expected error for bad JSON body")
|
||||
}
|
||||
}
|
||||
|
||||
// setupHandlerAccessDB sets up SQLite in-memory DB for handler tests that need
|
||||
// datasetService.Accessible to work.
|
||||
func setupHandlerAccessDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
|
||||
TranslateError: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to open sqlite: %v", err)
|
||||
}
|
||||
|
||||
if err := db.AutoMigrate(
|
||||
&entity.User{},
|
||||
&entity.Tenant{},
|
||||
&entity.UserTenant{},
|
||||
&entity.Knowledgebase{},
|
||||
); err != nil {
|
||||
t.Fatalf("failed to migrate: %v", err)
|
||||
}
|
||||
|
||||
// Insert user
|
||||
db.Create(&entity.User{ID: "user-1", Nickname: "test", Email: "test@test.com", Password: sptr("x")})
|
||||
// Insert tenant
|
||||
db.Create(&entity.Tenant{ID: "tenant-1", LLMID: "llm-1", EmbdID: "embd-1", ASRID: "asr-1"})
|
||||
// Insert user_tenant mapping
|
||||
db.Create(&entity.UserTenant{ID: "ut-1", UserID: "user-1", TenantID: "tenant-1", Role: "admin"})
|
||||
// Insert knowledgebase
|
||||
db.Create(&entity.Knowledgebase{
|
||||
ID: "ds-1", TenantID: "tenant-1", Name: "test-kb", EmbdID: "embd-1",
|
||||
CreatedBy: "user-1", Permission: string(entity.TenantPermissionMe),
|
||||
Status: sptr(string(entity.StatusValid)),
|
||||
})
|
||||
|
||||
return db
|
||||
}
|
||||
|
||||
// sptr returns a pointer to the given string (copy of service test helper).
|
||||
func sptr(s string) *string { return &s }
|
||||
|
||||
func TestStopParseDocumentsHandler_Success(t *testing.T) {
|
||||
db := setupHandlerAccessDB(t)
|
||||
orig := dao.DB
|
||||
dao.DB = db
|
||||
t.Cleanup(func() { dao.DB = orig })
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
fake := &fakeDocumentService{
|
||||
stopResult: map[string]interface{}{"success_count": 1},
|
||||
}
|
||||
h := &DocumentHandler{
|
||||
documentService: fake,
|
||||
datasetService: service.NewDatasetService(),
|
||||
}
|
||||
|
||||
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `{"document_ids": ["doc-1"]}`)
|
||||
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
||||
|
||||
h.StopParseDocuments(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var resp map[string]interface{}
|
||||
json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
if resp["code"] != float64(common.CodeSuccess) {
|
||||
t.Fatalf("expected code 0, got %v: %v", resp["code"], resp)
|
||||
}
|
||||
data := resp["data"].(map[string]interface{})
|
||||
if data["success_count"] != float64(1) {
|
||||
t.Fatalf("expected success_count=1, got %v", data["success_count"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopParseDocumentsHandler_ServiceError(t *testing.T) {
|
||||
db := setupHandlerAccessDB(t)
|
||||
orig := dao.DB
|
||||
dao.DB = db
|
||||
t.Cleanup(func() { dao.DB = orig })
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
fake := &fakeDocumentService{
|
||||
stopErr: fmt.Errorf("internal failure"),
|
||||
}
|
||||
h := &DocumentHandler{
|
||||
documentService: fake,
|
||||
datasetService: service.NewDatasetService(),
|
||||
}
|
||||
|
||||
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `{"document_ids": ["doc-1"]}`)
|
||||
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
||||
|
||||
h.StopParseDocuments(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
code, _ := resp["code"].(float64)
|
||||
if code == float64(common.CodeSuccess) {
|
||||
t.Fatal("expected error code for service error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopParseDocumentsHandler_NotAccessible(t *testing.T) {
|
||||
db := setupHandlerAccessDB(t)
|
||||
orig := dao.DB
|
||||
dao.DB = db
|
||||
t.Cleanup(func() { dao.DB = orig })
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
fake := &fakeDocumentService{}
|
||||
h := &DocumentHandler{
|
||||
documentService: fake,
|
||||
datasetService: service.NewDatasetService(),
|
||||
}
|
||||
|
||||
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `{"document_ids": ["doc-1"]}`)
|
||||
// Use a user that doesn't have access to ds-1
|
||||
c.Set("user_id", "other-user")
|
||||
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
||||
|
||||
h.StopParseDocuments(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d", w.Code)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
code, _ := resp["code"].(float64)
|
||||
if code == float64(common.CodeSuccess) {
|
||||
t.Fatal("expected error for no authorization")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user