mtfyun/ragflow
1
0 Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-11-10 21:32:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

refa: derive service identity from request context (#21923)

Browse Source
This commit is contained in:
buua436 authored and GitHub committed 2026-11-04 10:59:05 +08:00
1 parent d2353a8c36
commit 7dd0164fb3
53 files changed
+719 -668
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/handler/chunk.go
+16 -68
View File
@@ -34,14 +34,14 @@ import (
// chunkService is the consumer-side interface for ChunkHandler's service dependency.
type chunkService interface {
Get(ctx context.Context, req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error)
List(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error)
SwitchChunks(ctx context.Context, userID string, datasetID, documentID string, availableInt int, chunkIDs []string) error
UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest, userID string) error
RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest, userID string) (int64, error)
Parse(ctx context.Context, userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error)
AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error)
StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
Get(ctx context.Context, req *service.GetChunkRequest) (*service.GetChunkResponse, error)
List(ctx context.Context, req *service.ListChunksRequest) (*service.ListChunksResponse, error)
SwitchChunks(ctx context.Context, datasetID, documentID string, availableInt int, chunkIDs []string) error
UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest) error
RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest) (int64, error)
Parse(ctx context.Context, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error)
AddChunk(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error)
StopParsing(ctx context.Context, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
}
// ChunkHandler chunk handler
@@ -68,11 +68,6 @@ func NewChunkHandler(chunkService chunkService, userService *service.UserService
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/datasets/{dataset_id}/documents/{document_id}/chunks/{chunk_id} [get]
func (h *ChunkHandler) Get(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
datasetID := c.Param("dataset_id")
documentID := c.Param("document_id")
@@ -89,7 +84,7 @@ func (h *ChunkHandler) Get(c *gin.Context) {
}
ctx := c.Request.Context()
resp, err := h.chunkService.Get(ctx, req, user.ID)
resp, err := h.chunkService.Get(ctx, req)
if err != nil {
if respondPermissionErrorIf(c, err, false) {
return
@@ -103,17 +98,7 @@ func (h *ChunkHandler) Get(c *gin.Context) {
// Parse reparse the datasets' files
func (h *ChunkHandler) Parse(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
userID := strings.TrimSpace(user.ID)
if userID == "" {
common.RespondError(c, common.CodeArgumentError, "user_id is required")
return
}
datasetId := strings.TrimSpace(c.Param("dataset_id"))
if datasetId == "" {
common.Respond(c, http.StatusBadRequest, common.CodeBadRequest, nil, "dataset_id is required")
@@ -127,7 +112,7 @@ func (h *ChunkHandler) Parse(c *gin.Context) {
}
ctx := c.Request.Context()
data, code, err := h.chunkService.Parse(ctx, userID, datasetId, &req)
data, code, err := h.chunkService.Parse(ctx, datasetId, &req)
if code != common.CodeSuccess {
common.RespondError(c, code, err.Error())
return
@@ -138,11 +123,6 @@ func (h *ChunkHandler) Parse(c *gin.Context) {
// ListChunks retrieves chunks for a document from path/query parameters.
func (h *ChunkHandler) ListChunks(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
datasetID := c.Param("dataset_id")
documentID := c.Param("document_id")
@@ -183,7 +163,7 @@ func (h *ChunkHandler) ListChunks(c *gin.Context) {
}
ctx := c.Request.Context()
resp, err := h.chunkService.List(ctx, &req, user.ID)
resp, err := h.chunkService.List(ctx, &req)
if err != nil {
if respondPermissionErrorIf(c, err, false) {
return
@@ -241,11 +221,6 @@ func parseAvailableQuery(raw string) (int, bool, error) {
}
func (h *ChunkHandler) StopParsing(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
datasetID := c.Param("dataset_id")
if datasetID == "" {
@@ -264,7 +239,7 @@ func (h *ChunkHandler) StopParsing(c *gin.Context) {
}
ctx := c.Request.Context()
resp, code, err := h.chunkService.StopParsing(ctx, user.ID, datasetID, req)
resp, code, err := h.chunkService.StopParsing(ctx, datasetID, req)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -284,17 +259,6 @@ func (h *ChunkHandler) StopParsing(c *gin.Context) {
// SwitchChunks enable or disable a chunk
func (h *ChunkHandler) SwitchChunks(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
userID := strings.TrimSpace(user.ID)
if userID == "" {
common.Respond(c, http.StatusBadRequest, common.CodeAuthenticationError, nil, "user_id is required")
return
}
// Get required ID
datasetID := strings.TrimSpace(c.Param("dataset_id"))
@@ -333,7 +297,7 @@ func (h *ChunkHandler) SwitchChunks(c *gin.Context) {
}
ctx := c.Request.Context()
if err = h.chunkService.SwitchChunks(ctx, userID, datasetID, documentID, availableInt, chunkIDs); err != nil {
if err = h.chunkService.SwitchChunks(ctx, datasetID, documentID, availableInt, chunkIDs); err != nil {
if respondPermissionErrorIf(c, err, false) {
return
}
@@ -418,11 +382,6 @@ func optionalBodyString(rawBody map[string]interface{}, field, typeMessage strin
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/datasets/{dataset_id}/documents/{document_id}/chunks/{chunk_id} [patch]
func (h *ChunkHandler) UpdateChunk(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
// Decode the update fields.
body, err := io.ReadAll(c.Request.Body)
@@ -535,7 +494,7 @@ func (h *ChunkHandler) UpdateChunk(c *gin.Context) {
req.ChunkID = chunkID
ctx := c.Request.Context()
err = h.chunkService.UpdateChunk(ctx, &req, user.ID)
err = h.chunkService.UpdateChunk(ctx, &req)
if err != nil {
if respondPermissionErrorIf(c, err, false) {
return
@@ -566,11 +525,6 @@ func (h *ChunkHandler) UpdateChunk(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/datasets/{dataset_id}/documents/{document_id}/chunks [delete]
func (h *ChunkHandler) RemoveChunks(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
datasetID := c.Param("dataset_id")
docID := c.Param("document_id")
@@ -589,7 +543,7 @@ func (h *ChunkHandler) RemoveChunks(c *gin.Context) {
req.DocID = docID
ctx := c.Request.Context()
deletedCount, err := h.chunkService.RemoveChunks(ctx, &req, user.ID)
deletedCount, err := h.chunkService.RemoveChunks(ctx, &req)
if err != nil {
if respondPermissionErrorIf(c, err, false) {
return
@@ -658,13 +612,7 @@ func addChunkResponseMessage(code common.ErrorCode, err error) string {
func (h *ChunkHandler) AddChunk(c *gin.Context) {
ctx := c.Request.Context()
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
userID := user.ID
datasetID, documentID := strings.TrimSpace(c.Param("dataset_id")), strings.TrimSpace(c.Param("document_id"))
var rawBody map[string]json.RawMessage
@@ -716,7 +664,7 @@ func (h *ChunkHandler) AddChunk(c *gin.Context) {
ImageBase64: imageBase64,
}
resp, err := h.chunkService.AddChunk(ctx, &req, userID)
resp, err := h.chunkService.AddChunk(ctx, &req)
if err != nil {
var codedErr service.ErrorCoder
if errors.As(err, &codedErr) {
internal/handler/chunk_sdk_contract_test.go
+2 -2
View File
@@ -28,7 +28,7 @@ import (
func TestChunkUpdateRejectsMalformedSDKFields(t *testing.T) {
for _, body := range []string{`{"content":null}`, `{"content":1}`, `{"important_keywords":[1]}`, `{"important_keywords":null}`, `{"questions":"abc"}`, `{"questions":[false]}`, `{"available":1}`, `{"available":"true"}`} {
t.Run(body, func(t *testing.T) {
mock := &mockChunkSvc{updateChunkFn: func(context.Context, *service.UpdateChunkRequest, string) error {
mock := &mockChunkSvc{updateChunkFn: func(context.Context, *service.UpdateChunkRequest) error {
t.Error("invalid update reached service")
return nil
}}
@@ -50,7 +50,7 @@ func TestChunkUpdateRejectsMalformedSDKFields(t *testing.T) {
}
func TestChunkListAcceptsSDKID(t *testing.T) {
mock := &mockChunkSvc{listFn: func(_ context.Context, req *service.ListChunksRequest, _ string) (*service.ListChunksResponse, error) {
mock := &mockChunkSvc{listFn: func(_ context.Context, req *service.ListChunksRequest) (*service.ListChunksResponse, error) {
if len(req.ChunkIDs) != 1 || req.ChunkIDs[0] != "sdk-id" {
t.Errorf("filter=%v", req.ChunkIDs)
}
internal/handler/chunk_test.go
+62 -63
View File
@@ -11,7 +11,6 @@ import (
"testing"
"ragflow/internal/common"
"ragflow/internal/entity"
"ragflow/internal/permission"
"ragflow/internal/service"
@@ -21,12 +20,12 @@ import (
// mockChunkSvc implements chunkSvcIface for testing ChunkHandler.
// Only the methods actually called by the test are set; others panic.
type mockChunkSvc struct {
addChunkFn func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error)
listFn func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error)
switchChunksFn func(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error
updateChunkFn func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error
stopParsingFn func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
removeChunksFn func(context.Context, *service.RemoveChunksRequest, string) (int64, error)
addChunkFn func(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error)
listFn func(ctx context.Context, req *service.ListChunksRequest) (*service.ListChunksResponse, error)
switchChunksFn func(ctx context.Context, datasetID, documentID string, availableInt int, chunkIDs []string) error
updateChunkFn func(ctx context.Context, req *service.UpdateChunkRequest) error
stopParsingFn func(ctx context.Context, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
removeChunksFn func(context.Context, *service.RemoveChunksRequest) (int64, error)
}
type codedTestError struct {
@@ -42,30 +41,30 @@ func (e codedTestError) Code() common.ErrorCode {
return e.code
}
func (m *mockChunkSvc) Get(context.Context, *service.GetChunkRequest, string) (*service.GetChunkResponse, error) {
func (m *mockChunkSvc) Get(context.Context, *service.GetChunkRequest) (*service.GetChunkResponse, error) {
panic("not implemented")
}
func (m *mockChunkSvc) List(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
func (m *mockChunkSvc) List(ctx context.Context, req *service.ListChunksRequest) (*service.ListChunksResponse, error) {
if m.listFn != nil {
return m.listFn(ctx, req, userID)
return m.listFn(ctx, req)
}
panic("not implemented")
}
func (m *mockChunkSvc) SwitchChunks(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error {
func (m *mockChunkSvc) SwitchChunks(ctx context.Context, datasetID, documentID string, availableInt int, chunkIDs []string) error {
if m.switchChunksFn != nil {
return m.switchChunksFn(ctx, userID, datasetID, documentID, availableInt, chunkIDs)
return m.switchChunksFn(ctx, datasetID, documentID, availableInt, chunkIDs)
}
panic("not implemented")
}
func (m *mockChunkSvc) UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
func (m *mockChunkSvc) UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest) error {
if m.updateChunkFn != nil {
return m.updateChunkFn(ctx, req, userID)
return m.updateChunkFn(ctx, req)
}
panic("not implemented")
}
func (m *mockChunkSvc) RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest, userID string) (int64, error) {
func (m *mockChunkSvc) RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest) (int64, error) {
if m.removeChunksFn != nil {
return m.removeChunksFn(ctx, req, userID)
return m.removeChunksFn(ctx, req)
}
panic("not implemented")
}
@@ -83,10 +82,10 @@ func TestChunkRemoveUsesRouteIDs(t *testing.T) {
} {
t.Run(tc.name, func(t *testing.T) {
called := false
mock := &mockChunkSvc{removeChunksFn: func(_ context.Context, req *service.RemoveChunksRequest, userID string) (int64, error) {
mock := &mockChunkSvc{removeChunksFn: func(ctx context.Context, req *service.RemoveChunksRequest) (int64, error) {
called = true
if req.DatasetID != "kb-1" || req.DocID != "doc-1" || userID != "user-1" {
t.Fatalf("unexpected delete target: %#v, user %q", req, userID)
if req.DatasetID != "kb-1" || req.DocID != "doc-1" || permission.UserIDFromContext(ctx) != "user-1" {
t.Fatalf("unexpected delete target: %#v, user %q", req, permission.UserIDFromContext(ctx))
}
if !reflect.DeepEqual(req.ChunkIDs, []string{"chunk-1"}) {
t.Fatalf("unexpected chunk IDs: %v", req.ChunkIDs)
@@ -122,18 +121,18 @@ func TestChunkRemoveUsesRouteIDs(t *testing.T) {
})
}
}
func (m *mockChunkSvc) StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
func (m *mockChunkSvc) StopParsing(ctx context.Context, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
if m.stopParsingFn != nil {
return m.stopParsingFn(ctx, userID, datasetID, req)
return m.stopParsingFn(ctx, datasetID, req)
}
panic("not implemented")
}
func (m *mockChunkSvc) Parse(context.Context, string, string, *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) {
func (m *mockChunkSvc) Parse(context.Context, string, *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) {
panic("not implemented")
}
func (m *mockChunkSvc) AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
func (m *mockChunkSvc) AddChunk(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error) {
if m.addChunkFn != nil {
return m.addChunkFn(ctx, req, userID)
return m.addChunkFn(ctx, req)
}
return &service.AddChunkResponse{Chunk: map[string]interface{}{"id": "chunk-1"}}, nil
}
@@ -144,7 +143,7 @@ func setupChunkStopParsingTest(userID string) (*gin.Engine, *mockChunkSvc) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user", &entity.User{ID: userID})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), userID, ""))
})
r.DELETE("/api/v1/datasets/:dataset_id/chunks", h.StopParsing)
return r, mock
@@ -155,7 +154,7 @@ func setupChunkHandlerWithUser(userID string, mock *mockChunkSvc) (*gin.Engine,
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user", &entity.User{ID: userID})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), userID, ""))
})
return r, h
}
@@ -165,9 +164,9 @@ func TestChunkHandlerListChunksMapsPathAndQuery(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.GET("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.ListChunks)
mock.listFn = func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
if userID != "user-1" {
t.Fatalf("userID = %q, want user-1", userID)
mock.listFn = func(ctx context.Context, req *service.ListChunksRequest) (*service.ListChunksResponse, error) {
if permission.UserIDFromContext(ctx) != "user-1" {
t.Fatalf("context userID = %q, want user-1", permission.UserIDFromContext(ctx))
}
if req.DatasetID != "kb-1" || req.DocID != "doc-1" {
t.Fatalf("req ids = %q/%q, want kb-1/doc-1", req.DatasetID, req.DocID)
@@ -217,9 +216,9 @@ func TestChunkHandlerListChunksMapsAvailableFalse(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.GET("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.ListChunks)
mock.listFn = func(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
if userID != "user-1" {
t.Fatalf("userID = %q, want user-1", userID)
mock.listFn = func(ctx context.Context, req *service.ListChunksRequest) (*service.ListChunksResponse, error) {
if permission.UserIDFromContext(ctx) != "user-1" {
t.Fatalf("context userID = %q, want user-1", permission.UserIDFromContext(ctx))
}
if req.AvailableInt == nil || *req.AvailableInt != 0 {
t.Fatalf("available_int = %v, want 0", req.AvailableInt)
@@ -245,9 +244,9 @@ func TestChunkHandlerSwitchChunksCallsService(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.SwitchChunks)
mock.switchChunksFn = func(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error {
if userID != "user-1" || datasetID != "kb-1" || documentID != "doc-1" {
t.Fatalf("ids = %q/%q/%q, want user-1/kb-1/doc-1", userID, datasetID, documentID)
mock.switchChunksFn = func(ctx context.Context, datasetID, documentID string, availableInt int, chunkIDs []string) error {
if permission.UserIDFromContext(ctx) != "user-1" || datasetID != "kb-1" || documentID != "doc-1" {
t.Fatalf("identity and ids = %q/%q/%q, want user-1/kb-1/doc-1", permission.UserIDFromContext(ctx), datasetID, documentID)
}
if availableInt != 0 {
t.Fatalf("availableInt = %d, want 0", availableInt)
@@ -299,9 +298,9 @@ func TestChunkHandlerUpdateChunkUsesPathIDs(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
if userID != "user-1" {
t.Fatalf("userID = %q, want user-1", userID)
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest) error {
if permission.UserIDFromContext(ctx) != "user-1" {
t.Fatalf("context userID = %q, want user-1", permission.UserIDFromContext(ctx))
}
if req.DatasetID != "kb-1" || req.DocumentID != "doc-1" || req.ChunkID != "chunk-1" {
t.Fatalf("ids = %q/%q/%q, want kb-1/doc-1/chunk-1", req.DatasetID, req.DocumentID, req.ChunkID)
@@ -327,7 +326,7 @@ func TestChunkHandlerUpdateChunkValidationErrorIsBadRequest(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest) error {
return codedTestError{code: common.CodeArgumentError, msg: "`tag_feas` values must be finite numbers greater than 0"}
}
@@ -346,7 +345,7 @@ func TestChunkHandlerUpdateChunkForwardsImageFields(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest) error {
if req.ImageBase64 == nil || *req.ImageBase64 != "aGVsbG8=" {
t.Fatalf("image_base64 = %v, want aGVsbG8=", req.ImageBase64)
}
@@ -383,7 +382,7 @@ func TestChunkHandlerUpdateChunkRejectsNonStringImageFields(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
mock.updateChunkFn = func(context.Context, *service.UpdateChunkRequest, string) error {
mock.updateChunkFn = func(context.Context, *service.UpdateChunkRequest) error {
t.Fatal("service must not be called for a malformed image field")
return nil
}
@@ -429,7 +428,7 @@ func TestChunkHandlerUpdateChunkForwardsTagKwd(t *testing.T) {
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
called := false
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
mock.updateChunkFn = func(ctx context.Context, req *service.UpdateChunkRequest) error {
called = true
if !reflect.DeepEqual(req.TagKwd, tc.want) {
t.Fatalf("tag_kwd = %#v, want %#v", req.TagKwd, tc.want)
@@ -468,7 +467,7 @@ func TestChunkHandlerUpdateChunkRejectsInvalidTagKwd(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
mock.updateChunkFn = func(context.Context, *service.UpdateChunkRequest, string) error {
mock.updateChunkFn = func(context.Context, *service.UpdateChunkRequest) error {
t.Fatal("service must not be called for a malformed tag_kwd")
return nil
}
@@ -500,7 +499,7 @@ func TestChunkHandlerUpdateChunkStillRejectsUnknownFields(t *testing.T) {
r, h := setupChunkHandlerWithUser("user-1", mock)
r.PATCH("/api/v1/datasets/:dataset_id/documents/:document_id/chunks/:chunk_id", h.UpdateChunk)
mock.updateChunkFn = func(context.Context, *service.UpdateChunkRequest, string) error {
mock.updateChunkFn = func(context.Context, *service.UpdateChunkRequest) error {
t.Fatal("service must not be called for an unsupported field")
return nil
}
@@ -527,9 +526,9 @@ func TestChunkHandlerUpdateChunkStillRejectsUnknownFields(t *testing.T) {
func TestChunkStopParsing_Success(t *testing.T) {
r, mock := setupChunkStopParsingTest("user1")
mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
if userID != "user1" {
t.Fatalf("expected user1, got %q", userID)
mock.stopParsingFn = func(ctx context.Context, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
if permission.UserIDFromContext(ctx) != "user1" {
t.Fatalf("expected user1, got %q", permission.UserIDFromContext(ctx))
}
if datasetID != "kb1" {
t.Fatalf("expected kb1, got %q", datasetID)
@@ -562,7 +561,7 @@ func TestChunkStopParsing_Success(t *testing.T) {
func TestChunkStopParsingRouteRequiresDocumentIDs(t *testing.T) {
r, mock := setupChunkStopParsingTest("user1")
mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
mock.stopParsingFn = func(ctx context.Context, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
t.Fatal("service should not be called when document_ids is missing")
return nil, common.CodeSuccess, nil
}
@@ -589,7 +588,7 @@ func TestChunkStopParsingRouteRequiresDocumentIDs(t *testing.T) {
func TestChunkStopParsing_InvalidStateOmitsData(t *testing.T) {
r, mock := setupChunkStopParsingTest("user1")
mock.stopParsingFn = func(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
mock.stopParsingFn = func(ctx context.Context, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
return &service.StopParsingResponse{
Data: map[string]interface{}{"error_code": "DOC_STOP_PARSING_INVALID_STATE"},
}, common.CodeDataError, errors.New("can't stop parsing document that has not started or already completed")
@@ -628,9 +627,9 @@ func (e addChunkTestError) Code() common.ErrorCode { return e.code }
func TestChunkHandlerAddChunkSuccess(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
if userID != "user1" {
t.Fatalf("userID = %q, want user1", userID)
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error) {
if permission.UserIDFromContext(ctx) != "user1" {
t.Fatalf("context userID = %q, want user1", permission.UserIDFromContext(ctx))
}
if req.DatasetID != "kb1" || req.DocumentID != "doc1" || req.Content != "chunk body" {
t.Fatalf("unexpected request: %#v", req)
@@ -642,7 +641,7 @@ func TestChunkHandlerAddChunkSuccess(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user", &entity.User{ID: "user1"})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), "user1", ""))
})
h := &ChunkHandler{chunkService: mock}
r.POST("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.AddChunk)
@@ -666,7 +665,7 @@ func TestChunkHandlerAddChunkSuccess(t *testing.T) {
func TestChunkHandlerAddChunkPathIDsOverrideBody(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error) {
if req.DatasetID != "kb1" || req.DocumentID != "doc1" {
t.Fatalf("path IDs were not preserved: %#v", req)
}
@@ -680,7 +679,7 @@ func TestChunkHandlerAddChunkPathIDsOverrideBody(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user", &entity.User{ID: "user1"})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), "user1", ""))
})
h := &ChunkHandler{chunkService: mock}
r.POST("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.AddChunk)
@@ -697,7 +696,7 @@ func TestChunkHandlerAddChunkPathIDsOverrideBody(t *testing.T) {
func TestChunkHandlerAddChunkCodedError(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error) {
return nil, addChunkTestError{code: common.CodeDataError, msg: "`content` is required"}
},
}
@@ -705,7 +704,7 @@ func TestChunkHandlerAddChunkCodedError(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user", &entity.User{ID: "user1"})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), "user1", ""))
})
h := &ChunkHandler{chunkService: mock}
r.POST("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.AddChunk)
@@ -758,7 +757,7 @@ func TestChunkHandlerAddChunkValidatesListFields(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error) {
t.Fatal("service should not be called for invalid request")
return nil, nil
},
@@ -766,7 +765,7 @@ func TestChunkHandlerAddChunkValidatesListFields(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user", &entity.User{ID: "user1"})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), "user1", ""))
})
h := &ChunkHandler{chunkService: mock}
r.POST("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.AddChunk)
@@ -819,7 +818,7 @@ func TestChunkHandlerAddChunkForwardsTagKwd(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
var got *service.AddChunkRequest
mock := &mockChunkSvc{
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error) {
got = req
return &service.AddChunkResponse{Chunk: map[string]interface{}{"id": "chunk-1"}}, nil
},
@@ -827,7 +826,7 @@ func TestChunkHandlerAddChunkForwardsTagKwd(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user", &entity.User{ID: "user1"})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), "user1", ""))
})
h := &ChunkHandler{chunkService: mock}
r.POST("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.AddChunk)
@@ -858,7 +857,7 @@ func TestChunkHandlerAddChunkForwardsTagKwd(t *testing.T) {
func TestChunkHandlerAddChunkHidesServerErrorDetails(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error) {
return nil, addChunkTestError{code: common.CodeServerError, msg: "encode chunk embedding: provider secret"}
},
}
@@ -866,7 +865,7 @@ func TestChunkHandlerAddChunkHidesServerErrorDetails(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user", &entity.User{ID: "user1"})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), "user1", ""))
})
h := &ChunkHandler{chunkService: mock}
r.POST("/api/v1/datasets/:dataset_id/documents/:document_id/chunks", h.AddChunk)
internal/handler/dataset.go
+36 -50
View File
@@ -48,11 +48,11 @@ type DatasetsHandler struct {
}
type searchDatasetsService interface {
SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error)
SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest) (*service.SearchDatasetsResponse, error)
}
type searchDatasetService interface {
SearchDataset(ctx context.Context, datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error)
SearchDataset(ctx context.Context, datasetID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error)
}
// NewDatasetsHandler creates a new datasets' handler.
@@ -70,7 +70,7 @@ func NewDatasetsHandler(datasetsService *dataset.DatasetService, metadataService
// ListDatasets handles GET /api/v1/datasets.
func (h *DatasetsHandler) ListDatasets(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -78,7 +78,7 @@ func (h *DatasetsHandler) ListDatasets(c *gin.Context) {
if c.Query("type") == "filter" {
ctx := c.Request.Context()
data, code, err := h.datasetsService.ListDatasetFilters(ctx, user.ID, c.Query("canvas_id"))
data, code, err := h.datasetsService.ListDatasetFilters(ctx, c.Query("canvas_id"))
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -216,7 +216,6 @@ func (h *DatasetsHandler) ListDatasets(c *gin.Context) {
keywords,
ownerIDs,
parserID,
user.ID,
ids,
c.Query("canvas_id"),
)
@@ -341,7 +340,7 @@ func (h *DatasetsHandler) CreateDataset(c *gin.Context) {
// GetDataset handles GET /api/v1/datasets/:dataset_id.
func (h *DatasetsHandler) GetDataset(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -350,7 +349,7 @@ func (h *DatasetsHandler) GetDataset(c *gin.Context) {
ctx := c.Request.Context()
datasetID := c.Param("dataset_id")
result, code, err := h.datasetsService.GetDataset(ctx, datasetID, user.ID)
result, code, err := h.datasetsService.GetDataset(ctx, datasetID)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -547,7 +546,7 @@ func (h *DatasetsHandler) UpdateDataset(c *gin.Context) {
ctx := c.Request.Context()
result, code, err := h.datasetsService.UpdateDataset(ctx, datasetID, userID, req)
result, code, err := h.datasetsService.UpdateDataset(ctx, datasetID, req)
if err != nil {
common.RespondError(c, code, err.Error())
operationLog.Code = int(code)
@@ -565,7 +564,7 @@ func (h *DatasetsHandler) UpdateDataset(c *gin.Context) {
}
func (h *DatasetsHandler) GetMetadataConfig(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -574,7 +573,7 @@ func (h *DatasetsHandler) GetMetadataConfig(c *gin.Context) {
ctx := c.Request.Context()
datasetID := c.Param("dataset_id")
result, code, err := h.datasetsService.GetMetadataConfig(ctx, datasetID, user.ID)
result, code, err := h.datasetsService.GetMetadataConfig(ctx, datasetID)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -584,7 +583,7 @@ func (h *DatasetsHandler) GetMetadataConfig(c *gin.Context) {
}
func (h *DatasetsHandler) UpdateMetadataConfig(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -599,7 +598,7 @@ func (h *DatasetsHandler) UpdateMetadataConfig(c *gin.Context) {
ctx := c.Request.Context()
result, code, err := h.datasetsService.UpdateMetadataConfig(ctx, datasetID, user.ID, &req)
result, code, err := h.datasetsService.UpdateMetadataConfig(ctx, datasetID, &req)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -610,14 +609,14 @@ func (h *DatasetsHandler) UpdateMetadataConfig(c *gin.Context) {
// GetIngestionSummary handles GET /api/v1/datasets/:dataset_id/ingestions/summary.
func (h *DatasetsHandler) GetIngestionSummary(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
ctx := c.Request.Context()
datasetID := c.Param("dataset_id")
result, code, err := h.datasetsService.GetIngestionSummary(ctx, datasetID, user.ID)
result, code, err := h.datasetsService.GetIngestionSummary(ctx, datasetID)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -629,7 +628,7 @@ func (h *DatasetsHandler) GetIngestionSummary(c *gin.Context) {
// ListIngestionLogs handles GET /api/v1/datasets/:dataset_id/ingestions.
// The log_type=datasource query returns logs for connectors linked to the dataset.
func (h *DatasetsHandler) ListIngestionLogs(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -659,7 +658,7 @@ func (h *DatasetsHandler) ListIngestionLogs(c *gin.Context) {
ctx := c.Request.Context()
if strings.EqualFold(strings.TrimSpace(c.Query("log_type")), "datasource") {
result, code, err := h.datasetsService.ListSyncLogs(ctx, datasetID, user.ID, page, pageSize)
result, code, err := h.datasetsService.ListSyncLogs(ctx, datasetID, page, pageSize)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -681,7 +680,7 @@ func (h *DatasetsHandler) ListIngestionLogs(c *gin.Context) {
// equivalent; the frontend only sends it on the Go backend.
documentID := c.Query("document_id")
result, code, err := h.datasetsService.ListIngestionLogs(ctx, datasetID, user.ID, page, pageSize, terms, operationStatus, createDateFrom, createDateTo, logType, keywords, documentID)
result, code, err := h.datasetsService.ListIngestionLogs(ctx, datasetID, page, pageSize, terms, operationStatus, createDateFrom, createDateTo, logType, keywords, documentID)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -692,7 +691,7 @@ func (h *DatasetsHandler) ListIngestionLogs(c *gin.Context) {
// GetIngestionLog handles GET /api/v1/datasets/:dataset_id/ingestions/:log_id.
func (h *DatasetsHandler) GetIngestionLog(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -702,7 +701,7 @@ func (h *DatasetsHandler) GetIngestionLog(c *gin.Context) {
datasetID := c.Param("dataset_id")
logID := c.Param("log_id")
result, code, err := h.datasetsService.GetIngestionLog(ctx, datasetID, user.ID, logID)
result, code, err := h.datasetsService.GetIngestionLog(ctx, datasetID, logID)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -714,7 +713,7 @@ func (h *DatasetsHandler) GetIngestionLog(c *gin.Context) {
// ListIngestionMessages handles GET
// /api/v1/datasets/:dataset_id/ingestions/:log_id/messages.
func (h *DatasetsHandler) ListIngestionMessages(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -739,7 +738,7 @@ func (h *DatasetsHandler) ListIngestionMessages(c *gin.Context) {
}
result, code, err := h.datasetsService.ListIngestionMessages(
c.Request.Context(), c.Param("dataset_id"), user.ID, c.Param("log_id"), limit, afterID, beforeID,
c.Request.Context(), c.Param("dataset_id"), c.Param("log_id"), limit, afterID, beforeID,
)
if err != nil {
common.RespondError(c, code, err.Error())
@@ -852,7 +851,7 @@ func (h *DatasetsHandler) DeleteDatasets(c *gin.Context) {
// GetKnowledgeGraph handles GET /api/v1/datasets/:dataset_id/graph.
func (h *DatasetsHandler) GetKnowledgeGraph(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -865,7 +864,7 @@ func (h *DatasetsHandler) GetKnowledgeGraph(c *gin.Context) {
}
ctx := c.Request.Context()
datasetInstance, code, err := h.datasetsService.GetDataset(ctx, datasetID, user.ID)
datasetInstance, code, err := h.datasetsService.GetDataset(ctx, datasetID)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -955,7 +954,7 @@ func (h *DatasetsHandler) GetKnowledgeGraph(c *gin.Context) {
// DeleteKnowledgeGraph handles DELETE /api/v1/datasets/:dataset_id/graph.
func (h *DatasetsHandler) DeleteKnowledgeGraph(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -968,7 +967,7 @@ func (h *DatasetsHandler) DeleteKnowledgeGraph(c *gin.Context) {
}
ctx := c.Request.Context()
datasetInstance, code, err := h.datasetsService.GetDataset(ctx, datasetID, user.ID)
datasetInstance, code, err := h.datasetsService.GetDataset(ctx, datasetID)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -1001,7 +1000,7 @@ func (h *DatasetsHandler) DeleteKnowledgeGraph(c *gin.Context) {
// CheckEmbedding Check embedding model compatibility by sampling random chunks,
// re-embedding them with the new model, and computing cosine similarity.
func (h *DatasetsHandler) CheckEmbedding(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -1013,12 +1012,6 @@ func (h *DatasetsHandler) CheckEmbedding(c *gin.Context) {
return
}
userID := strings.TrimSpace(user.ID)
if userID == "" {
common.RespondError(c, common.CodeDataError, "user_id is required")
return
}
var req service.CheckEmbeddingRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.RespondError(c, common.CodeDataError, err.Error())
@@ -1030,7 +1023,7 @@ func (h *DatasetsHandler) CheckEmbedding(c *gin.Context) {
}
ctx := c.Request.Context()
data, code, err := h.datasetsService.CheckEmbedding(ctx, userID, datasetID, &req)
data, code, err := h.datasetsService.CheckEmbedding(ctx, datasetID, &req)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -1048,7 +1041,7 @@ func (h *DatasetsHandler) CheckEmbedding(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/datasets/tags/aggregation [get]
func (h *DatasetsHandler) AggregateTags(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -1070,7 +1063,7 @@ func (h *DatasetsHandler) AggregateTags(c *gin.Context) {
ctx := c.Request.Context()
result, code, err := h.datasetsService.AggregateTags(ctx, datasetIDs, user.ID)
result, code, err := h.datasetsService.AggregateTags(ctx, datasetIDs)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -1082,7 +1075,7 @@ func (h *DatasetsHandler) AggregateTags(c *gin.Context) {
// state. The optional `kind` query parameter scopes the status to one compile
// type.
func (h *DatasetsHandler) GetCompilationStatus(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -1092,10 +1085,9 @@ func (h *DatasetsHandler) GetCompilationStatus(c *gin.Context) {
common.RespondError(c, common.CodeDataError, "dataset_id is required")
return
}
userID := strings.TrimSpace(user.ID)
ctx := c.Request.Context()
kind := strings.TrimSpace(c.Query("kind"))
st, code, err := h.datasetsService.GetDatasetCompilationStatus(ctx, userID, datasetID, kind)
st, code, err := h.datasetsService.GetDatasetCompilationStatus(ctx, datasetID, kind)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -1152,7 +1144,7 @@ func (h *DatasetsHandler) ListMetadataFlattened(c *gin.Context) {
}
func (h *DatasetsHandler) UpdateDocumentMetadataConfig(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -1168,12 +1160,6 @@ func (h *DatasetsHandler) UpdateDocumentMetadataConfig(c *gin.Context) {
common.RespondError(c, common.CodeArgumentError, "document_id is required")
return
}
userID := strings.TrimSpace(user.ID)
if userID == "" {
common.RespondError(c, common.CodeArgumentError, "user_id is required")
return
}
var req map[string]interface{}
if err := c.ShouldBindJSON(&req); err != nil {
common.RespondError(c, common.CodeDataError, err.Error())
@@ -1181,7 +1167,7 @@ func (h *DatasetsHandler) UpdateDocumentMetadataConfig(c *gin.Context) {
}
ctx := c.Request.Context()
data, code, err := h.datasetsService.UpdateDocumentMetadataConfig(ctx, userID, datasetID, documentID, req)
data, code, err := h.datasetsService.UpdateDocumentMetadataConfig(ctx, datasetID, documentID, req)
if err != nil {
common.RespondError(c, code, err.Error())
return
@@ -1198,7 +1184,7 @@ func (h *DatasetsHandler) UpdateDocumentMetadataConfig(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/datasets/search [post]
func (h *DatasetsHandler) SearchDatasets(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -1261,7 +1247,7 @@ func (h *DatasetsHandler) SearchDatasets(c *gin.Context) {
ctx := c.Request.Context()
resp, err := searchService.SearchDatasets(ctx, &req, user.ID)
resp, err := searchService.SearchDatasets(ctx, &req)
if err != nil {
if respondPermissionErrorIf(c, err, false) {
return
@@ -1282,7 +1268,7 @@ func (h *DatasetsHandler) SearchDatasets(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/datasets/{dataset_id}/search [post]
func (h *DatasetsHandler) SearchDataset(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -1320,7 +1306,7 @@ func (h *DatasetsHandler) SearchDataset(c *gin.Context) {
ctx := c.Request.Context()
resp, err := searchService.SearchDataset(ctx, datasetID, user.ID, &req)
resp, err := searchService.SearchDataset(ctx, datasetID, &req)
if err != nil {
if respondPermissionErrorIf(c, err, false) {
return
internal/handler/dataset_search_test.go
+6 -4
View File
@@ -27,9 +27,9 @@ type fakeSearchDatasetService struct {
err error
}
func (f *fakeSearchDatasetService) SearchDataset(ctx context.Context, datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) {
func (f *fakeSearchDatasetService) SearchDataset(ctx context.Context, datasetID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) {
f.datasetID = datasetID
f.userID = userID
f.userID = permission.UserIDFromContext(ctx)
f.req = req
return f.resp, f.err
}
@@ -41,8 +41,8 @@ type fakeSearchDatasetsService struct {
err error
}
func (f *fakeSearchDatasetsService) SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) {
f.userID = userID
func (f *fakeSearchDatasetsService) SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest) (*service.SearchDatasetsResponse, error) {
f.userID = permission.UserIDFromContext(ctx)
f.req = req
return f.resp, f.err
}
@@ -55,6 +55,7 @@ func TestDatasetsHandlerSearchDataset(t *testing.T) {
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/ds-1/search", strings.NewReader(`{"question":"hello","document_ids":["doc-1"],"page_size":9,"metadata_condition":{"logic":"and","conditions":[{"name":"author","comparison_operator":"=","value":"Luo"}]},"knn_top_k":7,"knn_num_candidates":14,"keywords_similarity_weight":0.7}`))
req.Header.Set("Content-Type", "application/json")
req = req.WithContext(permission.WithIdentity(req.Context(), "user-1", ""))
c, _ := gin.CreateTestContext(rec)
c.Request = req
c.Set("user", &entity.User{ID: "user-1"})
@@ -171,6 +172,7 @@ func TestDatasetsHandlerSearchDatasetsSuccess(t *testing.T) {
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/datasets/search", strings.NewReader(`{"question":" hello ","dataset_ids":["ds-1"],"document_ids":["doc-1"],"page_size":9,"metadata_condition":{"logic":"and","conditions":[]},"top_k":7,"include_knowledge_compilation":false}`))
req.Header.Set("Content-Type", "application/json")
req = req.WithContext(permission.WithIdentity(req.Context(), "user-1", ""))
c, _ := gin.CreateTestContext(rec)
c.Request = req
c.Set("user", &entity.User{ID: "user-1"})
internal/handler/document.go
+37 -45
View File
@@ -52,8 +52,8 @@ type documentServiceIface interface {
GetDocumentByID(ctx context.Context, id string) (*document.DocumentResponse, error)
UpdateDocument(ctx context.Context, id string, req *document.UpdateDocumentRequest) (common.ErrorCode, error)
DeleteDocument(ctx context.Context, id string) error
DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error)
ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID string) (int, error)
ParseDocuments(ctx context.Context, datasetID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
StopParseDocuments(ctx context.Context, datasetID string, docIDs []string) (map[string]interface{}, error)
StopParseDocumentsAll(ctx context.Context, datasetID string, runStatuses []string) (map[string]interface{}, error)
ListDocuments(ctx context.Context, page, pageSize int) ([]*document.DocumentResponse, int64, error)
@@ -63,27 +63,27 @@ type documentServiceIface interface {
GetDocumentFiltersByDatasetID(ctx context.Context, opts dao.DocumentListOptions) (map[string]interface{}, int64, error)
GetMetadataByKBs(ctx context.Context, kbIDs []string) (map[string]interface{}, error)
GetDocumentsByAuthorID(ctx context.Context, authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error)
GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error)
GetDocumentImage(ctx context.Context, userID, imageID string) ([]byte, error)
GetDocumentImageForDocument(ctx context.Context, userID, docID, imageID string) ([]byte, error)
GetThumbnails(ctx context.Context, docIDs []string) (map[string]string, error)
GetDocumentImage(ctx context.Context, imageID string) ([]byte, error)
GetDocumentImageForDocument(ctx context.Context, docID, imageID string) ([]byte, error)
GetDocumentThumbnail(ctx context.Context, docID string) ([]byte, error)
GetMetadataSummary(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error)
GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]interface{}, error)
GetDocumentArtifact(ctx context.Context, filename, userID string) (*document.ArtifactResponse, error)
GetDocumentArtifact(ctx context.Context, filename string) (*document.ArtifactResponse, error)
GetDocumentPreview(ctx context.Context, docID string) (*document.DocumentPreview, error)
UploadLocalDocuments(ctx context.Context, kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string)
UploadWebDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error)
UploadEmptyDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error)
DownloadDocument(ctx context.Context, datasetID, docID string) (*document.DownloadDocumentResp, error)
UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error)
ApplyDatasetParserConfig(ctx context.Context, userID, datasetID string) (*document.ApplyDatasetParserConfigResponse, common.ErrorCode, error)
UpdateDatasetDocument(ctx context.Context, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error)
ApplyDatasetParserConfig(ctx context.Context, datasetID string) (*document.ApplyDatasetParserConfigResponse, common.ErrorCode, error)
BatchUpdateDocumentMetadatas(ctx context.Context, datasetID string, selector *document.MetadataSelector, updates []document.MetadataUpdate, deletes []document.MetadataDelete) (*document.BatchUpdateMetadatasResponse, common.ErrorCode, error)
ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error)
IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error)
Ingest(ctx context.Context, userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error)
RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error)
BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, DocumentIDs []string) (map[string]interface{}, common.ErrorCode, error)
ListIngestionTasks(ctx context.Context, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error)
IngestDocuments(ctx context.Context, datasetID string, docIDs []string) ([]*service.ParseDocumentResponse, error)
StopIngestionTasks(ctx context.Context, tasks []string) ([]*entity.IngestionTask, error)
Ingest(ctx context.Context, req *document.IngestDocumentRequest) (common.ErrorCode, error)
RemoveIngestionTasks(ctx context.Context, tasks []string) ([]map[string]string, error)
BatchUpdateDocumentStatus(ctx context.Context, datasetID, status string, DocumentIDs []string) (map[string]interface{}, common.ErrorCode, error)
HasActiveIngestionTasks(ctx context.Context, datasetID string) (bool, error)
}
@@ -145,7 +145,7 @@ func (h *DocumentHandler) GetDocumentByID(c *gin.Context) {
// GetThumbnail Get thumbnails for documents.
func (h *DocumentHandler) GetThumbnail(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
@@ -158,7 +158,7 @@ func (h *DocumentHandler) GetThumbnail(c *gin.Context) {
}
ctx := c.Request.Context()
result, err := h.documentService.GetThumbnails(ctx, user.ID, docIDs)
result, err := h.documentService.GetThumbnails(ctx, docIDs)
if err != nil {
common.RespondError(c, common.CodeServerError, err.Error())
return
@@ -189,14 +189,14 @@ func parseThumbnailDocIDs(c *gin.Context) []string {
// GetDocumentImage returns a document image from object storage.
func (h *DocumentHandler) GetDocumentImage(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
imageID := c.Param("image_id")
ctx := c.Request.Context()
data, err := h.documentService.GetDocumentImage(ctx, user.ID, imageID)
data, err := h.documentService.GetDocumentImage(ctx, imageID)
if err != nil {
common.RespondError(c, common.CodeDataError, "Image not found.")
return
@@ -205,12 +205,12 @@ func (h *DocumentHandler) GetDocumentImage(c *gin.Context) {
}
func (h *DocumentHandler) GetDocumentImageForDocument(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
data, err := h.documentService.GetDocumentImageForDocument(c.Request.Context(), user.ID, c.Param("id"), c.Param("image_id"))
data, err := h.documentService.GetDocumentImageForDocument(c.Request.Context(), c.Param("id"), c.Param("image_id"))
if err != nil {
common.RespondError(c, common.CodeDataError, "Image not found.")
return
@@ -255,14 +255,14 @@ func documentImageContentType(data []byte) string {
}
func (h *DocumentHandler) GetDocumentArtifact(c *gin.Context) {
user, code, msg := GetUser(c)
_, code, msg := GetUser(c)
if code != common.CodeSuccess {
common.RespondError(c, code, msg)
return
}
filename := c.Param("filename")
ctx := c.Request.Context()
artifact, err := h.documentService.GetDocumentArtifact(ctx, filename, user.ID)
artifact, err := h.documentService.GetDocumentArtifact(ctx, filename)
if err != nil {
switch {
case errors.Is(err, document.ErrArtifactInvalidFilename),
@@ -516,9 +516,8 @@ func (h *DocumentHandler) DeleteDocuments(c *gin.Context) {
}
}
userID := c.GetString("user_id")
ctx := c.Request.Context()
deleted, err := h.documentService.DeleteDocuments(ctx, ids, req.DeleteAll, datasetID, userID)
deleted, err := h.documentService.DeleteDocuments(ctx, ids, req.DeleteAll, datasetID)
if err != nil {
if respondPermissionErrorIf(c, err, false) {
operationLog.Code = int(common.CodePermissionError)
@@ -536,14 +535,14 @@ func (h *DocumentHandler) DeleteDocuments(c *gin.Context) {
// BatchUpdateDocumentStatus Batch update status of documents within a dataset.
func (h *DocumentHandler) BatchUpdateDocumentStatus(c *gin.Context) {
user, code, errorMessage := GetUser(c)
_, code, errorMessage := GetUser(c)
if code != common.CodeSuccess {
common.RespondError(c, code, errorMessage)
return
}
userID := strings.TrimSpace(user.ID)
if userID == "" {
ctx := c.Request.Context()
if permission.UserIDFromContext(ctx) == "" {
common.RespondError(c, common.CodeArgumentError, "invalid user id")
return
}
@@ -586,8 +585,7 @@ func (h *DocumentHandler) BatchUpdateDocumentStatus(c *gin.Context) {
return
}
ctx := c.Request.Context()
result, code, err := h.documentService.BatchUpdateDocumentStatus(ctx, userID, datasetID, status, documentIDs)
result, code, err := h.documentService.BatchUpdateDocumentStatus(ctx, datasetID, status, documentIDs)
if err != nil {
message := err.Error()
if code == common.CodeServerError {
@@ -1344,14 +1342,14 @@ func stringValue(value *string) string {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/documents/ingest [post]
func (h *DocumentHandler) Ingest(c *gin.Context) {
user, errorCode, errorMessage := GetUser(c)
_, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.RespondError(c, errorCode, errorMessage)
return
}
userID := strings.TrimSpace(user.ID)
if userID == "" {
ctx := c.Request.Context()
if permission.UserIDFromContext(ctx) == "" {
common.RespondError(c, common.CodeAuthenticationError, "no authentication")
return
}
@@ -1361,8 +1359,7 @@ func (h *DocumentHandler) Ingest(c *gin.Context) {
common.RespondError(c, common.CodeBadRequest, err.Error())
return
}
ctx := c.Request.Context()
if code, err := h.documentService.Ingest(ctx, userID, &req); err != nil {
if code, err := h.documentService.Ingest(ctx, &req); err != nil {
common.RespondError(c, code, err.Error())
return
}
@@ -1381,8 +1378,6 @@ func (h *DocumentHandler) ListIngestionTasks(c *gin.Context) {
return
}
userID := c.GetString("user_id")
var parseResult []*entity.IngestionTask
var err error
ctx := c.Request.Context()
@@ -1392,7 +1387,7 @@ func (h *DocumentHandler) ListIngestionTasks(c *gin.Context) {
return
}
}
parseResult, err = h.documentService.ListIngestionTasks(ctx, userID, req.DatasetID, 0, 0)
parseResult, err = h.documentService.ListIngestionTasks(ctx, req.DatasetID, 0, 0)
if err != nil {
common.RespondError(c, IngestionTaskErrorCode(err), err.Error())
return
@@ -1450,7 +1445,7 @@ func (h *DocumentHandler) StartIngestionTask(c *gin.Context) {
return
}
parseResult, err := h.documentService.IngestDocuments(ctx, datasetID, userID, req.DocumentIDs)
parseResult, err := h.documentService.IngestDocuments(ctx, datasetID, req.DocumentIDs)
if err != nil {
common.RespondError(c, common.CodeExceptionError, err.Error())
return
@@ -1476,9 +1471,8 @@ func (h *DocumentHandler) StopIngestionTasks(c *gin.Context) {
return
}
userID := c.GetString("user_id")
ctx := c.Request.Context()
parseResult, err := h.documentService.StopIngestionTasks(ctx, req.Tasks, userID)
parseResult, err := h.documentService.StopIngestionTasks(ctx, req.Tasks)
if err != nil {
common.RespondError(c, IngestionTaskErrorCode(err), err.Error())
return
@@ -1502,10 +1496,9 @@ func (h *DocumentHandler) RemoveIngestionTasks(c *gin.Context) {
return
}
userID := c.GetString("user_id")
ctx := c.Request.Context()
deletedTasks, err := h.documentService.RemoveIngestionTasks(ctx, req.Tasks, userID)
deletedTasks, err := h.documentService.RemoveIngestionTasks(ctx, req.Tasks)
if err != nil {
common.RespondError(c, IngestionTaskErrorCode(err), err.Error())
return
@@ -1526,14 +1519,13 @@ func (h *DocumentHandler) ParseDocuments(c *gin.Context) {
return
}
userID := c.GetString("user_id")
ctx := c.Request.Context()
if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.ResourceRef{Kind: permission.ResourceKindDataset, ID: datasetID}, permission.OperationRun); err != nil {
respondPermissionError(c, err, false)
return
}
parseResult, err := h.documentService.ParseDocuments(ctx, datasetID, userID, req.Documents)
parseResult, err := h.documentService.ParseDocuments(ctx, datasetID, req.Documents)
if err != nil {
common.RespondError(c, common.CodeExceptionError, err.Error())
return
@@ -1713,7 +1705,7 @@ func (h *DocumentHandler) UpdateDatasetDocument(c *gin.Context) {
}
}
ctx := c.Request.Context()
data, code, err := h.documentService.UpdateDatasetDocument(ctx, user.ID, datasetID, documentID, &req, present)
data, code, err := h.documentService.UpdateDatasetDocument(ctx, datasetID, documentID, &req, present)
if err != nil {
common.RespondError(c, code, err.Error())
return
internal/handler/document_apply_parser_config.go
+2 -2
View File
@@ -10,7 +10,7 @@ import (
// ApplyDatasetParserConfig handles POST /api/v1/datasets/:dataset_id/documents/apply-parser-config.
func (h *DocumentHandler) ApplyDatasetParserConfig(c *gin.Context) {
user, code, message := GetUser(c)
_, code, message := GetUser(c)
if code != common.CodeSuccess {
common.RespondError(c, code, message)
return
@@ -20,7 +20,7 @@ func (h *DocumentHandler) ApplyDatasetParserConfig(c *gin.Context) {
common.RespondError(c, common.CodeArgumentError, "dataset_id is required")
return
}
result, code, err := h.documentService.ApplyDatasetParserConfig(c.Request.Context(), user.ID, datasetID)
result, code, err := h.documentService.ApplyDatasetParserConfig(c.Request.Context(), datasetID)
if err != nil {
common.RespondError(c, code, err.Error())
return
internal/handler/document_apply_parser_config_test.go
+7 -5
View File
@@ -10,12 +10,13 @@ import (
"ragflow/internal/common"
"ragflow/internal/entity"
"ragflow/internal/permission"
"ragflow/internal/service/document"
"github.com/gin-gonic/gin"
)
func (f *fakeDocumentService) ApplyDatasetParserConfig(ctx context.Context, userID, datasetID string) (*document.ApplyDatasetParserConfigResponse, common.ErrorCode, error) {
func (f *fakeDocumentService) ApplyDatasetParserConfig(ctx context.Context, datasetID string) (*document.ApplyDatasetParserConfigResponse, common.ErrorCode, error) {
return &document.ApplyDatasetParserConfigResponse{DatasetID: datasetID}, common.CodeSuccess, nil
}
@@ -27,10 +28,10 @@ type applyParserConfigService struct {
ctx context.Context
}
func (s *applyParserConfigService) ApplyDatasetParserConfig(ctx context.Context, userID, datasetID string) (*document.ApplyDatasetParserConfigResponse, common.ErrorCode, error) {
func (s *applyParserConfigService) ApplyDatasetParserConfig(ctx context.Context, datasetID string) (*document.ApplyDatasetParserConfigResponse, common.ErrorCode, error) {
s.called = true
s.ctx = ctx
if userID != "user-1" || datasetID != "kb-1" {
if permission.UserIDFromContext(ctx) != "user-1" || datasetID != "kb-1" {
return nil, common.CodeArgumentError, errors.New("wrong user or dataset")
}
return &document.ApplyDatasetParserConfigResponse{DatasetID: datasetID, TotalDocuments: 3, UpdatedDocuments: 2, UnchangedDocuments: 1}, s.code, s.err
@@ -55,6 +56,7 @@ func TestApplyDatasetParserConfigHandler(t *testing.T) {
r.POST("/api/v1/datasets/:dataset_id/documents/apply-parser-config", func(c *gin.Context) {
if tc.auth {
c.Set("user", &entity.User{ID: "user-1"})
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), "user-1", ""))
}
h.ApplyDatasetParserConfig(c)
})
@@ -77,8 +79,8 @@ func TestApplyDatasetParserConfigHandler(t *testing.T) {
}
return
}
if body.Code != tc.code || svc.ctx != req.Context() {
t.Fatalf("response=%s context propagated=%t", rec.Body.String(), svc.ctx == req.Context())
if body.Code != tc.code || permission.UserIDFromContext(svc.ctx) != "user-1" {
t.Fatalf("response=%s context identity=%q", rec.Body.String(), permission.UserIDFromContext(svc.ctx))
}
if tc.err == nil && (body.Data == nil || body.Data.UpdatedDocuments != 2 || body.Data.UnchangedDocuments != 1) {
t.Fatalf("response=%s", rec.Body.String())
internal/handler/document_test.go
+18 -17
View File
@@ -190,8 +190,8 @@ func TestMapDocumentListItemIncludesLatestIngestionEvent(t *testing.T) {
}
}
func (f *fakeDocumentService) Ingest(ctx context.Context, userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error) {
f.ingestUserID = userID
func (f *fakeDocumentService) Ingest(ctx context.Context, req *document.IngestDocumentRequest) (common.ErrorCode, error) {
f.ingestUserID = permission.UserIDFromContext(ctx)
f.ingestReq = req
if f.ingestCode != 0 || f.ingestErr != nil {
return f.ingestCode, f.ingestErr
@@ -201,14 +201,14 @@ func (f *fakeDocumentService) Ingest(ctx context.Context, userID string, req *do
const uploadTestDatasetID = "123e4567-e89b-12d3-a456-426614174000"
func (f *fakeDocumentService) UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) {
func (f *fakeDocumentService) UpdateDatasetDocument(ctx context.Context, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) {
return nil, common.CodeSuccess, nil
}
func (f *fakeDocumentService) BatchUpdateDocumentMetadatas(ctx context.Context, datasetID string, selector *document.MetadataSelector, updates []document.MetadataUpdate, deletes []document.MetadataDelete) (*document.BatchUpdateMetadatasResponse, common.ErrorCode, error) {
return nil, common.CodeSuccess, nil
}
func (f *fakeDocumentService) GetDocumentArtifact(ctx context.Context, filename, _ string) (*document.ArtifactResponse, error) {
func (f *fakeDocumentService) GetDocumentArtifact(ctx context.Context, filename string) (*document.ArtifactResponse, error) {
if filename == "error.txt" {
return nil, document.ErrArtifactNotFound
}
@@ -270,10 +270,10 @@ func (f *fakeDocumentService) DeleteDocument(ctx context.Context, id string) err
f.deletedID = id
return nil
}
func (f *fakeDocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error) {
func (f *fakeDocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID string) (int, error) {
return f.deleted, f.err
}
func (f *fakeDocumentService) ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
func (f *fakeDocumentService) ParseDocuments(ctx context.Context, datasetID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
return nil, nil
}
func (f *fakeDocumentService) StopParseDocuments(ctx context.Context, datasetID string, docIDs []string) (map[string]interface{}, error) {
@@ -319,21 +319,21 @@ func (f *fakeDocumentService) GetMetadataByKBs(ctx context.Context, kbIDs []stri
}
return map[string]interface{}{}, nil
}
func (f *fakeDocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) {
func (f *fakeDocumentService) BatchUpdateDocumentStatus(ctx context.Context, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) {
return map[string]interface{}{}, common.CodeSuccess, nil
}
func (f *fakeDocumentService) GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error) {
f.thumbnailUserID = userID
func (f *fakeDocumentService) GetThumbnails(ctx context.Context, docIDs []string) (map[string]string, error) {
f.thumbnailUserID = permission.UserIDFromContext(ctx)
f.thumbnailDocIDs = append([]string(nil), docIDs...)
return f.thumbnails, f.thumbnailErr
}
func (f *fakeDocumentService) GetDocumentImage(ctx context.Context, userID, imageID string) ([]byte, error) {
f.imageUserID = userID
func (f *fakeDocumentService) GetDocumentImage(ctx context.Context, imageID string) ([]byte, error) {
f.imageUserID = permission.UserIDFromContext(ctx)
f.imageID = imageID
return f.imageData, f.imageErr
}
func (f *fakeDocumentService) GetDocumentImageForDocument(ctx context.Context, userID, docID, imageID string) ([]byte, error) {
f.imageUserID = userID
func (f *fakeDocumentService) GetDocumentImageForDocument(ctx context.Context, docID, imageID string) ([]byte, error) {
f.imageUserID = permission.UserIDFromContext(ctx)
f.imageDocumentID = docID
f.imageID = imageID
return f.imageData, f.imageErr
@@ -370,16 +370,16 @@ func (f *fakeDocumentService) UploadEmptyDocument(ctx context.Context, kb *entit
return nil, common.CodeServerError, fmt.Errorf("not implemented")
}
func (f *fakeDocumentService) ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) {
func (f *fakeDocumentService) ListIngestionTasks(ctx context.Context, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) {
return nil, nil
}
func (f *fakeDocumentService) IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
func (f *fakeDocumentService) IngestDocuments(ctx context.Context, datasetID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
return nil, nil
}
func (f *fakeDocumentService) StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error) {
func (f *fakeDocumentService) StopIngestionTasks(ctx context.Context, tasks []string) ([]*entity.IngestionTask, error) {
return f.stopIngestionTasks, f.stopIngestionTaskErr
}
func (f *fakeDocumentService) RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error) {
func (f *fakeDocumentService) RemoveIngestionTasks(ctx context.Context, tasks []string) ([]map[string]string, error) {
return f.removeIngestionTasks, f.removeIngestionTaskErr
}
@@ -830,6 +830,7 @@ func setupDocumentIngestRoute(userID string, svc *fakeDocumentService) *gin.Engi
}
r := gin.New()
r.Use(func(c *gin.Context) {
c.Request = c.Request.WithContext(permission.WithIdentity(c.Request.Context(), userID, ""))
c.Set("user", &entity.User{ID: userID})
c.Set("user_id", userID)
})
internal/handler/mcp_retrieval.go
+14 -6
View File
@@ -8,6 +8,7 @@ import (
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/mcp"
"ragflow/internal/permission"
"ragflow/internal/service"
"ragflow/internal/service/document"
)
@@ -39,15 +40,22 @@ func validateRetrievalWindow(page, pageSize int) error {
// the MCP retrieval entry point can be tested against a recording fake without
// standing up the full DatasetService.
type mcpDatasetService interface {
ListDatasets(ctx context.Context, id, name string, page, pageSize int, terms []dao.OrderTerm, keywords string, ownerIDs []string, parserID, userID string, ids []string, canvasID string) ([]map[string]interface{}, int64, common.ErrorCode, error)
SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error)
GetDataset(ctx context.Context, datasetID, userID string) (map[string]interface{}, common.ErrorCode, error)
ListDatasets(ctx context.Context, id, name string, page, pageSize int, terms []dao.OrderTerm, keywords string, ownerIDs []string, parserID string, ids []string, canvasID string) ([]map[string]interface{}, int64, common.ErrorCode, error)
SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest) (*service.SearchDatasetsResponse, error)
GetDataset(ctx context.Context, datasetID string) (map[string]interface{}, common.ErrorCode, error)
}
// mcpRetrieval executes a retrieval request on behalf of the MCP tool handler.
// It translates the mcp.RetrievalRequest into a service.SearchDatasetsRequest
// and calls DatasetService.SearchDatasets. The result is serialized as JSON.
func mcpRetrieval(ctx context.Context, ds mcpDatasetService, userID string, req mcp.RetrievalRequest) (string, error) {
contextUserID := permission.UserIDFromContext(ctx)
if contextUserID == "" {
ctx = permission.WithIdentity(ctx, userID, permission.TenantIDFromContext(ctx))
} else if contextUserID != userID {
return "", permission.ErrIdentityMismatch
}
if err := validateRetrievalWindow(req.Page, req.PageSize); err != nil {
return "", err
}
@@ -59,7 +67,7 @@ func mcpRetrieval(ctx context.Context, ds mcpDatasetService, userID string, req
ids, err := fetchAllDatasetIDs(func(page, pageSize int) ([]map[string]interface{}, int64, error) {
data, total, _, err := ds.ListDatasets(ctx,
"", "", page, pageSize, []dao.OrderTerm{{Column: "create_time", Desc: true}},
"", nil, "", userID, nil, "",
"", nil, "", nil, "",
)
return data, total, err
}, maxPageSize)
@@ -111,7 +119,7 @@ func mcpRetrieval(ctx context.Context, ds mcpDatasetService, userID string, req
searchReq.Keyword = &v
}
resp, err := ds.SearchDatasets(ctx, searchReq, userID)
resp, err := ds.SearchDatasets(ctx, searchReq)
if err != nil {
return "", err
}
@@ -121,7 +129,7 @@ func mcpRetrieval(ctx context.Context, ds mcpDatasetService, userID string, req
datasetNames := make(map[string]string)
documentMetadata := make(map[string]map[string]any)
for _, id := range datasetIDs {
info, code, err := ds.GetDataset(ctx, id, userID)
info, code, err := ds.GetDataset(ctx, id)
if err != nil || code != common.CodeSuccess {
continue
}
internal/handler/mcp_retrieval_test.go
+5 -4
View File
@@ -13,6 +13,7 @@ import (
"ragflow/internal/dao"
"ragflow/internal/ingestion/knowledge_compile"
"ragflow/internal/mcp"
"ragflow/internal/permission"
"ragflow/internal/service"
)
@@ -35,17 +36,17 @@ type recordingDatasetService struct {
searchErr error
}
func (f *recordingDatasetService) ListDatasets(_ context.Context, _, _ string, page, pageSize int, _ []dao.OrderTerm, _ string, _ []string, _, _ string, _ []string, _ string) ([]map[string]interface{}, int64, common.ErrorCode, error) {
func (f *recordingDatasetService) ListDatasets(_ context.Context, _, _ string, page, pageSize int, _ []dao.OrderTerm, _ string, _ []string, _ string, _ []string, _ string) ([]map[string]interface{}, int64, common.ErrorCode, error) {
f.listCalls++
f.listPages = append(f.listPages, page)
f.listSizes = append(f.listSizes, pageSize)
return f.listRows, int64(len(f.listRows)), common.CodeSuccess, nil
}
func (f *recordingDatasetService) SearchDatasets(_ context.Context, req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) {
func (f *recordingDatasetService) SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest) (*service.SearchDatasetsResponse, error) {
f.searchCalls++
f.request = req
f.searchUser = userID
f.searchUser = permission.UserIDFromContext(ctx)
if f.searchErr != nil {
return nil, f.searchErr
}
@@ -55,7 +56,7 @@ func (f *recordingDatasetService) SearchDatasets(_ context.Context, req *service
return f.response, nil
}
func (f *recordingDatasetService) GetDataset(_ context.Context, datasetID, _ string) (map[string]interface{}, common.ErrorCode, error) {
func (f *recordingDatasetService) GetDataset(_ context.Context, datasetID string) (map[string]interface{}, common.ErrorCode, error) {
f.getCalls++
f.getIDs = append(f.getIDs, datasetID)
return map[string]interface{}{"name": "dataset-" + datasetID}, common.CodeSuccess, nil
internal/handler/mcp_server.go
+7 -1
View File
@@ -118,9 +118,15 @@ func (h *MCPServerHandler) HandleMCP(c *gin.Context) {
// mcpListDatasets wraps DatasetService.ListDatasets for the MCP tool handler,
// filling in default values for parameters that the MCP tool does not expose.
func mcpListDatasets(ctx context.Context, ds *dataset.DatasetService, userID string, page, pageSize int, orderby string, desc bool) ([]map[string]interface{}, int64, error) {
contextUserID := permission.UserIDFromContext(ctx)
if contextUserID == "" {
ctx = permission.WithIdentity(ctx, userID, permission.TenantIDFromContext(ctx))
} else if contextUserID != userID {
return nil, 0, permission.ErrIdentityMismatch
}
data, total, _, err := ds.ListDatasets(ctx,
"", "", page, pageSize, []dao.OrderTerm{{Column: orderby, Desc: desc}},
"", nil, "", userID, nil, "",
"", nil, "", nil, "",
)
return data, total, err
}
internal/service/chunk/chunk.go
+22 -22
View File
@@ -86,13 +86,13 @@ type ChunkService struct {
ingestionTaskDAO *dao.IngestionTaskDAO
searchService *service.SearchService
checkDatasetAccessFunc func(string, string) error
checkDatasetAccessFunc func(string) error
getKnowledgebaseByIDFunc func(string) (*entity.Knowledgebase, error)
getDocumentsByIDsFunc func([]string) ([]*entity.Document, error)
// startParseDocumentsFunc overrides the DSL start-parse flow. Production
// uses service.DocumentService.StartParseDocuments; tests inject a fake
// to avoid the MQ publisher.
startParseDocumentsFunc func(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, userID string, opts document.StartParseOptions) error
startParseDocumentsFunc func(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, opts document.StartParseOptions) error
// cancelIngestionTaskFunc overrides the document-parsing cancellation.
// Production uses service.DocumentService.CancelDocParse; tests inject
// a fake to avoid the MQ publisher.
@@ -557,7 +557,7 @@ func hydrateChunkVectors(ctx context.Context, engine engine.DocEngine, chunks []
}
// Get retrieves a chunk by ID
func (s *ChunkService) Get(ctx context.Context, req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error) {
func (s *ChunkService) Get(ctx context.Context, req *service.GetChunkRequest) (*service.GetChunkResponse, error) {
if s.docEngine == nil {
return nil, fmt.Errorf("doc engine not initialized")
}
@@ -566,7 +566,7 @@ func (s *ChunkService) Get(ctx context.Context, req *service.GetChunkRequest, us
return nil, fmt.Errorf("chunk_id is required")
}
kb, err := s.getAccessibleKnowledgebase(ctx, req.DatasetID, userID, permission.OperationRead)
kb, err := s.getAccessibleKnowledgebase(ctx, req.DatasetID, permission.OperationRead)
if err != nil {
return nil, err
}
@@ -634,7 +634,7 @@ func (s *ChunkService) cancelAllTasksOfDoc(ctx context.Context, doc *entity.Docu
return cancel(ctx, doc)
}
func (s *ChunkService) StopParsing(ctx context.Context, userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
func (s *ChunkService) StopParsing(ctx context.Context, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.ResourceRef{Kind: permission.ResourceKindDataset, ID: datasetID}, permission.OperationRun); err != nil {
code, permissionErr := permissionresponse.Normalize(err)
return nil, code, permissionErr
@@ -706,9 +706,9 @@ func checkDuplicateIDs(documentIDs []string, idTypes string) ([]string, []string
return uniqueDocIDs, duplicateMessages
}
func (s *ChunkService) checkDatasetAccess(ctx context.Context, datasetID, userID string, operation permission.Operation) error {
func (s *ChunkService) checkDatasetAccess(ctx context.Context, datasetID string, operation permission.Operation) error {
if s.checkDatasetAccessFunc != nil {
return s.checkDatasetAccessFunc(datasetID, userID)
return s.checkDatasetAccessFunc(datasetID)
}
return permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.ResourceRef{Kind: permission.ResourceKindDataset, ID: datasetID}, operation)
}
@@ -720,8 +720,8 @@ func (s *ChunkService) getKnowledgebaseByID(ctx context.Context, datasetID strin
return s.kbDAO.GetByID(ctx, dao.DB, datasetID)
}
func (s *ChunkService) getAccessibleKnowledgebase(ctx context.Context, datasetID, userID string, operation permission.Operation) (*entity.Knowledgebase, error) {
if err := s.checkDatasetAccess(ctx, datasetID, userID, operation); err != nil {
func (s *ChunkService) getAccessibleKnowledgebase(ctx context.Context, datasetID string, operation permission.Operation) (*entity.Knowledgebase, error) {
if err := s.checkDatasetAccess(ctx, datasetID, operation); err != nil {
return nil, err
}
kb, err := s.getKnowledgebaseByID(ctx, datasetID)
@@ -738,8 +738,8 @@ func (s *ChunkService) getDocumentsByIDs(ctx context.Context, docIDs []string) (
return s.documentDAO.GetByIDs(ctx, dao.DB, docIDs)
}
func (s *ChunkService) Parse(ctx context.Context, userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) {
if err := s.checkDatasetAccess(ctx, datasetID, userID, permission.OperationRun); err != nil {
func (s *ChunkService) Parse(ctx context.Context, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) {
if err := s.checkDatasetAccess(ctx, datasetID, permission.OperationRun); err != nil {
code, permissionErr := permissionresponse.Normalize(err)
return nil, code, permissionErr
}
@@ -789,7 +789,7 @@ func (s *ChunkService) Parse(ctx context.Context, userID, datasetID string, req
for _, docID := range docIDs {
doc := docByID[docID]
if err = startParse(ctx, doc, kb, userID, document.StartParseOptions{RerunWithDelete: true}); err != nil {
if err = startParse(ctx, doc, kb, document.StartParseOptions{RerunWithDelete: true}); err != nil {
return nil, common.CodeServerError, err
}
successCount++
@@ -808,7 +808,7 @@ func (s *ChunkService) Parse(ctx context.Context, userID, datasetID string, req
}
// List retrieves chunks for a document
func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) {
func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest) (*service.ListChunksResponse, error) {
if s.docEngine == nil {
return nil, fmt.Errorf("doc engine not initialized")
}
@@ -827,7 +827,7 @@ func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest,
return nil, fmt.Errorf("document not found")
}
kb, err := s.getAccessibleKnowledgebase(ctx, doc.KbID, userID, permission.OperationRead)
kb, err := s.getAccessibleKnowledgebase(ctx, doc.KbID, permission.OperationRead)
if err != nil {
return nil, err
}
@@ -1049,7 +1049,7 @@ func (s *ChunkService) List(ctx context.Context, req *service.ListChunksRequest,
}, nil
}
func (s *ChunkService) SwitchChunks(ctx context.Context, userID, datasetID, documentID string, availableInt int, chunkIDs []string) error {
func (s *ChunkService) SwitchChunks(ctx context.Context, datasetID, documentID string, availableInt int, chunkIDs []string) error {
if s.docEngine == nil {
return fmt.Errorf("doc engine not initialized")
}
@@ -1062,7 +1062,7 @@ func (s *ChunkService) SwitchChunks(ctx context.Context, userID, datasetID, docu
return fmt.Errorf("req is null")
}
kb, err := s.getAccessibleKnowledgebase(ctx, datasetID, userID, permission.OperationUpdate)
kb, err := s.getAccessibleKnowledgebase(ctx, datasetID, permission.OperationUpdate)
if err != nil {
return err
}
@@ -1095,7 +1095,7 @@ func (s *ChunkService) SwitchChunks(ctx context.Context, userID, datasetID, docu
return nil
}
func (s *ChunkService) UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest, userID string) error {
func (s *ChunkService) UpdateChunk(ctx context.Context, req *service.UpdateChunkRequest) error {
if s.docEngine == nil {
return fmt.Errorf("doc engine not initialized")
}
@@ -1104,7 +1104,7 @@ func (s *ChunkService) UpdateChunk(ctx context.Context, req *service.UpdateChunk
return fmt.Errorf("chunk_id is required")
}
kb, err := s.getAccessibleKnowledgebase(ctx, req.DatasetID, userID, permission.OperationUpdate)
kb, err := s.getAccessibleKnowledgebase(ctx, req.DatasetID, permission.OperationUpdate)
if err != nil {
return err
}
@@ -1271,7 +1271,7 @@ func (s *ChunkService) UpdateChunk(ctx context.Context, req *service.UpdateChunk
return nil
}
func (s *ChunkService) RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest, userID string) (int64, error) {
func (s *ChunkService) RemoveChunks(ctx context.Context, req *service.RemoveChunksRequest) (int64, error) {
if s.docEngine == nil {
return 0, fmt.Errorf("doc engine not initialized")
}
@@ -1290,7 +1290,7 @@ func (s *ChunkService) RemoveChunks(ctx context.Context, req *service.RemoveChun
return 0, permission.ErrResourceNotFound
}
kb, err := s.getAccessibleKnowledgebase(ctx, doc.KbID, userID, permission.OperationUpdate)
kb, err := s.getAccessibleKnowledgebase(ctx, doc.KbID, permission.OperationUpdate)
if err != nil {
return 0, err
}
@@ -1381,14 +1381,14 @@ func (s *ChunkService) nextSourceChunkOrder(ctx context.Context, indexName, data
return 0
}
func (s *ChunkService) AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
func (s *ChunkService) AddChunk(ctx context.Context, req *service.AddChunkRequest) (*service.AddChunkResponse, error) {
if s.docEngine == nil {
return nil, addChunkError{code: common.CodeServerError, message: "doc engine not initialized"}
}
if req == nil {
return nil, addChunkError{code: common.CodeDataError, message: "invalid request payload"}
}
if err := s.checkDatasetAccess(ctx, req.DatasetID, userID, permission.OperationUpdate); err != nil {
if err := s.checkDatasetAccess(ctx, req.DatasetID, permission.OperationUpdate); err != nil {
code, permissionErr := permissionresponse.Normalize(err)
return nil, addChunkError{code: code, message: permissionErr.Error(), cause: permissionErr}
}
internal/service/chunk/chunk_permissions_test.go
+5 -5
View File
@@ -68,17 +68,17 @@ func TestChunkManagementDatasetPermissions(t *testing.T) {
}}
_, err := s.Get(ctx, &service.GetChunkRequest{
DatasetID: "kb-1", DocumentID: "doc-1", ChunkID: "chunk-1",
}, tt.userID)
})
return err
}},
{"list", func(s *ChunkService) error {
s.docEngine = &listChunksSearchEngine{}
_, err := s.List(ctx, &service.ListChunksRequest{DatasetID: "kb-1", DocID: "doc-1"}, tt.userID)
_, err := s.List(ctx, &service.ListChunksRequest{DatasetID: "kb-1", DocID: "doc-1"})
return err
}},
{"switch", func(s *ChunkService) error {
s.docEngine = &updateChunkTestEngine{}
return s.SwitchChunks(ctx, tt.userID, "kb-1", "doc-1", 0, []string{"chunk-1"})
return s.SwitchChunks(ctx, "kb-1", "doc-1", 0, []string{"chunk-1"})
}},
{"update", func(s *ChunkService) error {
s.docEngine = &updateChunkTestEngine{existingChunk: map[string]interface{}{
@@ -86,11 +86,11 @@ func TestChunkManagementDatasetPermissions(t *testing.T) {
}}
return s.UpdateChunk(ctx, &service.UpdateChunkRequest{
DatasetID: "kb-1", DocumentID: "doc-1", ChunkID: "chunk-1",
}, tt.userID)
})
}},
{"remove", func(s *ChunkService) error {
s.docEngine = &parseTestDocEngine{}
_, err := s.RemoveChunks(ctx, &service.RemoveChunksRequest{DatasetID: "kb-1", DocID: "doc-1", DeleteAll: true}, tt.userID)
_, err := s.RemoveChunks(ctx, &service.RemoveChunksRequest{DatasetID: "kb-1", DocID: "doc-1", DeleteAll: true})
return err
}},
} {
internal/service/chunk/chunk_sdk_contract_test.go
+2 -2
View File
@@ -32,7 +32,7 @@ func TestSDKChunkListSchemaAndEmptyDeletion(t *testing.T) {
engine.chunks = append(engine.chunks, map[string]interface{}{"id": "chunk-2", "content_with_weight": "plain"})
svc := &ChunkService{docEngine: engine, kbDAO: dao.NewKnowledgebaseDAO()}
ctx := chunkTestContext(t, "user-1", "tenant-1")
resp, err := svc.List(ctx, &service.ListChunksRequest{DatasetID: "kb-1", DocID: "doc-1"}, "user-1")
resp, err := svc.List(ctx, &service.ListChunksRequest{DatasetID: "kb-1", DocID: "doc-1"})
if err != nil {
t.Fatal(err)
}
@@ -47,7 +47,7 @@ func TestSDKChunkListSchemaAndEmptyDeletion(t *testing.T) {
}
}
for _, ids := range [][]string{nil, {}} {
n, err := svc.RemoveChunks(ctx, &service.RemoveChunksRequest{DatasetID: "kb-1", DocID: "doc-1", ChunkIDs: ids}, "user-1")
n, err := svc.RemoveChunks(ctx, &service.RemoveChunksRequest{DatasetID: "kb-1", DocID: "doc-1", ChunkIDs: ids})
if err != nil || n != 0 {
t.Errorf("empty selection: count=%d err=%v", n, err)
}
internal/service/chunk/chunk_test.go
+56 -56
View File
@@ -181,7 +181,7 @@ func TestParsePrevalidatesDocumentsBeforeMutating(t *testing.T) {
}
ctx := chunkTestContext(t, userID, userID)
_, code, err := svc.Parse(ctx, userID, datasetID, &service.ParseFileRequest{
_, code, err := svc.Parse(ctx, datasetID, &service.ParseFileRequest{
DocumentIDs: []string{"doc-1", "missing-doc", "doc-2"},
})
if err == nil {
@@ -213,9 +213,9 @@ func TestParsePrevalidatesDocumentsBeforeMutating(t *testing.T) {
func TestParseRejectsInaccessibleDataset(t *testing.T) {
svc := newParseTestService(t)
svc.checkDatasetAccessFunc = func(string, string) error { return permission.ErrPermissionDenied }
svc.checkDatasetAccessFunc = func(string) error { return permission.ErrPermissionDenied }
ctx := t.Context()
_, code, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
_, code, err := svc.Parse(ctx, "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
if err == nil {
t.Fatal("expected parse to fail")
}
@@ -229,10 +229,10 @@ func TestParseRejectsInaccessibleDataset(t *testing.T) {
func TestParseRequiresDocumentIDs(t *testing.T) {
svc := newParseTestService(t)
svc.checkDatasetAccessFunc = func(string, string) error { return nil }
svc.checkDatasetAccessFunc = func(string) error { return nil }
ctx := t.Context()
for _, req := range []*service.ParseFileRequest{nil, {DocumentIDs: nil}, {DocumentIDs: []string{}}} {
_, code, err := svc.Parse(ctx, "user-1", "kb-1", req)
_, code, err := svc.Parse(ctx, "kb-1", req)
if err == nil {
t.Fatal("expected parse to fail")
}
@@ -247,10 +247,10 @@ func TestParseRequiresDocumentIDs(t *testing.T) {
func TestParseReturnsDataErrorWhenDatasetMissing(t *testing.T) {
svc := newParseTestService(t)
svc.checkDatasetAccessFunc = func(string, string) error { return nil }
svc.checkDatasetAccessFunc = func(string) error { return nil }
svc.getKnowledgebaseByIDFunc = func(string) (*entity.Knowledgebase, error) { return nil, nil }
ctx := t.Context()
_, code, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
_, code, err := svc.Parse(ctx, "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
if err == nil {
t.Fatal("expected parse to fail")
}
@@ -265,7 +265,7 @@ func TestParseReturnsDataErrorWhenDatasetMissing(t *testing.T) {
func TestParseReturnsServerErrorWhenDocumentsQueryFails(t *testing.T) {
svc := newParseTestService(t)
queryErr := errors.New("documents query failed")
svc.checkDatasetAccessFunc = func(string, string) error { return nil }
svc.checkDatasetAccessFunc = func(string) error { return nil }
svc.getKnowledgebaseByIDFunc = func(string) (*entity.Knowledgebase, error) {
return &entity.Knowledgebase{ID: "kb-1", TenantID: "user-1"}, nil
}
@@ -273,7 +273,7 @@ func TestParseReturnsServerErrorWhenDocumentsQueryFails(t *testing.T) {
return nil, queryErr
}
ctx := t.Context()
_, code, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
_, code, err := svc.Parse(ctx, "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
if !errors.Is(err, queryErr) {
t.Fatalf("expected query error, got %v", err)
}
@@ -294,7 +294,7 @@ func TestParseRejectsRunningDocument(t *testing.T) {
svc := newParseTestService(t)
ctx := chunkTestContext(t, userID, userID)
_, code, err := svc.Parse(ctx, userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
_, code, err := svc.Parse(ctx, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
if err == nil {
t.Fatal("expected parse to fail")
}
@@ -333,7 +333,7 @@ func TestListSortsChunksByDocumentPosition(t *testing.T) {
DocID: documentID,
Page: &page,
Size: &size,
}, userID)
})
if err != nil {
t.Fatalf("List() error = %v", err)
}
@@ -408,7 +408,7 @@ func TestListReportsChunkAvailability(t *testing.T) {
DocID: documentID,
Page: &page,
Size: &size,
}, userID)
})
if err != nil {
t.Fatalf("List() error = %v", err)
}
@@ -463,7 +463,7 @@ func TestListScopesSearchToDocumentDataset(t *testing.T) {
DocID: documentID,
Page: &page,
Size: &size,
}, userID); err != nil {
}); err != nil {
t.Fatalf("List() error = %v", err)
}
@@ -510,7 +510,7 @@ func TestGetScopesChunkLookupToDocumentDataset(t *testing.T) {
DatasetID: datasetID,
DocumentID: documentID,
ChunkID: "chunk-1",
}, userID)
})
if err != nil {
t.Fatalf("Get() error = %v", err)
}
@@ -572,7 +572,7 @@ func TestGetRejectsChunkOfAnotherDocument(t *testing.T) {
DatasetID: datasetID,
DocumentID: documentID,
ChunkID: "chunk-1",
}, userID)
})
if err == nil || !strings.Contains(err.Error(), "chunk not found") {
t.Fatalf("Get() error = %v, want chunk not found", err)
}
@@ -607,7 +607,7 @@ func TestGetRequiresDocumentInRouteDataset(t *testing.T) {
DatasetID: datasetID,
DocumentID: "doc-missing",
ChunkID: "chunk-1",
}, userID)
})
if err == nil || !strings.Contains(err.Error(), "document not found") {
t.Fatalf("Get() error = %v, want document not found", err)
}
@@ -616,7 +616,7 @@ func TestGetRequiresDocumentInRouteDataset(t *testing.T) {
DatasetID: datasetID,
DocumentID: "doc-2",
ChunkID: "chunk-1",
}, userID)
})
if err == nil || !strings.Contains(err.Error(), "document does not belong to this dataset") {
t.Fatalf("Get() error = %v, want document does not belong to this dataset", err)
}
@@ -694,7 +694,7 @@ func TestListBuildsMatchTextExprForKeywords(t *testing.T) {
Page: &page,
Size: &size,
Keywords: " invoice terms ",
}, userID)
})
if err != nil {
t.Fatalf("List() error = %v", err)
}
@@ -773,7 +773,7 @@ func TestUpdateChunkRejectsChunkFromAnotherDocument(t *testing.T) {
DatasetID: "kb-1",
DocumentID: "doc-a",
ChunkID: "chunk-1",
}, "user-1")
})
if err == nil {
t.Fatal("expected UpdateChunk to reject a chunk from another document")
}
@@ -808,7 +808,7 @@ func TestUpdateChunkUpdatesSameDocumentWithDocumentCondition(t *testing.T) {
DatasetID: "kb-1",
DocumentID: "doc-a",
ChunkID: "chunk-1",
}, "user-1")
})
if err != nil {
t.Fatalf("UpdateChunk() error = %v", err)
}
@@ -863,7 +863,7 @@ func TestUpdateChunkStoresTagKwdOnlyWhenSent(t *testing.T) {
DocumentID: "doc-a",
ChunkID: "chunk-1",
TagKwd: tt.tagKwd,
}, "user-1")
})
if err != nil {
t.Fatalf("UpdateChunk() error = %v", err)
}
@@ -922,7 +922,7 @@ func TestUpdateChunkStoresImageAndFlagsImageChunk(t *testing.T) {
ChunkID: "chunk-1",
ImageBase64: &imageBase64,
ImageUpdateMode: tc.mode,
}, "user-1")
})
if err != nil {
t.Fatalf("UpdateChunk() error = %v", err)
}
@@ -987,7 +987,7 @@ func TestUpdateChunkRemoveModeClearsImageAfterIndexUpdate(t *testing.T) {
DocumentID: "doc-a",
ChunkID: "chunk-1",
ImageUpdateMode: &removeMode,
}, "user-1")
})
if err != nil {
t.Fatalf("UpdateChunk() error = %v", err)
}
@@ -1043,7 +1043,7 @@ func TestUpdateChunkMarksWikiDirtyWhenImageRemovalFails(t *testing.T) {
ChunkID: "chunk-1",
Content: &content,
ImageUpdateMode: &removeMode,
}, "user-1")
})
if err == nil || !strings.Contains(err.Error(), "Failed to remove chunk image") {
t.Fatalf("UpdateChunk() error = %v, want the removal failure", err)
}
@@ -1098,7 +1098,7 @@ func TestUpdateChunkRejectsInvalidImageInput(t *testing.T) {
ChunkID: "chunk-1",
ImageBase64: tc.imageB64,
ImageUpdateMode: tc.mode,
}, "user-1")
})
if err == nil || !strings.Contains(err.Error(), tc.wantError) {
t.Fatalf("UpdateChunk() error = %v, want %q", err, tc.wantError)
}
@@ -1152,8 +1152,8 @@ func TestAddChunkSuccess(t *testing.T) {
docEngine: engine,
kbDAO: dao.NewKnowledgebaseDAO(),
documentDAO: dao.NewDocumentDAO(),
checkDatasetAccessFunc: func(datasetIDArg, userIDArg string) error {
if datasetIDArg == datasetID && userIDArg == userID {
checkDatasetAccessFunc: func(datasetIDArg string) error {
if datasetIDArg == datasetID {
return nil
}
return permission.ErrPermissionDenied
@@ -1197,7 +1197,7 @@ func TestAddChunkSuccess(t *testing.T) {
ImportantKeywords: []string{"k1"},
Questions: []string{" q1 ", ""},
TagFeas: map[string]interface{}{"tag1": float64(0.5)},
}, userID)
})
if err != nil {
t.Fatalf("AddChunk() error = %v", err)
}
@@ -1306,8 +1306,8 @@ func TestAddChunkPersistsTagKwdOnlyWhenProvided(t *testing.T) {
docEngine: engine,
kbDAO: dao.NewKnowledgebaseDAO(),
documentDAO: dao.NewDocumentDAO(),
checkDatasetAccessFunc: func(datasetIDArg, userIDArg string) error {
if datasetIDArg != datasetID || userIDArg != userID {
checkDatasetAccessFunc: func(datasetIDArg string) error {
if datasetIDArg != datasetID {
return permission.ErrPermissionDenied
}
return nil
@@ -1339,7 +1339,7 @@ func TestAddChunkPersistsTagKwdOnlyWhenProvided(t *testing.T) {
DatasetID: datasetID,
DocumentID: documentID,
Content: "chunk body",
}, userID)
})
if err != nil {
t.Fatalf("AddChunk() error = %v", err)
}
@@ -1361,7 +1361,7 @@ func TestAddChunkPersistsTagKwdOnlyWhenProvided(t *testing.T) {
DocumentID: documentID,
Content: "chunk body",
TagKwd: []string{},
}, userID)
})
if err != nil {
t.Fatalf("AddChunk() error = %v", err)
}
@@ -1385,7 +1385,7 @@ func TestAddChunkPersistsTagKwdOnlyWhenProvided(t *testing.T) {
DocumentID: documentID,
Content: "chunk body",
TagKwd: want,
}, userID)
})
if err != nil {
t.Fatalf("AddChunk() error = %v", err)
}
@@ -1434,7 +1434,7 @@ func TestAddChunkValidationErrors(t *testing.T) {
name: "empty content",
req: &service.AddChunkRequest{DatasetID: "kb-1", DocumentID: "doc-1", Content: " "},
setup: func(svc *ChunkService) {
svc.checkDatasetAccessFunc = func(string, string) error { return nil }
svc.checkDatasetAccessFunc = func(string) error { return nil }
svc.getKnowledgebaseByIDFunc = func(string) (*entity.Knowledgebase, error) {
return &entity.Knowledgebase{ID: "kb-1", TenantID: "user-1", EmbdID: "embed-1"}, nil
}
@@ -1448,7 +1448,7 @@ func TestAddChunkValidationErrors(t *testing.T) {
if tt.setup != nil {
tt.setup(svc)
}
_, err := svc.AddChunk(ctx, tt.req, "user-1")
_, err := svc.AddChunk(ctx, tt.req)
if err == nil || !strings.Contains(err.Error(), tt.wantMsg) {
t.Fatalf("error = %v, want substring %q", err, tt.wantMsg)
}
@@ -1469,7 +1469,7 @@ func TestAddChunkImageAndTagFeatureValidation(t *testing.T) {
docEngine: &addChunkTestEngine{},
kbDAO: dao.NewKnowledgebaseDAO(),
documentDAO: dao.NewDocumentDAO(),
checkDatasetAccessFunc: func(string, string) error { return nil },
checkDatasetAccessFunc: func(string) error { return nil },
getKnowledgebaseByIDFunc: func(id string) (*entity.Knowledgebase, error) {
return &entity.Knowledgebase{ID: id, TenantID: userID, EmbdID: "embed-1"}, nil
},
@@ -1504,7 +1504,7 @@ func TestAddChunkImageAndTagFeatureValidation(t *testing.T) {
DocumentID: documentID,
Content: "chunk body",
ImageBase64: strPtr("not-base64"),
}, userID)
})
if err == nil || !strings.Contains(err.Error(), "invalid `image_base64`") {
t.Fatalf("expected invalid image error, got %v", err)
}
@@ -1516,7 +1516,7 @@ func TestAddChunkImageAndTagFeatureValidation(t *testing.T) {
DocumentID: documentID,
Content: "chunk body",
TagFeas: map[string]interface{}{"tag1": "bad"},
}, userID)
})
if err == nil || !strings.Contains(err.Error(), "`tag_feas` values must be finite numbers") || resp != nil {
t.Fatalf("expected tag_feas validation error, got resp=%#v err=%v", resp, err)
}
@@ -1527,7 +1527,7 @@ func TestAddChunkImageAndTagFeatureValidation(t *testing.T) {
Content: "chunk body",
TagFeas: map[string]interface{}{"tag1": float64(1)},
ImageBase64: strPtr(validJPEG),
}, userID)
})
if err != nil {
t.Fatalf("AddChunk() with image error = %v", err)
}
@@ -1553,7 +1553,7 @@ func TestAddChunkIncrementsStatsAfterInsert(t *testing.T) {
docEngine: engine,
kbDAO: dao.NewKnowledgebaseDAO(),
documentDAO: dao.NewDocumentDAO(),
checkDatasetAccessFunc: func(string, string) error { return nil },
checkDatasetAccessFunc: func(string) error { return nil },
getKnowledgebaseByIDFunc: func(id string) (*entity.Knowledgebase, error) {
return &entity.Knowledgebase{ID: id, TenantID: userID, EmbdID: "embed-1"}, nil
},
@@ -1580,7 +1580,7 @@ func TestAddChunkIncrementsStatsAfterInsert(t *testing.T) {
DatasetID: datasetID,
DocumentID: documentID,
Content: "chunk body",
}, userID)
})
if err != nil {
t.Fatalf("AddChunk() error = %v", err)
}
@@ -1773,7 +1773,7 @@ func TestRemoveChunksRejectsInvalidDatasetBeforeMutation(t *testing.T) {
if !deleteAll {
req.ChunkIDs = []string{"chunk-1"}
}
count, err := svc.RemoveChunks(t.Context(), req, "owner")
count, err := svc.RemoveChunks(t.Context(), req)
if err == nil || count != 0 {
t.Fatalf("invalid dataset deletion = (%d, %v), want error and zero count", count, err)
}
@@ -1806,7 +1806,7 @@ func TestRemoveChunksDecrementsStatsAfterDelete(t *testing.T) {
DatasetID: "kb-1",
DocID: "doc-1",
ChunkIDs: []string{"chunk-1", "chunk-2", "chunk-3"},
}, "user-1")
})
if err != nil {
t.Fatalf("RemoveChunks() error = %v", err)
}
@@ -1871,7 +1871,7 @@ func TestRemoveChunksSkipsStatsWhenNothingDeleted(t *testing.T) {
DatasetID: "kb-1",
DocID: "doc-1",
DeleteAll: true,
}, "user-1")
})
if err != nil {
t.Fatalf("RemoveChunks() error = %v", err)
}
@@ -1905,7 +1905,7 @@ func TestRemoveChunksReturnsStatsError(t *testing.T) {
DatasetID: "kb-1",
DocID: "doc-1",
ChunkIDs: []string{"chunk-1", "chunk-2"},
}, "user-1")
})
if err == nil || !strings.Contains(err.Error(), "failed to update chunk stats") {
t.Fatalf("expected stats update error, got count=%d err=%v", deletedCount, err)
}
@@ -2543,7 +2543,7 @@ func TestSwitchChunksUpdatesDocEngineWithAvailableInt(t *testing.T) {
}
ctx := chunkTestContext(t, "user-1", "tenant-1")
if err = svc.SwitchChunks(ctx, "user-1", "kb-1", "doc-1", 0, []string{"chunk-1", "chunk-2"}); err != nil {
if err = svc.SwitchChunks(ctx, "kb-1", "doc-1", 0, []string{"chunk-1", "chunk-2"}); err != nil {
t.Fatalf("SwitchChunks() error = %v", err)
}
@@ -2679,9 +2679,9 @@ func TestChunkServiceParse_RejectsBatchWithRunningIngestionTask(t *testing.T) {
insertChunkTestIngestionTask(t, "task-2", "user-1", "doc-2", "kb-1", common.RUNNING)
svc := newParseTestService(t)
svc.checkDatasetAccessFunc = func(string, string) error { return nil }
svc.checkDatasetAccessFunc = func(string) error { return nil }
ctx := t.Context()
_, _, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1", "doc-2"}})
_, _, err := svc.Parse(ctx, "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1", "doc-2"}})
if err == nil {
t.Fatal("expected error for RUNNING ingestion task, got nil")
}
@@ -2695,16 +2695,16 @@ func TestChunkServiceParse_CallsStartParseDocumentsWithRerunWithDelete(t *testin
insertChunkTestDoc(t, "doc-1", "kb-1")
svc := newParseTestService(t)
svc.checkDatasetAccessFunc = func(string, string) error { return nil }
svc.checkDatasetAccessFunc = func(string) error { return nil }
var calledOpts document.StartParseOptions
var calledDocID string
svc.startParseDocumentsFunc = func(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, userID string, opts document.StartParseOptions) error {
svc.startParseDocumentsFunc = func(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, opts document.StartParseOptions) error {
calledOpts = opts
calledDocID = doc.ID
return nil
}
ctx := t.Context()
_, code, err := svc.Parse(ctx, "user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
_, code, err := svc.Parse(ctx, "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
if err != nil {
t.Fatalf("Parse: %v", err)
}
@@ -2727,12 +2727,12 @@ func TestChunkServiceParse_ReturnsPartialSuccessForDuplicateDocumentIDs(t *testi
insertChunkTestDoc(t, "doc-1", datasetID)
svc := newParseTestService(t)
svc.startParseDocumentsFunc = func(context.Context, *entity.Document, *entity.Knowledgebase, string, document.StartParseOptions) error {
svc.startParseDocumentsFunc = func(context.Context, *entity.Document, *entity.Knowledgebase, document.StartParseOptions) error {
return nil
}
ctx := chunkTestContext(t, userID, userID)
result, code, err := svc.Parse(ctx, userID, datasetID, &service.ParseFileRequest{
result, code, err := svc.Parse(ctx, datasetID, &service.ParseFileRequest{
DocumentIDs: []string{"doc-1", "doc-1"},
})
if err == nil {
@@ -2763,7 +2763,7 @@ func TestStopParsing_CallsCancelIngestionTask(t *testing.T) {
return nil
}
ctx := chunkTestContext(t, "user-1", "user-1")
_, _, err := svc.StopParsing(ctx, "user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
_, _, err := svc.StopParsing(ctx, "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
if err != nil {
t.Fatalf("StopParsing: %v", err)
}
@@ -2786,7 +2786,7 @@ func TestStopParsing_RejectsDocumentWithoutIngestionTask(t *testing.T) {
return document.ErrParseNotRunning
}
ctx := chunkTestContext(t, "user-1", "user-1")
_, code, err := svc.StopParsing(ctx, "user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
_, code, err := svc.StopParsing(ctx, "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
if err == nil {
t.Fatal("expected error for document without ingestion task")
}
@@ -2821,7 +2821,7 @@ func TestStopParsing_DoesNotDeleteChunksOrResetCountersAfterCancel(t *testing.T)
engine := &parseTestDocEngine{chunkStoreExists: true}
svc.docEngine = engine
ctx := chunkTestContext(t, "user-1", "user-1")
_, _, err := svc.StopParsing(ctx, "user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
_, _, err := svc.StopParsing(ctx, "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
if err != nil {
t.Fatalf("StopParsing: %v", err)
}
internal/service/dataset/compilation_status.go
+2 -2
View File
@@ -36,9 +36,9 @@ type CompilationStatus struct {
}
// GetDatasetCompilationStatus returns the scheduling-row lifecycle state for a
// dataset after verifying the calling user owns it. When no row exists the
// dataset after verifying the caller can read it. When no row exists the
// dataset has never had any compile work, so the state is idle.
func (d *DatasetService) GetDatasetCompilationStatus(ctx context.Context, userID, datasetID, kind string) (CompilationStatus, common.ErrorCode, error) {
func (d *DatasetService) GetDatasetCompilationStatus(ctx context.Context, datasetID, kind string) (CompilationStatus, common.ErrorCode, error) {
if datasetID == "" {
return CompilationStatus{}, common.CodeDataError, errors.New("dataset_id is required")
}
internal/service/dataset/compilation_status_test.go
+4 -4
View File
@@ -99,7 +99,7 @@ func TestGetDatasetCompilationStatus_NoRowIsIdle(t *testing.T) {
insertCompilationOwnerKB(t, "kb-no-row", "user-1")
st, code, err := testCompilationStatusService().GetDatasetCompilationStatus(
withDatasetTestContext(t.Context(), "user-1", "user-1"), "user-1", "kb-no-row", "graph")
withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-no-row", "graph")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -156,7 +156,7 @@ func TestGetDatasetCompilationStatus_FullOutput(t *testing.T) {
}
st, code, err := testCompilationStatusService().GetDatasetCompilationStatus(
withDatasetTestContext(t.Context(), "user-1", "user-1"), "user-1", "kb-full", "graph")
withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-full", "graph")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -198,7 +198,7 @@ func TestGetDatasetCompilationStatus_Unauthorized(t *testing.T) {
}
st, code, err := testCompilationStatusService().GetDatasetCompilationStatus(
withDatasetTestContext(t.Context(), "user-2", "user-1"), "user-2", "kb-other", "graph")
withDatasetTestContext(t.Context(), "user-2", "user-1"), "kb-other", "graph")
if err == nil {
t.Fatalf("expected authorization error, got nil (status=%+v)", st)
}
@@ -211,7 +211,7 @@ func TestGetDatasetCompilationStatus_Unauthorized(t *testing.T) {
func TestGetDatasetCompilationStatus_EmptyID(t *testing.T) {
setupCompilationStatusTestDB(t)
_, code, err := testCompilationStatusService().GetDatasetCompilationStatus(
t.Context(), "user-1", "", "graph")
t.Context(), "", "graph")
if err == nil {
t.Fatal("expected error for empty dataset_id")
}
internal/service/dataset/crud.go
+4 -4
View File
@@ -328,7 +328,7 @@ func (d *DatasetService) dedupeDatasetName(ctx context.Context, name, tenantID s
})
}
func (d *DatasetService) GetDataset(ctx context.Context, datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) GetDataset(ctx context.Context, datasetID string) (map[string]interface{}, common.ErrorCode, error) {
datasetID = strings.TrimSpace(datasetID)
if datasetID == "" {
return nil, common.CodeDataError, errors.New("lack of \"Dataset ID\"")
@@ -588,7 +588,7 @@ func datasetNameAndID(kb *entity.Knowledgebase) string {
return fmt.Sprintf("%s (%s)", kb.Name, kb.ID)
}
func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page, pageSize int, terms []dao.OrderTerm, keywords string, ownerIDs []string, parserID, userID string, ids []string, canvasID string) ([]map[string]interface{}, int64, common.ErrorCode, error) {
func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page, pageSize int, terms []dao.OrderTerm, keywords string, ownerIDs []string, parserID string, ids []string, canvasID string) ([]map[string]interface{}, int64, common.ErrorCode, error) {
id = strings.TrimSpace(id)
canvasID = strings.TrimSpace(canvasID)
if id != "" && len(ids) > 0 {
@@ -722,7 +722,7 @@ func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page
}
if len(deniedIDs) > 0 {
common.Warn("User lacks permission for datasets",
zap.String("user_id", userID),
zap.String("user_id", permission.UserIDFromContext(ctx)),
zap.Strings("dataset_ids", common.Deduplicate(deniedIDs)),
)
}
@@ -751,7 +751,7 @@ func (d *DatasetService) ListDatasets(ctx context.Context, id, name string, page
return data, total, common.CodeSuccess, nil
}
func (d *DatasetService) ListDatasetFilters(ctx context.Context, userID, canvasID string) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) ListDatasetFilters(ctx context.Context, canvasID string) (map[string]interface{}, common.ErrorCode, error) {
owners := make([]*entity.DatasetOwnerFilter, 0)
err := dao.DB.Transaction(func(tx *gorm.DB) error {
query := permission.ScopeQuery{Kind: permission.ResourceKindDataset, Operation: permission.OperationRead}
internal/service/dataset/embedding.go
+1 -1
View File
@@ -37,7 +37,7 @@ type embeddingCheckSample struct {
// CheckEmbedding verifies that a candidate embedding model is compatible with a
// dataset's existing vectors (the standard "switch embedding model" validation).
// It is independent of the retired RunIndex/graph_rag_queue scheduling path.
func (d *DatasetService) CheckEmbedding(ctx context.Context, userID, datasetID string, req *service.CheckEmbeddingRequest) (*service.EmbeddingCheckResponse, common.ErrorCode, error) {
func (d *DatasetService) CheckEmbedding(ctx context.Context, datasetID string, req *service.CheckEmbeddingRequest) (*service.EmbeddingCheckResponse, common.ErrorCode, error) {
if datasetID == "" {
return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`)
}
internal/service/dataset/empty_version_test.go
+3 -3
View File
@@ -114,7 +114,7 @@ func TestUpdateDataset_ExplicitEmptyVersionRejected(t *testing.T) {
parseType := 2
empty := ""
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-empty-upd", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-empty-upd", service.UpdateDatasetRequest{
PipelineID: &canvas1,
ParseType: &parseType,
PipelineVersionID: &empty,
@@ -145,7 +145,7 @@ func TestUpdateDataset_PipelineChangeWithoutVersionRejected(t *testing.T) {
insertDatasetUpdateKB(t, "kb-switch", "tenant-1", "Original") // builtin, no pipeline bound
parseType := 2
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-switch", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-switch", service.UpdateDatasetRequest{
PipelineID: &canvas1,
ParseType: &parseType,
})
@@ -170,7 +170,7 @@ func TestUpdateDataset_PreservePinnedVersionOnPartialUpdate(t *testing.T) {
seedBoundKB(t, db, "kb-preserve", canvas1, ver1)
name := "renamed-partial"
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-preserve", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-preserve", service.UpdateDatasetRequest{
Name: &name,
})
if err != nil || code != common.CodeSuccess {
internal/service/dataset/gap_b1b2_test.go
+2 -2
View File
@@ -103,7 +103,7 @@ func TestUpdateDataset_ParserConfigRebuildUsesPinnedVersionDSL(t *testing.T) {
seedBoundKBWithVersion(t, db, "kb-1", canvasID, "v-old")
// Request omits pipeline_id AND pipeline_version_id; only parser_config changes.
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1",
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1",
service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{"Parser:Shared": map[string]interface{}{"k": "override"}},
})
@@ -144,7 +144,7 @@ func TestUpdateDataset_SingleFieldVersionPatchDoesNotClearPin(t *testing.T) {
seedBoundKBWithVersion(t, db, "kb-1", canvasID, "v-old")
// Single-field PATCH: only pipeline_version_id, pipeline_id omitted.
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1",
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1",
service.UpdateDatasetRequest{
PipelineVersionID: sptr("v-new"),
})
internal/service/dataset/ingestion.go
+5 -5
View File
@@ -44,7 +44,7 @@ type IngestionMessagesResponse struct {
Terminal bool `json:"terminal"`
}
func (d *DatasetService) GetIngestionSummary(ctx context.Context, datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) GetIngestionSummary(ctx context.Context, datasetID string) (map[string]interface{}, common.ErrorCode, error) {
if datasetID == "" {
return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`)
}
@@ -149,7 +149,7 @@ func (d *DatasetService) getDownloadStatus(ctx context.Context, datasetID string
// ListIngestionMessages returns events owned by the requested immutable
// pipeline log.
func (d *DatasetService) ListIngestionMessages(ctx context.Context, datasetID, userID, logID string, limit int, afterID, beforeID *int) (*IngestionMessagesResponse, common.ErrorCode, error) {
func (d *DatasetService) ListIngestionMessages(ctx context.Context, datasetID, logID string, limit int, afterID, beforeID *int) (*IngestionMessagesResponse, common.ErrorCode, error) {
if datasetID == "" {
return nil, common.CodeArgumentError, errors.New(`lack of "Dataset ID"`)
}
@@ -229,7 +229,7 @@ func isReadableIngestionLog(log *entity.PipelineOperationLog) bool {
}
// ListSyncLogs returns paginated sync logs for connectors currently linked to an accessible dataset.
func (d *DatasetService) ListSyncLogs(ctx context.Context, datasetID, userID string, page, pageSize int) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) ListSyncLogs(ctx context.Context, datasetID string, page, pageSize int) (map[string]interface{}, common.ErrorCode, error) {
if datasetID == "" {
return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`)
}
@@ -287,7 +287,7 @@ func (d *DatasetService) ListSyncLogs(ctx context.Context, datasetID, userID str
}, common.CodeSuccess, nil
}
func (d *DatasetService) ListIngestionLogs(ctx context.Context, datasetID, userID string, page, pageSize int, terms []dao.OrderTerm, operationStatus []string, createDateFrom, createDateTo, logType, keywords, documentID string) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) ListIngestionLogs(ctx context.Context, datasetID string, page, pageSize int, terms []dao.OrderTerm, operationStatus []string, createDateFrom, createDateTo, logType, keywords, documentID string) (map[string]interface{}, common.ErrorCode, error) {
if datasetID == "" {
return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`)
}
@@ -349,7 +349,7 @@ func (d *DatasetService) ListIngestionLogs(ctx context.Context, datasetID, userI
}, common.CodeSuccess, nil
}
func (d *DatasetService) GetIngestionLog(ctx context.Context, datasetID, userID, logID string) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) GetIngestionLog(ctx context.Context, datasetID, logID string) (map[string]interface{}, common.ErrorCode, error) {
if datasetID == "" {
return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`)
}
internal/service/dataset/ingestion_messages_test.go
+19 -19
View File
@@ -24,7 +24,7 @@ func TestListIngestionMessagesUsesRunScopedKeysets(t *testing.T) {
insertMessageEvent(t, db, "run-1", "task-1", dao.EventTypeMessage, "detail")
}
response, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "run-1", 2, nil, nil)
response, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "run-1", 2, nil, nil)
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListIngestionMessages = (%+v, %v, %v), want success", response, code, err)
}
@@ -34,7 +34,7 @@ func TestListIngestionMessagesUsesRunScopedKeysets(t *testing.T) {
assertMessageEventIDs(t, response.Items, 4, 5)
afterID := 2
response, code, err = NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "run-1", 2, &afterID, nil)
response, code, err = NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "run-1", 2, &afterID, nil)
if err != nil || code != common.CodeSuccess {
t.Fatalf("after page = (%+v, %v, %v), want success", response, code, err)
}
@@ -44,7 +44,7 @@ func TestListIngestionMessagesUsesRunScopedKeysets(t *testing.T) {
}
beforeID := 4
if _, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "run-1", 2, &afterID, &beforeID); err == nil || code != common.CodeArgumentError {
if _, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "run-1", 2, &afterID, &beforeID); err == nil || code != common.CodeArgumentError {
t.Fatalf("mutually exclusive cursors = code:%v err:%v, want argument error", code, err)
}
}
@@ -59,7 +59,7 @@ func TestListIngestionMessagesRejectsUnnumberedRun(t *testing.T) {
insertMessageRun(t, "legacy-run", "kb-1", "doc-1", entity.TaskStatusRunning, 0)
insertMessageEvent(t, db, "legacy-run", "task-1", dao.EventTypeMessage, "old detail")
response, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "legacy-run", 200, nil, nil)
response, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "legacy-run", 200, nil, nil)
if response != nil || err == nil || code != common.CodeDataError {
t.Fatalf("ListIngestionMessages = (%+v, %v, %v), want not-found data error", response, code, err)
}
@@ -89,7 +89,7 @@ func TestDatasetIngestionLogUsesEventStream(t *testing.T) {
insertMessageRun(t, "numbered-sentinel-run", "kb-1", entity.DatasetLogDocumentID, entity.TaskStatusDone, 1)
insertMessageEvent(t, db, "dataset-run", "dataset-run", dao.EventTypeTerminal, "Knowledge compilation completed")
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", 1, 30, nil, nil, "", "", "dataset", "", "")
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", 1, 30, nil, nil, "", "", "dataset", "", "")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListIngestionLogs = (%+v, %v, %v), want success", result, code, err)
}
@@ -99,7 +99,7 @@ func TestDatasetIngestionLogUsesEventStream(t *testing.T) {
}
assertLatestEventMap(t, logs[0], 1, "Knowledge compilation completed")
messages, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "dataset-run", 200, nil, nil)
messages, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "dataset-run", 200, nil, nil)
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListIngestionMessages = (%+v, %v, %v), want success", messages, code, err)
}
@@ -107,7 +107,7 @@ func TestDatasetIngestionLogUsesEventStream(t *testing.T) {
t.Fatalf("dataset messages = %+v, want terminal event stream", messages)
}
log, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "dataset-run")
log, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "dataset-run")
if err != nil || code != common.CodeSuccess {
t.Fatalf("GetIngestionLog = (%+v, %v, %v), want success", log, code, err)
}
@@ -160,7 +160,7 @@ func TestListSyncLogsReturnsOnlyCurrentlyLinkedDatasetConnectors(t *testing.T) {
permission.OperationUse, permission.OperationRead); err != nil {
t.Fatalf("check linked sync log access fixture: %v", err)
}
result, code, err := NewDatasetService().ListSyncLogs(ctx, "kb-1", "user-1", 1, 30)
result, code, err := NewDatasetService().ListSyncLogs(ctx, "kb-1", 1, 30)
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListSyncLogs = (%+v, %v, %v), want success", result, code, err)
}
@@ -192,7 +192,7 @@ func TestListIngestionLogsIncludesPythonFileLogsAndStatusFilters(t *testing.T) {
t.Fatalf("insert Python run: %v", err)
}
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", 1, 30, nil, []string{"DONE"}, "", "", "file", "", "")
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", 1, 30, nil, []string{"DONE"}, "", "", "file", "", "")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListIngestionLogs = (%+v, %v, %v), want success", result, code, err)
}
@@ -214,7 +214,7 @@ func TestListIngestionLogsIncludesPythonFileLogsAndStatusFilters(t *testing.T) {
t.Fatalf("Python file log = %#v, want file-log fields", byID["python-run"])
}
detail, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "python-run")
detail, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "python-run")
if err != nil || code != common.CodeSuccess {
t.Fatalf("GetIngestionLog = (%+v, %v, %v), want success", detail, code, err)
}
@@ -225,7 +225,7 @@ func TestListIngestionLogsIncludesPythonFileLogsAndStatusFilters(t *testing.T) {
t.Fatalf("Python detail = %#v, want file-log fields", detail)
}
messages, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "python-run", 200, nil, nil)
messages, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "python-run", 200, nil, nil)
if err != nil || code != common.CodeSuccess || !messages.Terminal || len(messages.Items) != 0 {
t.Fatalf("Python ListIngestionMessages = (%+v, %v, %v), want empty terminal success", messages, code, err)
}
@@ -247,7 +247,7 @@ func TestPythonDatasetLogIsReadableAndKeepsProgressMessage(t *testing.T) {
t.Fatalf("insert Python dataset run: %v", err)
}
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", 1, 30, nil, []string{"DONE"}, "", "", "dataset", "", "")
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", 1, 30, nil, []string{"DONE"}, "", "", "dataset", "", "")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListIngestionLogs = (%+v, %v, %v), want success", result, code, err)
}
@@ -259,7 +259,7 @@ func TestPythonDatasetLogIsReadableAndKeepsProgressMessage(t *testing.T) {
t.Fatalf("Python dataset log fields = %#v, want stored status and progress_msg", logs[0])
}
log, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "python-dataset-run")
log, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "python-dataset-run")
if err != nil || code != common.CodeSuccess {
t.Fatalf("GetIngestionLog = (%+v, %v, %v), want success", log, code, err)
}
@@ -267,7 +267,7 @@ func TestPythonDatasetLogIsReadableAndKeepsProgressMessage(t *testing.T) {
t.Fatalf("detail progress_msg = %#v, want %q", log["progress_msg"], progressMessage)
}
messages, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "python-dataset-run", 200, nil, nil)
messages, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "python-dataset-run", 200, nil, nil)
if err != nil || code != common.CodeSuccess || !messages.Terminal || len(messages.Items) != 0 {
t.Fatalf("legacy ListIngestionMessages = (%+v, %v, %v), want empty terminal success", messages, code, err)
}
@@ -289,7 +289,7 @@ func TestListIngestionLogsEmbedsLatestRunEventWithoutProgressMessage(t *testing.
insertMessageEvent(t, db, "run-1", "task-1", dao.EventTypeTerminal, "latest first")
insertMessageEvent(t, db, "run-2", "task-2", dao.EventTypeMessage, "latest second")
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", 1, 30, nil, nil, "", "", "file", "", "")
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", 1, 30, nil, nil, "", "", "file", "", "")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListIngestionLogs = (%+v, %v, %v), want success", result, code, err)
}
@@ -320,7 +320,7 @@ func TestGetIngestionLogEmbedsLatestRunEvent(t *testing.T) {
insertMessageRun(t, "run-1", "kb-1", "doc-1", entity.TaskStatusDone, 1)
insertMessageEvent(t, db, "run-1", "task-1", dao.EventTypeMessage, "latest detail")
result, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "run-1")
result, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "run-1")
if err != nil || code != common.CodeSuccess {
t.Fatalf("GetIngestionLog = (%+v, %v, %v), want success", result, code, err)
}
@@ -344,7 +344,7 @@ func TestIngestionEventWriterFeedsRunScopedLogReaders(t *testing.T) {
t.Fatalf("record progress event: %v", err)
}
messages, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "run-1", 200, nil, nil)
messages, code, err := NewDatasetService().ListIngestionMessages(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "run-1", 200, nil, nil)
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListIngestionMessages = (%+v, %v, %v), want success", messages, code, err)
}
@@ -352,7 +352,7 @@ func TestIngestionEventWriterFeedsRunScopedLogReaders(t *testing.T) {
t.Fatalf("messages = %+v, want lifecycle then latest progress", messages.Items)
}
log, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", "run-1")
log, code, err := NewDatasetService().GetIngestionLog(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "run-1")
if err != nil || code != common.CodeSuccess {
t.Fatalf("GetIngestionLog = (%+v, %v, %v), want success", log, code, err)
}
@@ -446,7 +446,7 @@ func TestListIngestionLogsRejectsUnknownLogType(t *testing.T) {
{name: "empty", logType: "", wantCode: common.CodeDataError},
} {
t.Run(tc.name, func(t *testing.T) {
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", "user-1", 1, 30, nil, nil, "", "", tc.logType, "", "")
result, code, err := NewDatasetService().ListIngestionLogs(withDatasetTestContext(t.Context(), "user-1", "user-1"), "kb-1", 1, 30, nil, nil, "", "", tc.logType, "", "")
if tc.wantCode == common.CodeSuccess {
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListIngestionLogs(%q) = (%+v, %v, %v), want success", tc.logType, result, code, err)
internal/service/dataset/list_test.go
+12 -12
View File
@@ -29,7 +29,7 @@ func TestDatasetServiceListDatasetsFiltersByIDs(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
data, total, code, err := testDatasetListService(t).ListDatasets(ctx,
"", "", 1, 30, []dao.OrderTerm{{Column: "create_time", Desc: true}},
"", nil, "", "tenant-1", []string{"kb-1"}, "",
"", nil, "", []string{"kb-1"}, "",
)
if err != nil {
t.Fatalf("ListDatasets failed: %v", err)
@@ -69,7 +69,7 @@ func TestDatasetServiceListDatasetsIDsAccessibleViaTeamTenant(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "user-1", "owner-1")
data, total, code, err := testDatasetListService(t).ListDatasets(ctx,
"", "", 1, 30, []dao.OrderTerm{{Column: "create_time", Desc: true}},
"", nil, "", "user-1", []string{"kb-team"}, "",
"", nil, "", []string{"kb-team"}, "",
)
if err != nil {
t.Fatalf("ListDatasets failed: %v", err)
@@ -98,7 +98,7 @@ func TestDatasetServiceListDatasetsAndFiltersUseAuthorizedScope(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "user-1", "owner-1")
data, total, code, err := testDatasetListService(t).ListDatasets(ctx,
"", "", 1, 30, []dao.OrderTerm{{Column: "create_time", Desc: true}},
"", []string{"owner-1"}, "", "user-1", nil, "",
"", []string{"owner-1"}, "", nil, "",
)
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListDatasets = (%#v, %d, %v), want success", data, code, err)
@@ -107,7 +107,7 @@ func TestDatasetServiceListDatasetsAndFiltersUseAuthorizedScope(t *testing.T) {
t.Fatalf("ListDatasets returned %#v (total=%d), want only the shared dataset", data, total)
}
filters, code, err := testDatasetListService(t).ListDatasetFilters(ctx, "user-1", "")
filters, code, err := testDatasetListService(t).ListDatasetFilters(ctx, "")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListDatasetFilters = (%#v, %v, %v), want success", filters, code, err)
}
@@ -146,7 +146,7 @@ func TestDatasetServiceListDatasetsForCanvasUsesCallerOwnerIntersection(t *testi
service := testDatasetListService(t)
data, total, code, err := service.ListDatasets(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "", "", 1, 30,
[]dao.OrderTerm{{Column: "create_time", Desc: true}}, "", nil, "", "user-1", nil, "canvas-shared")
[]dao.OrderTerm{{Column: "create_time", Desc: true}}, "", nil, "", nil, "canvas-shared")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListDatasets(canvas) = (%#v, %d, %v), want success", data, code, err)
}
@@ -155,7 +155,7 @@ func TestDatasetServiceListDatasetsForCanvasUsesCallerOwnerIntersection(t *testi
}
data, total, code, err = service.ListDatasets(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "", "", 1, 30,
[]dao.OrderTerm{{Column: "create_time", Desc: true}}, "", nil, "", "user-1",
[]dao.OrderTerm{{Column: "create_time", Desc: true}}, "", nil, "",
[]string{"kb-caller-only", "kb-shared", "kb-owner-only"}, "canvas-shared")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListDatasets(canvas, ids) = (%#v, %d, %v), want success", data, code, err)
@@ -188,7 +188,7 @@ func TestDatasetServiceListDatasetsPaginatesInsidePermissionScope(t *testing.T)
service := testDatasetListService(t)
terms := []dao.OrderTerm{{Column: "name"}}
pageOne, total, code, err := service.ListDatasets(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "", "", 1, 1, terms, "", nil, "", "user-1", nil, "")
pageOne, total, code, err := service.ListDatasets(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "", "", 1, 1, terms, "", nil, "", nil, "")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListDatasets(page 1) = (%#v, %d, %v), want success", pageOne, code, err)
}
@@ -196,7 +196,7 @@ func TestDatasetServiceListDatasetsPaginatesInsidePermissionScope(t *testing.T)
t.Fatalf("page 1 = %#v (total=%d), want kb-own and two accessible datasets", pageOne, total)
}
pageTwo, secondTotal, code, err := service.ListDatasets(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "", "", 2, 1, terms, "", nil, "", "user-1", nil, "")
pageTwo, secondTotal, code, err := service.ListDatasets(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "", "", 2, 1, terms, "", nil, "", nil, "")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListDatasets(page 2) = (%#v, %d, %v), want success", pageTwo, code, err)
}
@@ -210,7 +210,7 @@ func TestDatasetServiceListDatasetFiltersReturnsEmptyOwnerList(t *testing.T) {
pushServiceDB(t, db)
ensureDatasetTestMembership(t, "user-without-datasets", "user-without-datasets", "owner")
filters, code, err := testDatasetListService(t).ListDatasetFilters(withDatasetTestContext(t.Context(), "user-without-datasets", "user-without-datasets"), "user-without-datasets", "")
filters, code, err := testDatasetListService(t).ListDatasetFilters(withDatasetTestContext(t.Context(), "user-without-datasets", "user-without-datasets"), "")
if err != nil || code != common.CodeSuccess {
t.Fatalf("ListDatasetFilters = (%#v, %v, %v), want success", filters, code, err)
}
@@ -245,7 +245,7 @@ func TestDatasetServiceListDatasetsRejectsIDAndIDsTogether(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "user-1", "tenant-1")
_, _, code, err := testDatasetListService(t).ListDatasets(ctx,
"kb-1", "", 1, 30, []dao.OrderTerm{{Column: "create_time", Desc: true}},
"", nil, "", "tenant-1", []string{"kb-1"}, "",
"", nil, "", []string{"kb-1"}, "",
)
if err == nil {
t.Fatal("expected id/ids conflict error")
@@ -268,7 +268,7 @@ func TestDatasetServiceListDatasetsFiltersInaccessibleIDs(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "user-1", "user-1")
data, total, code, err := testDatasetListService(t).ListDatasets(ctx,
"", "", 1, 30, []dao.OrderTerm{{Column: "create_time", Desc: true}},
"", nil, "", "user-1", []string{"kb-valid", "kb-private", "kb-missing"}, "",
"", nil, "", []string{"kb-valid", "kb-private", "kb-missing"}, "",
)
if err != nil {
t.Fatalf("ListDatasets failed: %v", err)
@@ -289,7 +289,7 @@ func TestDatasetServiceListDatasetsReturnsEmptyForAllInaccessibleIDs(t *testing.
ctx := withDatasetTestContext(t.Context(), "user-1", "user-1")
data, total, code, err := testDatasetListService(t).ListDatasets(ctx,
"", "", 1, 30, []dao.OrderTerm{{Column: "create_time", Desc: true}},
"", nil, "", "user-1", []string{"kb-private", "kb-missing"}, "",
"", nil, "", []string{"kb-private", "kb-missing"}, "",
)
if err != nil {
t.Fatalf("ListDatasets failed: %v", err)
internal/service/dataset/metadata.go
+5 -6
View File
@@ -15,8 +15,8 @@ import (
)
// UpdateDocumentMetadataConfig updates the metadata config for a document in a dataset.
func (d *DatasetService) UpdateDocumentMetadataConfig(ctx context.Context, userID, datasetID, documentID string, req map[string]interface{}) (*entity.Document, common.ErrorCode, error) {
userID = strings.TrimSpace(userID)
func (d *DatasetService) UpdateDocumentMetadataConfig(ctx context.Context, datasetID, documentID string, req map[string]interface{}) (*entity.Document, common.ErrorCode, error) {
userID := strings.TrimSpace(permission.UserIDFromContext(ctx))
datasetID = strings.TrimSpace(datasetID)
if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.ResourceRef{Kind: permission.ResourceKindDataset, ID: datasetID}, permission.OperationUpdate); err != nil {
code, permissionErr := permissionresponse.Normalize(err)
@@ -69,9 +69,8 @@ func (d *DatasetService) UpdateDocumentMetadataConfig(ctx context.Context, userI
}
// GetMetadataConfig gets the auto-metadata configuration for a dataset.
func (d *DatasetService) GetMetadataConfig(ctx context.Context, datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) GetMetadataConfig(ctx context.Context, datasetID string) (map[string]interface{}, common.ErrorCode, error) {
datasetID = strings.TrimSpace(datasetID)
userID = strings.TrimSpace(userID)
if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.ResourceRef{Kind: permission.ResourceKindDataset, ID: datasetID}, permission.OperationRead); err != nil {
code, permissionErr := permissionresponse.Normalize(err)
return nil, code, permissionErr
@@ -100,9 +99,9 @@ func (d *DatasetService) GetMetadataConfig(ctx context.Context, datasetID, userI
}
// UpdateMetadataConfig updates the auto-metadata configuration for a dataset.
func (d *DatasetService) UpdateMetadataConfig(ctx context.Context, datasetID, userID string, req *service.MetadataConfigRequest) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) UpdateMetadataConfig(ctx context.Context, datasetID string, req *service.MetadataConfigRequest) (map[string]interface{}, common.ErrorCode, error) {
datasetID = strings.TrimSpace(datasetID)
userID = strings.TrimSpace(userID)
userID := strings.TrimSpace(permission.UserIDFromContext(ctx))
if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.ResourceRef{Kind: permission.ResourceKindDataset, ID: datasetID}, permission.OperationUpdate); err != nil {
code, permissionErr := permissionresponse.Normalize(err)
internal/service/dataset/metadata_config_test.go
+13 -13
View File
@@ -110,7 +110,7 @@ func TestDatasetServiceUpdateDocumentMetadataConfig(t *testing.T) {
metadata := map[string]interface{}{"author": "Alice", "year": float64(2026)}
doc, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig(
ctx,
"user-1",
"kb-1",
"doc-1",
map[string]interface{}{"metadata": metadata},
@@ -162,7 +162,7 @@ func TestDatasetServiceUpdateDocumentMetadataConfigRequiresMetadata(t *testing.T
ctx := withDatasetTestContext(t.Context(), "user-1", "user-1")
_, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig(
ctx,
"user-1",
"kb-1",
"doc-1",
map[string]interface{}{},
@@ -188,7 +188,7 @@ func TestDatasetServiceUpdateDocumentMetadataConfigRejectsNonOwner(t *testing.T)
ctx := withDatasetTestContext(t.Context(), "user-1", "owner-1")
_, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig(
ctx,
"user-1",
"kb-1",
"doc-1",
map[string]interface{}{"metadata": map[string]interface{}{"author": "Alice"}},
@@ -221,7 +221,7 @@ func TestDatasetServiceUpdateDocumentMetadataConfigAllowsTeamMember(t *testing.T
ctx := withDatasetTestContext(t.Context(), "user-1", "owner-1")
doc, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig(
ctx,
"user-1",
"kb-1",
"doc-1",
map[string]interface{}{"metadata": map[string]interface{}{"author": "Alice"}},
@@ -267,7 +267,7 @@ func TestDatasetServiceUpdateMetadataConfigSyncsExtractorSchema(t *testing.T) {
result, code, err := (&DatasetService{
kbDAO: dao.NewKnowledgebaseDAO(),
tenantDAO: dao.NewTenantDAO(),
}).UpdateMetadataConfig(ctx, "kb-1", "tenant-1", &service.MetadataConfigRequest{
}).UpdateMetadataConfig(ctx, "kb-1", &service.MetadataConfigRequest{
Metadata: []service.MetadataConfigField{
{Key: "author", Type: "string"},
},
@@ -321,7 +321,7 @@ func TestDatasetServiceUpdateMetadataConfigSyncsExtractorSchema(t *testing.T) {
_, code, err = (&DatasetService{
kbDAO: dao.NewKnowledgebaseDAO(),
tenantDAO: dao.NewTenantDAO(),
}).UpdateMetadataConfig(ctx, "kb-1", "tenant-1", &service.MetadataConfigRequest{
}).UpdateMetadataConfig(ctx, "kb-1", &service.MetadataConfigRequest{
Metadata: []service.MetadataConfigField{
{Key: "author", Type: "string"},
},
@@ -370,7 +370,7 @@ func TestDatasetServiceUpdateMetadataConfigSyncsExtractorSchema(t *testing.T) {
_, code, err = (&DatasetService{
kbDAO: dao.NewKnowledgebaseDAO(),
tenantDAO: dao.NewTenantDAO(),
}).UpdateMetadataConfig(ctx, "kb-1", "tenant-1", &service.MetadataConfigRequest{
}).UpdateMetadataConfig(ctx, "kb-1", &service.MetadataConfigRequest{
Metadata: []service.MetadataConfigField{},
BuiltInMetadata: []service.MetadataConfigField{},
})
@@ -413,7 +413,7 @@ func TestDatasetServiceGetMetadataConfig(t *testing.T) {
_, code, err := (&DatasetService{
kbDAO: dao.NewKnowledgebaseDAO(),
tenantDAO: dao.NewTenantDAO(),
}).UpdateMetadataConfig(ctx, "kb-1", "tenant-1", &service.MetadataConfigRequest{
}).UpdateMetadataConfig(ctx, "kb-1", &service.MetadataConfigRequest{
Metadata: []service.MetadataConfigField{
{Key: "author", Type: "string"},
},
@@ -431,7 +431,7 @@ func TestDatasetServiceGetMetadataConfig(t *testing.T) {
result, code, err := (&DatasetService{
kbDAO: dao.NewKnowledgebaseDAO(),
tenantDAO: dao.NewTenantDAO(),
}).GetMetadataConfig(ctx, "kb-1", "tenant-1")
}).GetMetadataConfig(ctx, "kb-1")
if err != nil {
t.Fatalf("GetMetadataConfig failed: %v", err)
}
@@ -466,7 +466,7 @@ func TestUpdateMetadataConfig_TenantLookupFailureDoesNotPersist(t *testing.T) {
result, code, err := (&DatasetService{
kbDAO: dao.NewKnowledgebaseDAO(),
tenantDAO: dao.NewTenantDAO(),
}).UpdateMetadataConfig(ctx, "kb-1", "missing-tenant", &service.MetadataConfigRequest{
}).UpdateMetadataConfig(ctx, "kb-1", &service.MetadataConfigRequest{
Metadata: []service.MetadataConfigField{
{Key: "author", Type: "string"},
},
@@ -496,7 +496,7 @@ func TestUpdateMetadataConfig_ExplicitEnabledPreservedWhenEmpty(t *testing.T) {
result, code, err := (&DatasetService{
kbDAO: dao.NewKnowledgebaseDAO(),
tenantDAO: dao.NewTenantDAO(),
}).UpdateMetadataConfig(ctx, "kb-1", "tenant-1", &service.MetadataConfigRequest{
}).UpdateMetadataConfig(ctx, "kb-1", &service.MetadataConfigRequest{
Enabled: &enabled,
Metadata: []service.MetadataConfigField{},
BuiltInMetadata: []service.MetadataConfigField{},
@@ -535,7 +535,7 @@ func TestUpdateDocumentMetadataConfig_KBAndTenantLookupFailure(t *testing.T) {
// Missing KB -> database layer should surface as CodeServerError / not found
_, code, err := testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig(
ctx, "owner-1", "kb-missing", "doc-ok",
ctx, "kb-missing", "doc-ok",
map[string]interface{}{"metadata": map[string]interface{}{"a": "b"}},
)
if code == common.CodeSuccess || err == nil {
@@ -547,7 +547,7 @@ func TestUpdateDocumentMetadataConfig_KBAndTenantLookupFailure(t *testing.T) {
insertDatasetMetadataConfigDoc(t, "doc2", "kb-bad-tenant", entity.JSONMap{})
missingTenantCtx := withDatasetTestContext(t.Context(), "missing-tenant", "missing-tenant")
_, code, err = testDatasetServiceForDocumentMetadataConfig(t).UpdateDocumentMetadataConfig(
missingTenantCtx, "missing-tenant", "kb-bad-tenant", "doc2",
missingTenantCtx, "kb-bad-tenant", "doc2",
map[string]interface{}{"metadata": map[string]interface{}{"a": "b"}},
)
if code != common.CodeServerError || err == nil {
internal/service/dataset/pipeline_version_binding_test.go
+7 -7
View File
@@ -156,7 +156,7 @@ func TestUpdateDataset_RewritesVersionOnPipelineChange(t *testing.T) {
// C5/G3: switching the bound pipeline now requires an explicit version (the
// empty version is no longer auto-backfilled to the new canvas latest). Supply
// canvas2's version to verify the rewrite still pins the chosen version.
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-change", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-change", service.UpdateDatasetRequest{
PipelineID: &canvas2,
ParseType: &parseType,
PipelineVersionID: &ver2,
@@ -185,7 +185,7 @@ func TestUpdateDataset_PreservesVersionWhenPipelineUnchanged(t *testing.T) {
parseType := 2
name := "kb-preserve-renamed"
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-preserve", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-preserve", service.UpdateDatasetRequest{
Name: &name,
PipelineID: &canvas1,
ParseType: &parseType,
@@ -351,7 +351,7 @@ func TestUpdateDataset_RebindsExplicitVersion(t *testing.T) {
seedBoundKB(t, db, "kb-repin", canvas1, verOld)
parseType := 2
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-repin", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-repin", service.UpdateDatasetRequest{
PipelineID: &canvas1,
PipelineVersionID: &verNew,
ParseType: &parseType,
@@ -411,7 +411,7 @@ func TestUpdateDataset_EmptyVersionClearsPin(t *testing.T) {
parseType := 2
empty := ""
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-clear", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-clear", service.UpdateDatasetRequest{
PipelineID: &canvas1,
ParseType: &parseType,
PipelineVersionID: &empty,
@@ -456,7 +456,7 @@ func TestUpdateDataset_VersionChangeResetsParserConfigToNewVersionDefault(t *tes
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
parseType := 2
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-c3", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-c3", service.UpdateDatasetRequest{
PipelineID: &canvas1,
PipelineVersionID: &verNew,
ParseType: &parseType,
@@ -528,7 +528,7 @@ func TestUpdateDataset_VersionChangeKeepsSubmittedParserConfig(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
parseType := 2
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-c3-echo", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-c3-echo", service.UpdateDatasetRequest{
PipelineID: &canvas1,
PipelineVersionID: &verNew,
ParseType: &parseType,
@@ -585,7 +585,7 @@ func TestUpdateDataset_SameVersionKeepsParserConfig(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
parseType := 2
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-c3-same", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-c3-same", service.UpdateDatasetRequest{
PipelineID: &canvas1,
PipelineVersionID: &verOld, // same version, no change
ParseType: &parseType,
internal/service/dataset/sdk_contract_test.go
+2 -2
View File
@@ -39,7 +39,7 @@ func TestSDKDatasetUnicodeCharacterLimit(t *testing.T) {
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{Name: new(strings.Repeat("资", n))})
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{Name: new(strings.Repeat("资", n))})
if (err == nil) != (n == 128) {
t.Fatalf("update characters=%d code=%d err=%v", n, code, err)
}
@@ -50,7 +50,7 @@ func TestSDKDatasetNegativePagerankDiagnostic(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1", service.UpdateDatasetRequest{Pagerank: new(int64(-1))})
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{Pagerank: new(int64(-1))})
if code != common.CodeDataError || err == nil || !strings.Contains(err.Error(), "greater than or equal to 0") {
t.Fatalf("code=%d err=%v", code, err)
}
internal/service/dataset/search.go
+4 -3
View File
@@ -17,14 +17,14 @@ import (
"ragflow/internal/service/nlp"
)
func (d *DatasetService) SearchDataset(ctx context.Context, datasetID, userID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) {
func (d *DatasetService) SearchDataset(ctx context.Context, datasetID string, req *service.SearchDatasetRequest) (*service.SearchDatasetsResponse, error) {
if datasetID == "" {
return nil, fmt.Errorf("dataset_id is required")
}
return d.SearchDatasets(ctx, req.ToSearchDatasetsRequest(datasetID), userID)
return d.SearchDatasets(ctx, req.ToSearchDatasetsRequest(datasetID))
}
func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest, userID string) (*service.SearchDatasetsResponse, error) {
func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.SearchDatasetsRequest) (*service.SearchDatasetsResponse, error) {
if req.Question == "" {
return nil, fmt.Errorf("question is required")
}
@@ -34,6 +34,7 @@ func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.Search
if err := service.ValidateKnowledgeCompilationScope(req.KnowledgeCompilationScope); err != nil {
return nil, err
}
userID := permission.UserIDFromContext(ctx)
common.Info("SearchDatasets started", zap.String("userID", userID), zap.Any("datasets", req.DatasetIDs), zap.String("question", req.Question))
page := 1
internal/service/dataset/search_integration_test.go
+2 -2
View File
@@ -119,7 +119,7 @@ func TestSearchDatasetsEmbeddingConsistencyIntegration(t *testing.T) {
}
req := f.request(0.7, 0)
req.DatasetIDs = []string{"kb-a", "kb-b"}
result, err := f.service.SearchDatasets(f.ctx, req, f.owner)
result, err := f.service.SearchDatasets(f.ctx, req)
if err == nil || !strings.Contains(err.Error(), "datasets use different embedding models") || result != nil {
t.Fatalf("different base models must be rejected before retrieval: result=%+v err=%v", result, err)
}
@@ -280,7 +280,7 @@ func (f *searchIntegrationFixture) request(weight, threshold float64) *service.S
func (f *searchIntegrationFixture) search(t *testing.T, req *service.SearchDatasetsRequest) *service.SearchDatasetsResponse {
t.Helper()
result, err := f.service.SearchDatasets(f.ctx, req, f.owner)
result, err := f.service.SearchDatasets(f.ctx, req)
if err != nil {
t.Fatalf("SearchDatasets: %v", err)
}
internal/service/dataset/tags.go
+1 -1
View File
@@ -14,7 +14,7 @@ import (
permissionresponse "ragflow/internal/permission/response"
)
func (d *DatasetService) AggregateTags(ctx context.Context, datasetIDs []string, userID string) ([]map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) AggregateTags(ctx context.Context, datasetIDs []string) ([]map[string]interface{}, common.ErrorCode, error) {
if len(datasetIDs) == 0 {
return nil, common.CodeDataError, errors.New("Lack of dataset_ids in query parameters")
}
internal/service/dataset/tags_aggregate_test.go
+13 -13
View File
@@ -105,7 +105,7 @@ func TestDatasetServiceAggregateTagsSeedsFromTagFile(t *testing.T) {
}
ctx := withDatasetTestContext(t.Context(), "user-1", "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(ctx, []string{kbInput}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(ctx, []string{kbInput})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -151,7 +151,7 @@ func TestDatasetServiceAggregateTagsMergesAcrossTagFiles(t *testing.T) {
}
ctx := withDatasetTestContext(t.Context(), "user-1", "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(ctx, []string{kb1Input, kb2Input}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(ctx, []string{kb1Input, kb2Input})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -188,7 +188,7 @@ func TestDatasetServiceAggregateTagsEmptyWhenNoTagFile(t *testing.T) {
}
ctx := withDatasetTestContext(t.Context(), "user-1", "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(ctx, []string{kbInput}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(ctx, []string{kbInput})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -209,7 +209,7 @@ func TestDatasetServiceAggregateTagsRejectsUnauthorizedDataset(t *testing.T) {
insertAggregateTagsKB(t, kbID, "tenant-9", string(entity.TenantPermissionMe), "", 0)
ctx := withDatasetTestContext(t.Context(), "user-1", "tenant-9")
_, code, err := testDatasetServiceForAggregateTags(t, nil).AggregateTags(ctx, []string{kbInput}, "user-1")
_, code, err := testDatasetServiceForAggregateTags(t, nil).AggregateTags(ctx, []string{kbInput})
if err == nil {
t.Fatal("expected authorization error")
}
@@ -246,7 +246,7 @@ func TestDatasetServiceAggregateTagsScopesLoaderToDatasetTenant(t *testing.T) {
return map[string]int{"finance": 1}, nil
}
if _, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(ctx, []string{kbInput}, "user-1"); err != nil {
if _, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(ctx, []string{kbInput}); err != nil {
t.Fatalf("AggregateTags failed: code=%d err=%v", code, err)
}
if gotOwner != kbTenant {
@@ -292,7 +292,7 @@ func TestDatasetServiceAggregateTagsReadsDocumentTagFile(t *testing.T) {
return nil, nil
}
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -324,7 +324,7 @@ func TestDatasetServiceAggregateTagsCountsSnapshotOnce(t *testing.T) {
return nil, nil
}
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -360,7 +360,7 @@ func TestDatasetServiceAggregateTagsUnionsDocumentTagFiles(t *testing.T) {
return nil, nil
}
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -396,7 +396,7 @@ func TestDatasetServiceAggregateTagsCountsSharedFileOnce(t *testing.T) {
return nil, nil
}
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kb1Input, kb2Input}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kb1Input, kb2Input})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -439,7 +439,7 @@ func TestDatasetServiceAggregateTagsSkipsUnresolvableDocumentTagSource(t *testin
return map[string]int{"finance": 2}, nil
}
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{deadKB, liveKB}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{deadKB, liveKB})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -467,7 +467,7 @@ func TestDatasetServiceAggregateTagsIgnoresUnappliedDatasetTagSource(t *testing.
return nil, nil
}
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput}, "user-1")
result, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput})
if err != nil {
t.Fatalf("AggregateTags failed: %v", err)
}
@@ -495,7 +495,7 @@ func TestDatasetServiceAggregateTagsFailsOnStorageErrorForDocumentTagSource(t *t
return nil, errors.New("tag source file: no storage backend registered")
}
_, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput}, "user-1")
_, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "user-1"), []string{kbInput})
if err == nil {
t.Fatal("expected an error for a storage failure on a document-level tag source")
}
@@ -524,7 +524,7 @@ func TestDatasetServiceAggregateTagsScopesDocumentTagSourceToDatasetTenant(t *te
return map[string]int{"finance": 1}, nil
}
if _, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "tenant-owner"), []string{kbInput}, "user-1"); err != nil {
if _, code, err := testDatasetServiceForAggregateTags(t, loader).AggregateTags(withDatasetTestContext(t.Context(), "user-1", "tenant-owner"), []string{kbInput}); err != nil {
t.Fatalf("AggregateTags failed: code=%d err=%v", code, err)
}
if gotOwner != kbTenant {
internal/service/dataset/update.go
+4 -4
View File
@@ -31,9 +31,9 @@ type datasetPagerankUpdate struct {
datasetID string
}
func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, userID string, req service.UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) {
func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID string, req service.UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) {
datasetID = strings.TrimSpace(datasetID)
userID = strings.TrimSpace(userID)
userID := strings.TrimSpace(permission.UserIDFromContext(ctx))
if _, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID); err != nil {
if dao.IsNotFoundErr(err) {
code, permissionErr := permissionresponse.NormalizeHidden(permission.ErrResourceNotFound)
@@ -184,7 +184,7 @@ func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, userID st
var scheduledTaskIDs []string
var pagerankUpdate *datasetPagerankUpdate
err = dao.DB.Transaction(func(tx *gorm.DB) error {
lockedKB, code, authErr := d.lockAccessibleDatasetForUpdate(ctx, tx, datasetID, userID)
lockedKB, code, authErr := d.lockAccessibleDatasetForUpdate(ctx, tx, datasetID)
if authErr != nil {
txCode = code
return authErr
@@ -547,7 +547,7 @@ func publishDatasetSyncerTasks(taskIDs []string) {
}
}
func (d *DatasetService) lockAccessibleDatasetForUpdate(ctx context.Context, tx *gorm.DB, datasetID, userID string) (*entity.Knowledgebase, common.ErrorCode, error) {
func (d *DatasetService) lockAccessibleDatasetForUpdate(ctx context.Context, tx *gorm.DB, datasetID string) (*entity.Knowledgebase, common.ErrorCode, error) {
var kb entity.Knowledgebase
err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
Where("id = ? AND status = ?", datasetID, string(entity.StatusValid)).
internal/service/dataset/update_test.go
+39 -39
View File
@@ -45,7 +45,7 @@ func TestDatasetServiceUpdateDatasetUpdatesFields(t *testing.T) {
parseType := 1
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Name: &name,
Description: &description,
Language: &language,
@@ -105,12 +105,12 @@ func TestUpdateDataset_ClearsAvatarWithEmptyString(t *testing.T) {
svc := testDatasetUpdateService(t)
initial := "data:image/png;base64,iVBORw0KGgo="
if _, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1", service.UpdateDatasetRequest{Avatar: &initial}); err != nil || code != common.CodeSuccess {
if _, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{Avatar: &initial}); err != nil || code != common.CodeSuccess {
t.Fatalf("set avatar failed: code=%d err=%v", code, err)
}
empty := ""
result, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1", service.UpdateDatasetRequest{Avatar: &empty})
result, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{Avatar: &empty})
if err != nil {
t.Fatalf("clear avatar failed: %v", err)
}
@@ -138,7 +138,7 @@ func TestUpdateDataset_ParentChildConfigReachesGeneralChunker(t *testing.T) {
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"GeneralChunker:SixApplesFall": map[string]interface{}{
"parent_child": map[string]interface{}{
@@ -188,7 +188,7 @@ func TestUpdateDatasetPreservesParentChildChunkerRuntimeConfig(t *testing.T) {
t.Fatalf("seed parent-child parser config: %v", err)
}
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
// Component-scoped override (flat parser-level keys are rejected now).
ParserConfig: map[string]interface{}{
"GeneralChunker:SixApplesFall": map[string]interface{}{"chunk_token_size": float64(256)},
@@ -220,7 +220,7 @@ func TestUpdateDataset_RejectsSimultaneousParserIDAndPipelineID(t *testing.T) {
pipelineID := "abcdef0123456789abcdef0123456789"
ctx := t.Context()
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
PipelineID: &pipelineID,
})
@@ -246,7 +246,7 @@ func TestUpdateDataset_ParseTypeBuiltinClearsPipelineID(t *testing.T) {
parseType := 1
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
ParseType: &parseType,
})
@@ -278,7 +278,7 @@ func TestUpdateDataset_ParseTypePipelineRejectsParserID(t *testing.T) {
pipelineID := "ABCDEF0123456789ABCDEF0123456789"
parseType := 2
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
PipelineID: &pipelineID,
ParseType: &parseType,
@@ -323,7 +323,7 @@ func TestUpdateDataset_PipelineSwitchKeepsSubmittedParserConfig(t *testing.T) {
seedCanvasVersion(t, db, pipelineID, verID, 0)
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
PipelineID: &pipelineID,
ParseType: &parseType,
PipelineVersionID: &verID,
@@ -385,7 +385,7 @@ func TestUpdateDataset_ParseTypePipelineCleansConfigAgainstCanvas(t *testing.T)
}
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
PipelineID: &pipelineID,
ParseType: &parseType,
ParserConfig: override,
@@ -428,7 +428,7 @@ func TestUpdateDatasetRejectsInvalidPages(t *testing.T) {
chunkMethod := "manual"
parseType := 1
ctx := t.Context()
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
ParseType: &parseType,
ParserConfig: map[string]interface{}{
@@ -454,7 +454,7 @@ func TestDatasetServiceGetDatasetReturnsEmptyConnectorList(t *testing.T) {
insertDatasetUpdateKB(t, datasetID, "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).GetDataset(ctx, "11111111-1111-4111-8111-111111111111", "tenant-1")
result, code, err := testDatasetUpdateService(t).GetDataset(ctx, "11111111-1111-4111-8111-111111111111")
if err != nil {
t.Fatalf("GetDataset failed: %v", err)
}
@@ -479,7 +479,7 @@ func TestDatasetServiceUpdateDatasetRejectsMissingDataset(t *testing.T) {
name := "Renamed"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "missing-kb", "tenant-1", service.UpdateDatasetRequest{Name: &name})
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "missing-kb", service.UpdateDatasetRequest{Name: &name})
if err == nil {
t.Fatal("expected missing dataset error")
}
@@ -499,7 +499,7 @@ func TestDatasetServiceUpdateDatasetRejectsNonOwner(t *testing.T) {
name := "Renamed"
ctx := withDatasetTestContext(t.Context(), "tenant-2", "tenant-2")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-2", service.UpdateDatasetRequest{Name: &name})
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{Name: &name})
if err == nil {
t.Fatal("expected permission error")
}
@@ -524,7 +524,7 @@ func TestDatasetServiceUpdateDatasetRejectsTeamMemberPermissionChange(t *testing
permission := string(entity.TenantPermissionMe)
ctx := withDatasetTestContext(t.Context(), "user-1", "owner-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "user-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Permission: &permission,
})
if err == nil {
@@ -557,7 +557,7 @@ func TestDatasetServiceUpdateDatasetAllowsTeamMemberToResubmitUnchangedPermissio
name := "Renamed by team member"
permission := string(entity.TenantPermissionTeam)
result, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "kb-1", "user-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "kb-1", service.UpdateDatasetRequest{
Name: &name,
Permission: &permission,
})
@@ -601,7 +601,7 @@ func TestDatasetServiceUpdateDatasetResolvesEmbeddingInDatasetTenant(t *testing.
enableDatasetTestModelAccess(t, "tenant-1", "user-1", "instance-1")
permission := string(entity.TenantPermissionTeam)
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "user-1", "tenant-1"), "kb-1", "user-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "user-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
Permission: &permission,
EmbeddingModel: &embeddingModelID,
})
@@ -628,7 +628,7 @@ func TestDatasetServiceUpdateDatasetValidatesName(t *testing.T) {
name := " "
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{Name: &name})
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{Name: &name})
if err == nil {
t.Fatal("expected name validation error")
}
@@ -648,7 +648,7 @@ func TestDatasetServiceUpdateDatasetRejectsDuplicateName(t *testing.T) {
name := "Existing"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{Name: &name})
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{Name: &name})
if err == nil {
t.Fatal("expected duplicate name error")
}
@@ -666,7 +666,7 @@ func TestDatasetServiceUpdateDatasetRejectsNoPropertiesModified(t *testing.T) {
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{})
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{})
if err == nil {
t.Fatal("expected no-op update error")
}
@@ -686,7 +686,7 @@ func TestDatasetServiceUpdateDatasetLinksConnectors(t *testing.T) {
autoParse := "0"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Connectors: &[]service.DatasetConnectorRequest{{ID: "connector-1", AutoParse: autoParse}},
})
if err != nil {
@@ -723,7 +723,7 @@ func TestDatasetServiceUpdateDatasetRejectsCrossTenantConnector(t *testing.T) {
autoParse := "0"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Connectors: &[]service.DatasetConnectorRequest{{ID: "connector-1", AutoParse: autoParse}},
})
if err == nil {
@@ -757,7 +757,7 @@ func TestDatasetServiceUpdateDatasetAcceptsProviderInstanceEmbedding(t *testing.
embeddingModel := "embedding-2@test@ZHIPU-AI"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
EmbeddingModel: &embeddingModel,
})
if err != nil {
@@ -789,7 +789,7 @@ func TestDatasetServiceUpdateDatasetAcceptsEmbeddingModelID(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
embeddingModelID := "aabbccdd11223344aabbccdd11223344"
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
EmbeddingModel: &embeddingModelID,
})
if err != nil {
@@ -821,7 +821,7 @@ func TestDatasetServiceUpdateDatasetRejectsEmptyConnectorID(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
connectors := []service.DatasetConnectorRequest{{ID: " "}}
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Connectors: &connectors,
})
if err == nil {
@@ -857,7 +857,7 @@ func TestDatasetServiceUpdateDatasetRejectsInvalidEmbeddingModelFormat(t *testin
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
embdModel := tc.embeddingModel
_, code, err := svc.UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := svc.UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
EmbeddingModel: &embdModel,
})
if err == nil {
@@ -881,7 +881,7 @@ func TestDatasetServiceUpdateDatasetRejectsDuplicateNameCaseInsensitive(t *testi
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
uppercaseName := "EXISTING"
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Name: &uppercaseName,
})
if err == nil {
@@ -909,7 +909,7 @@ func TestDatasetServiceUpdateDatasetPreservesUnmodifiedFields(t *testing.T) {
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
newName := "Renamed Only"
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Name: &newName,
})
if err != nil {
@@ -954,7 +954,7 @@ func TestDatasetServiceUpdateDatasetPreservesParserConfigOnEmptyUpdate(t *testin
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
name := "Updated Name"
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Name: &name,
})
if err != nil {
@@ -1241,7 +1241,7 @@ func TestUpdateDataset_StripsUnknownParam_Builtin(t *testing.T) {
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"no_such_param": 1,
@@ -1278,7 +1278,7 @@ func TestUpdateDataset_AcceptsValidComponentParams_Builtin(t *testing.T) {
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{"parse_method": "deepdoc"},
@@ -1327,7 +1327,7 @@ func TestUpdateDatasetUsesTenantModelIDForExtractorDefault(t *testing.T) {
}
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{"parse_method": "deepdoc"},
@@ -1399,7 +1399,7 @@ func TestUpdateDataset_PreservesIncomingMetadataWhenCleaningParserConfig(t *test
},
}
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
ParserConfig: parserConfig,
})
if err != nil || code != common.CodeSuccess {
@@ -1454,7 +1454,7 @@ func TestUpdateDataset_PreservesExistingMetadataWhenParserConfigOmitsIt(t *testi
t.Fatalf("seed parser_config: %v", err)
}
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{"parse_method": "deepdoc"},
@@ -1497,7 +1497,7 @@ func TestUpdateDataset_StripsCanvasUnknownParam(t *testing.T) {
insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original", "canvas-1")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:NoSuch": map[string]interface{}{
"pdf": map[string]interface{}{},
@@ -1536,7 +1536,7 @@ func TestUpdateDataset_AcceptsValidCanvasComponentParams(t *testing.T) {
insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original", "canvas-1")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:CustomRhyme": map[string]interface{}{
"pdf": map[string]interface{}{},
@@ -1568,7 +1568,7 @@ func TestUpdateDataset_SwitchCanvasToBuiltinValidatesAgainstBuiltin(t *testing.T
chunkMethod := "naive"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
parseType := 1
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
ParseType: &parseType,
ParserConfig: map[string]interface{}{
@@ -1602,13 +1602,13 @@ func TestUpdateDatasetPagerankUnchangedIsANoOpOnInfinity(t *testing.T) {
svc.docEngine = fakeChatDocEngine{}
unchanged := int64(0) // insertDatasetUpdateKB stores pagerank 0
if _, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1",
if _, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1",
service.UpdateDatasetRequest{Pagerank: &unchanged}); err != nil || code != common.CodeSuccess {
t.Fatalf("unchanged pagerank: code=%d err=%v, want success (a no-op, as in Python)", code, err)
}
changed := int64(7)
_, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", "tenant-1",
_, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1",
service.UpdateDatasetRequest{Pagerank: &changed})
if err == nil || !strings.Contains(err.Error(), "can only be set when doc_engine is elasticsearch") {
t.Fatalf("changed pagerank: code=%d err=%v, want the doc_engine message", code, err)
internal/service/document/document_apply_parser_config.go
+1 -1
View File
@@ -29,7 +29,7 @@ type ApplyDatasetParserConfigResponse struct {
// ApplyDatasetParserConfig resets document parser settings without changing
// parsed content or scheduling ingestion. Dataset grants cover its documents.
func (s *DocumentService) ApplyDatasetParserConfig(ctx context.Context, userID, datasetID string) (*ApplyDatasetParserConfigResponse, common.ErrorCode, error) {
func (s *DocumentService) ApplyDatasetParserConfig(ctx context.Context, datasetID string) (*ApplyDatasetParserConfigResponse, common.ErrorCode, error) {
datasetID = strings.TrimSpace(datasetID)
if datasetID == "" {
return nil, common.CodeArgumentError, fmt.Errorf("dataset_id is required")
internal/service/document/document_apply_parser_config_test.go
+9 -9
View File
@@ -40,7 +40,7 @@ func TestApplyDatasetParserConfigReplacesSettingsAndPreservesContent(t *testing.
t.Fatal(err)
}
svc := testDocumentService(t)
result, code, err := svc.ApplyDatasetParserConfig(documentTestContext(t.Context()), "user-1", "kb-1")
result, code, err := svc.ApplyDatasetParserConfig(documentTestContext(t.Context()), "kb-1")
if err != nil || code != common.CodeSuccess {
t.Fatalf("apply: code=%d err=%v", code, err)
}
@@ -68,7 +68,7 @@ func TestApplyDatasetParserConfigReplacesSettingsAndPreservesContent(t *testing.
if err != nil || other.ParserID != "naive" || len(other.ParserConfig) != 0 {
t.Fatalf("other dataset changed: doc=%+v err=%v", other, err)
}
result, code, err = svc.ApplyDatasetParserConfig(documentTestContext(t.Context()), "user-1", "kb-1")
result, code, err = svc.ApplyDatasetParserConfig(documentTestContext(t.Context()), "kb-1")
if err != nil || code != common.CodeSuccess || result.UpdatedDocuments != 0 || result.UnchangedDocuments != 205 {
t.Fatalf("repeat: result=%+v code=%d err=%v", result, code, err)
}
@@ -102,7 +102,7 @@ func TestApplyDatasetParserConfigRollsBackAllBatches(t *testing.T) {
} else {
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-200", "kb-1", failure)
}
result, code, err := testDocumentService(t).ApplyDatasetParserConfig(documentTestContext(t.Context()), "user-1", "kb-1")
result, code, err := testDocumentService(t).ApplyDatasetParserConfig(documentTestContext(t.Context()), "kb-1")
if err == nil || code != wantCode || result != nil {
t.Fatalf("result=%+v code=%d err=%v", result, code, err)
}
@@ -136,7 +136,7 @@ func TestApplyDatasetParserConfigResolvesMixedFileTypes(t *testing.T) {
t.Fatal(err)
}
}
_, code, err := testDocumentService(t).ApplyDatasetParserConfig(documentTestContext(t.Context()), "user-1", "kb-1")
_, code, err := testDocumentService(t).ApplyDatasetParserConfig(documentTestContext(t.Context()), "kb-1")
if err != nil || code != common.CodeSuccess {
t.Fatalf("apply: code=%d err=%v", code, err)
}
@@ -159,21 +159,21 @@ func TestApplyDatasetParserConfigAuthorizationAndEmptyDataset(t *testing.T) {
}
svc := testDocumentService(t)
for _, id := range []string{"kb-1", "missing"} {
result, code, err := svc.ApplyDatasetParserConfig(documentTestActorContext(t.Context(), "outsider", "tenant-other"), "outsider", id)
result, code, err := svc.ApplyDatasetParserConfig(documentTestActorContext(t.Context(), "outsider", "tenant-other"), id)
if err == nil || code == common.CodeSuccess || result != nil {
t.Fatalf("unauthorized apply: result=%+v code=%d err=%v", result, code, err)
}
}
result, code, err := svc.ApplyDatasetParserConfig(documentTestContext(t.Context()), "user-1", "kb-1")
result, code, err := svc.ApplyDatasetParserConfig(documentTestContext(t.Context()), "kb-1")
if err != nil || code != common.CodeSuccess || result.TotalDocuments != 0 {
t.Fatalf("empty dataset: result=%+v code=%d err=%v", result, code, err)
}
ctx, cancel := context.WithCancel(documentTestContext(t.Context()))
cancel()
if _, _, err := svc.ApplyDatasetParserConfig(ctx, "user-1", "kb-1"); err == nil {
if _, _, err := svc.ApplyDatasetParserConfig(ctx, "kb-1"); err == nil {
t.Fatal("canceled request succeeded")
}
if _, _, err := svc.ApplyDatasetParserConfig(t.Context(), "user-1", " "); err == nil || !strings.Contains(err.Error(), "dataset_id") {
if _, _, err := svc.ApplyDatasetParserConfig(t.Context(), " "); err == nil || !strings.Contains(err.Error(), "dataset_id") {
t.Fatalf("empty dataset ID: %v", err)
}
}
@@ -211,7 +211,7 @@ func TestApplyDatasetParserConfigCopiesPipelineVersion(t *testing.T) {
}).Error; err != nil {
t.Fatal(err)
}
_, code, err := testDocumentService(t).ApplyDatasetParserConfig(documentTestContext(t.Context()), "user-1", "kb-1")
_, code, err := testDocumentService(t).ApplyDatasetParserConfig(documentTestContext(t.Context()), "kb-1")
doc, readErr := dao.NewDocumentDAO().GetByID(t.Context(), db, "doc-1")
if readErr != nil {
t.Fatal(readErr)
internal/service/document/document_applykb_chain_test.go
+5 -5
View File
@@ -132,7 +132,7 @@ func TestIngest_ApplyKB_RealDatasetAPIChain(t *testing.T) {
ctx := documentTestContext(t.Context())
enabled := true
if _, code, err := newDatasetService().UpdateMetadataConfig(ctx, "kb-1", "user-1", &service.MetadataConfigRequest{
if _, code, err := newDatasetService().UpdateMetadataConfig(ctx, "kb-1", &service.MetadataConfigRequest{
Enabled: &enabled,
Metadata: []service.MetadataConfigField{{Key: "author", Type: "string"}},
BuiltInMetadata: []service.MetadataConfigField{{Key: "document_name", Type: "string"}},
@@ -163,7 +163,7 @@ func TestIngest_ApplyKB_RealDatasetAPIChain(t *testing.T) {
svc := testDocumentService(t)
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusRunning),
Delete: true,
@@ -247,7 +247,7 @@ func TestIngest_ApplyKB_DisabledDatasetAPIChain(t *testing.T) {
ctx := documentTestContext(t.Context())
disabled := false
if _, code, err := newDatasetService().UpdateMetadataConfig(ctx, "kb-1", "user-1", &service.MetadataConfigRequest{
if _, code, err := newDatasetService().UpdateMetadataConfig(ctx, "kb-1", &service.MetadataConfigRequest{
Enabled: &disabled,
}); err != nil || code != common.CodeSuccess {
t.Fatalf("UpdateMetadataConfig err=%v code=%v", err, code)
@@ -257,7 +257,7 @@ func TestIngest_ApplyKB_DisabledDatasetAPIChain(t *testing.T) {
svc := testDocumentService(t)
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusRunning),
Delete: true,
@@ -386,7 +386,7 @@ func TestIngest_ApplyKB_RealCreatedDatasetChain(t *testing.T) {
svc := testDocumentService(t)
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
code, err = svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err = svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusRunning),
Delete: true,
internal/service/document/document_artifact.go
+27 -46
View File
@@ -24,7 +24,7 @@ var ErrDocumentImageNotFound = errors.New("document image not found")
// GetDocumentImage serves legacy KB-scoped image IDs only when an indexed
// chunk proves which accessible document owns the exact image ID.
func (s *DocumentService) GetDocumentImage(ctx context.Context, userID, imageID string) ([]byte, error) {
func (s *DocumentService) GetDocumentImage(ctx context.Context, imageID string) ([]byte, error) {
bucket, _, ok := strings.Cut(imageID, "-")
if !ok || len(bucket) != 32 {
return nil, ErrDocumentImageNotFound
@@ -36,12 +36,12 @@ func (s *DocumentService) GetDocumentImage(ctx context.Context, userID, imageID
if err != nil || docID == "" {
return nil, ErrDocumentImageNotFound
}
return s.GetDocumentImageForDocument(ctx, userID, docID, imageID)
return s.GetDocumentImageForDocument(ctx, docID, imageID)
}
// GetDocumentImageForDocument retrieves an image after proving that the exact
// composite image ID belongs to an indexed chunk of an accessible document.
func (s *DocumentService) GetDocumentImageForDocument(ctx context.Context, userID, docID, imageID string) ([]byte, error) {
func (s *DocumentService) GetDocumentImageForDocument(ctx context.Context, docID, imageID string) ([]byte, error) {
doc, err := s.documentDAO.GetByID(ctx, dao.DB, docID)
if err != nil || doc == nil || service.CheckDocumentAccess(ctx, doc.ID, doc.KbID, permission.OperationRead) != nil {
return nil, ErrDocumentImageNotFound
@@ -119,14 +119,14 @@ func (s *DocumentService) imageDocumentID(ctx context.Context, kbID, imageID, do
// GetDocumentArtifact retrieves a sandbox artifact from object storage.
//
// userID scopes the lookup: a CodeExec sandbox artifact is only
// The request identity scopes the lookup: a CodeExec sandbox artifact is only
// returned when the caller owns (or has team access to) at least
// one agent session whose `message` references this filename (or
// its `documents/artifact/<name>` URL form). The authorization
// gate runs BEFORE the storage read so a probe of an unknown
// filename cannot distinguish "you cannot see it" from "it
// exists" — both return ErrArtifactNotFound. Mirrors PR #16169.
func (s *DocumentService) GetDocumentArtifact(ctx context.Context, filename, userID string) (*ArtifactResponse, error) {
func (s *DocumentService) GetDocumentArtifact(ctx context.Context, filename string) (*ArtifactResponse, error) {
basename := filepath.Base(filename)
if basename != filename || strings.Contains(filename, "/") || strings.Contains(filename, "\\") {
return nil, ErrArtifactInvalidFilename
@@ -138,7 +138,7 @@ func (s *DocumentService) GetDocumentArtifact(ctx context.Context, filename, use
return nil, ErrArtifactInvalidFileType
}
if !s.sandboxArtifactAccessible(ctx, basename, userID) {
if !s.sandboxArtifactAccessible(ctx, basename) {
// Same error as "object does not exist" to avoid leaking
// whether the artifact exists for a different user/agent.
return nil, ErrArtifactNotFound
@@ -170,34 +170,21 @@ func (s *DocumentService) GetDocumentArtifact(ctx context.Context, filename, use
}, nil
}
// sandboxArtifactDialogIDsForUser returns the distinct agent
// (canvas) dialog_ids for sessions owned by userID whose
// `message` blob references filename. A CodeExec artifact URL
// appears in `message` as either a bare filename or the
// `documents/artifact/<name>` form, so the helper matches both.
//
// Implemented as a direct GORM query on the
// API4Conversation table — GORM's `Contains` maps to MySQL
// `LIKE '%...%'` which is fine here because the storage path is
// short and indexed lookup on (user_id, exp_user_id) keeps the
// scan narrow.
func (s *DocumentService) sandboxArtifactDialogIDsForUser(ctx context.Context, filename, userID string) []string {
// sandboxArtifactDialogIDs returns the distinct agent (canvas) dialog IDs for
// sessions owned by the request identity whose message references filename.
// Artifact URLs appear as either a bare filename or the
// documents/artifact/<name> form.
func (s *DocumentService) sandboxArtifactDialogIDs(ctx context.Context, filename string) []string {
userID := permission.UserIDFromContext(ctx)
if filename == "" || userID == "" {
return nil
}
// Escape SQL LIKE wildcards (%, _) before building the pattern.
// Without escaping, a caller could submit a filename like
// "%.png" or "_" and the LIKE query would match arbitrary
// referenced artifacts in any user's conversation — letting the
// caller pass the authorization check against one filename and
// then GET another artifact by name (PR review round 5, Major #8).
//
// Escape character: '!'. We avoid '\\' because SQL string
// literal parsing of '\\' is driver-specific (SQLite treats
// it as a single backslash, MySQL treats it as one, Postgres
// rejects the unterminated string) — '!' is a benign character
// in real filenames (artifact names rarely contain '!') and
// parses identically in every driver.
// Escape SQL LIKE wildcards before building the pattern. Without this, a
// filename such as %.png could match another user's referenced artifact.
// Use ! as the escape character because backslash handling differs between
// database drivers, while ! is parsed consistently by SQLite, MySQL, and
// Postgres.
filenameSafe := escapeSQLLikePattern(filename)
artifactRefSafe := escapeSQLLikePattern("documents/artifact/" + filename)
filenamePattern := "%" + filenameSafe + "%"
@@ -207,34 +194,28 @@ func (s *DocumentService) sandboxArtifactDialogIDsForUser(ctx context.Context, f
Select("c.dialog_id").
Joins("JOIN api_4_conversation_message AS cm ON cm.conversation_id = c.id").
Where("c.user_id = ? OR c.exp_user_id = ?", userID, userID).
Where(`cm.content LIKE ? ESCAPE '!' OR cm.content LIKE ? ESCAPE '!'`,
Where("cm.content LIKE ? ESCAPE '!' OR cm.content LIKE ? ESCAPE '!'",
filenamePattern, artifactRefPattern).
Distinct("c.dialog_id").
Rows()
Distinct("c.dialog_id").Rows()
if err != nil {
return nil
}
defer rows.Close()
for rows.Next() {
var d string
if err = rows.Scan(&d); err == nil && d != "" {
dialogIDs[d] = struct{}{}
var dialogID string
if err = rows.Scan(&dialogID); err == nil && dialogID != "" {
dialogIDs[dialogID] = struct{}{}
}
}
out := make([]string, 0, len(dialogIDs))
for d := range dialogIDs {
out = append(out, d)
for dialogID := range dialogIDs {
out = append(out, dialogID)
}
return out
}
// sandboxArtifactAccessible reports whether userID may read at least one
// canvas whose session references filename.
func (s *DocumentService) sandboxArtifactAccessible(ctx context.Context, filename, userID string) bool {
if userID == "" {
return false
}
for _, dialogID := range s.sandboxArtifactDialogIDsForUser(ctx, filename, userID) {
func (s *DocumentService) sandboxArtifactAccessible(ctx context.Context, filename string) bool {
for _, dialogID := range s.sandboxArtifactDialogIDs(ctx, filename) {
if permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.ResourceRef{Kind: permission.ResourceKindCanvas, ID: dialogID}, permission.OperationRead) == nil {
return true
}
internal/service/document/document_crud.go
+1 -1
View File
@@ -216,7 +216,7 @@ func (s *DocumentService) DeleteDocument(ctx context.Context, id string) error {
//
// ids: specific document IDs; deleteAll: delete all docs in the dataset.
// Returns the number of successfully deleted documents.
func (s *DocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error) {
func (s *DocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID string) (int, error) {
// Deleting the entire dataset requires dataset-level access; deleting
// explicit documents also accepts direct document grants.
if deleteAll {
internal/service/document/document_dataset_update.go
+6 -6
View File
@@ -25,7 +25,7 @@ import (
"go.uber.org/zap"
)
func (s *DocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) {
func (s *DocumentService) BatchUpdateDocumentStatus(ctx context.Context, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) {
statusInt, convErr := strconv.Atoi(status)
if convErr != nil {
return nil, common.CodeArgumentError, fmt.Errorf("invalid status: %s", status)
@@ -138,7 +138,7 @@ func (s *DocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID,
return result, common.CodeSuccess, nil
}
func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *UpdateDatasetDocumentRequest, present map[string]bool) (*UpdateDatasetDocumentResponse, common.ErrorCode, error) {
func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, datasetID, documentID string, req *UpdateDatasetDocumentRequest, present map[string]bool) (*UpdateDatasetDocumentResponse, common.ErrorCode, error) {
if err := service.CheckDocumentAccess(ctx, documentID, datasetID, permission.OperationUpdate); err != nil {
code, permissionErr := permissionresponse.Normalize(err)
return nil, code, permissionErr
@@ -159,7 +159,7 @@ func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, dat
}
return nil, common.CodeServerError, err
}
if code, err := s.validateDatasetDocumentUpdate(ctx, datasetID, documentID, userID, doc, req, present); err != nil {
if code, err := s.validateDatasetDocumentUpdate(ctx, datasetID, documentID, doc, req, present); err != nil {
return nil, code, err
}
@@ -292,7 +292,7 @@ func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, dat
}
pipelinepkg.ApplyParentChildChunkerConfig(cleaned, parentChildOverrides)
llmID := ""
if modelInfo, modelErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{UserID: userID, TenantID: kb.TenantID}, entity.ModelTypeChat); modelErr == nil {
if modelInfo, modelErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{TenantID: kb.TenantID}, entity.ModelTypeChat); modelErr == nil {
llmID = service.ModelReference(modelInfo)
}
cleaned = service.ApplyComponentScopedParserConfig(cleaned, llmID)
@@ -323,7 +323,7 @@ func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, dat
zap.String("parserID", effParserID), zap.Error(cpErr))
} else {
llmID := ""
if modelInfo, llmErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{UserID: userID, TenantID: kb.TenantID}, entity.ModelTypeChat); llmErr == nil {
if modelInfo, llmErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{TenantID: kb.TenantID}, entity.ModelTypeChat); llmErr == nil {
llmID = service.ModelReference(modelInfo)
} else if !errors.Is(llmErr, service.ErrNoDefaultModel) {
common.Warn("document pipeline/version switch: failed to resolve default chat model, falling back to no llm binding",
@@ -420,7 +420,7 @@ func (s *DocumentService) validateDocumentModifiable(ctx context.Context, doc *e
}
}
func (s *DocumentService) validateDatasetDocumentUpdate(ctx context.Context, datasetID, documentID, userID string, doc *entity.Document, req *UpdateDatasetDocumentRequest, present map[string]bool) (common.ErrorCode, error) {
func (s *DocumentService) validateDatasetDocumentUpdate(ctx context.Context, datasetID, documentID string, doc *entity.Document, req *UpdateDatasetDocumentRequest, present map[string]bool) (common.ErrorCode, error) {
if req == nil {
return common.CodeDataError, errors.New("invalid request payload")
}
internal/service/document/document_dataset_update_test.go
+11 -11
View File
@@ -40,7 +40,7 @@ func TestUpdateDatasetDocumentParentChildPatchUsesRequestedFields(t *testing.T)
}).Error; err != nil {
t.Fatal(err)
}
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParserConfig: map[string]any{"GeneralChunker:SixApplesFall": tc.patch},
}, map[string]bool{"parser_config": true})
if code != common.CodeSuccess || err != nil {
@@ -94,7 +94,7 @@ func TestUpdateDatasetDocumentConfigPatchPreservesOmittedSettings(t *testing.T)
if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1").Updates(updates).Error; err != nil {
t.Fatal(err)
}
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", req, present)
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", req, present)
if err != nil || code != common.CodeSuccess {
t.Fatalf("update: code=%d err=%v", code, err)
}
@@ -155,7 +155,7 @@ func TestUpdateDatasetDocumentRejectsUnsupportedFileFamilyBeforeWrites(t *testin
svc := testDocumentService(t)
svc.docEngine = engine
svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine)
_, code, err := svc.UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", req, present)
_, code, err := svc.UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", req, present)
if code != common.CodeDataError || err == nil || !strings.Contains(err.Error(), "does not support") {
t.Fatalf("expected unsupported file rejection: code=%d err=%v", code, err)
}
@@ -207,7 +207,7 @@ func TestUpdateDatasetDocumentParserAliasPreservesConfig(t *testing.T) {
req.ParserConfig = map[string]any{"Parser:HipSignsRhyme": map[string]any{"enable_vision_enhancement": true}}
present["parser_config"] = true
}
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", req, present)
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", req, present)
if err != nil || code != common.CodeSuccess {
t.Fatalf("update: code=%d err=%v", code, err)
}
@@ -259,7 +259,7 @@ func TestUpdateDatasetDocumentSelectionWithoutConfig(t *testing.T) {
} else {
req.ChunkMethod = new("manual")
}
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", req, present)
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", req, present)
if err != nil || code != common.CodeSuccess {
t.Fatalf("update: code=%d err=%v", code, err)
}
@@ -298,7 +298,7 @@ func TestUpdateDatasetDocumentChunkMethodAcceptsCurrentTemplates(t *testing.T) {
pushServiceDB(t, db)
insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0)
insertNamedTestDoc(t, "doc-1", "kb-1", file.name, 0, 0)
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ChunkMethod: &file.method,
}, map[string]bool{"chunk_method": true})
if err != nil || code != common.CodeSuccess {
@@ -339,7 +339,7 @@ func TestUpdateDatasetDocumentUsesEffectiveParserForConfig(t *testing.T) {
} else if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("parser_id", "manual").Error; err != nil {
t.Fatal(err)
}
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", req, present)
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", req, present)
if err != nil || code != common.CodeSuccess {
t.Fatalf("update: code=%d err=%v", code, err)
}
@@ -379,7 +379,7 @@ func TestUpdateDatasetDocumentKeepsDocumentCanvasForConfigPatch(t *testing.T) {
if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("pipeline_id", pipelineID).Error; err != nil {
t.Fatal(err)
}
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParserConfig: map[string]any{"Parser:custom": map[string]any{"vlm": map[string]any{"llm_id": "vision-model"}}},
}, map[string]bool{"parser_config": true})
if err != nil || code != common.CodeSuccess {
@@ -455,7 +455,7 @@ func TestUpdateDatasetDocumentRejectsParserErrorsBeforeWrites(t *testing.T) {
svc := testDocumentService(t)
svc.docEngine = engine
svc.metadataSvc = service.NewMetadataServiceForTest(dao.NewKnowledgebaseDAO(), engine)
_, code, err := svc.UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", req, present)
_, code, err := svc.UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", req, present)
if code != common.CodeDataError || err == nil || !strings.Contains(err.Error(), wantError) {
t.Fatalf("expected %q, got code=%d err=%v", wantError, code, err)
}
@@ -497,7 +497,7 @@ func TestBatchUpdateDocumentStatusReturnsPermissionErrorWhenAllDenied(t *testing
svc := testDocumentService(t)
ctx := documentTestActorContext(t.Context(), "member", "tenant-1")
_, code, err := svc.BatchUpdateDocumentStatus(ctx, "member", "kb-1", "0", []string{"doc-1"})
_, code, err := svc.BatchUpdateDocumentStatus(ctx, "kb-1", "0", []string{"doc-1"})
if err == nil {
t.Fatal("expected permission error for an inaccessible document")
}
@@ -530,7 +530,7 @@ func TestBatchUpdateDocumentStatusRetainsPartialFailure(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestActorContext(t.Context(), "member", "tenant-1")
_, code, err := svc.BatchUpdateDocumentStatus(ctx, "member", "kb-1", "1", []string{"member-doc", "owner-doc"})
_, code, err := svc.BatchUpdateDocumentStatus(ctx, "kb-1", "1", []string{"member-doc", "owner-doc"})
if err == nil {
t.Fatal("expected a partial failure")
}
internal/service/document/document_ingest.go
+13 -7
View File
@@ -13,11 +13,15 @@ import (
permissionresponse "ragflow/internal/permission/response"
)
func (s *DocumentService) ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) {
return s.ingestionTaskSvc.ListByUser(ctx, userID, datasetID, page, pageSize)
func (s *DocumentService) ListIngestionTasks(ctx context.Context, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) {
return s.ingestionTaskSvc.ListByUser(ctx, permission.UserIDFromContext(ctx), datasetID, page, pageSize)
}
func (s *DocumentService) IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
func (s *DocumentService) IngestDocuments(ctx context.Context, datasetID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
userID := permission.UserIDFromContext(ctx)
if userID == "" {
return nil, permission.ErrUnauthenticated
}
responses, err := s.ingestionTaskSvc.CreateForDocuments(ctx, datasetID, userID, docIDs)
if err != nil {
return nil, err
@@ -26,15 +30,17 @@ func (s *DocumentService) IngestDocuments(ctx context.Context, datasetID, userID
return responses, nil
}
func (s *DocumentService) StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error) {
func (s *DocumentService) StopIngestionTasks(ctx context.Context, tasks []string) ([]*entity.IngestionTask, error) {
userID := permission.UserIDFromContext(ctx)
return s.ingestionTaskSvc.RequestStopMany(ctx, tasks, &userID)
}
func (s *DocumentService) RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error) {
func (s *DocumentService) RemoveIngestionTasks(ctx context.Context, tasks []string) ([]map[string]string, error) {
userID := permission.UserIDFromContext(ctx)
return s.ingestionTaskSvc.RemoveMany(ctx, tasks, &userID)
}
func (s *DocumentService) Ingest(ctx context.Context, userID string, req *IngestDocumentRequest) (common.ErrorCode, error) {
func (s *DocumentService) Ingest(ctx context.Context, req *IngestDocumentRequest) (common.ErrorCode, error) {
run := fmt.Sprint(req.Run)
docs, err := s.documentDAO.GetByIDs(ctx, dao.DB, req.DocIDs)
@@ -112,7 +118,7 @@ func (s *DocumentService) Ingest(ctx context.Context, userID string, req *Ingest
for _, vd := range toStart {
doc := vd.doc
kb := vd.kb
if err = s.StartParseDocuments(ctx, doc, kb, userID, StartParseOptions{
if err = s.StartParseDocuments(ctx, doc, kb, StartParseOptions{
ApplyKB: req.ApplyKB,
RerunWithDelete: req.Delete,
}); err != nil {
internal/service/document/document_list.go
+1 -1
View File
@@ -139,7 +139,7 @@ func (s *DocumentService) ListDocuments(ctx context.Context, page, pageSize int)
return responses, total, nil
}
func (s *DocumentService) GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error) {
func (s *DocumentService) GetThumbnails(ctx context.Context, docIDs []string) (map[string]string, error) {
if len(docIDs) == 0 {
return map[string]string{}, nil
}
internal/service/document/document_modify_guard_test.go
+4 -4
View File
@@ -92,7 +92,7 @@ func updateDatasetDocumentRejected(t *testing.T, status string, req *UpdateDatas
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
_, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", req, present)
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", req, present)
if err == nil {
t.Fatalf("expected task-state rejection, got nil error (code=%v)", code)
}
@@ -118,7 +118,7 @@ func updateDatasetDocumentAllowed(t *testing.T, status string, req *UpdateDatase
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
_, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", req, present)
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", req, present)
if err != nil && strings.Contains(err.Error(), "cannot be modified") {
t.Fatalf("editable doc was wrongly rejected: %v", err)
}
@@ -201,7 +201,7 @@ func TestUpdateDatasetDocumentRunningEmptyPresentNotRejected(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
_, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1",
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{}, map[string]bool{})
if err != nil {
t.Fatalf("empty PATCH on RUNNING doc should not be rejected: code=%v err=%v", code, err)
@@ -231,7 +231,7 @@ func TestUpdateDatasetDocumentAllowsDoneRename(t *testing.T) {
name := "renamed.txt"
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
_, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1",
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{Name: &name},
map[string]bool{"name": true})
if err != nil && strings.Contains(err.Error(), "cannot be modified") {
internal/service/document/document_parse.go
+7 -3
View File
@@ -17,6 +17,7 @@ import (
"ragflow/internal/entity"
"ragflow/internal/ingestion/knowledge_compile"
ingestionpipeline "ragflow/internal/ingestion/pipeline"
"ragflow/internal/permission"
"ragflow/internal/storage"
"gorm.io/gorm"
@@ -112,7 +113,10 @@ func (s *DocumentService) deletePreparedDocumentTask(ctx context.Context, docID
// KB config (ApplyKB), validates storage, and enqueues an ingestion task.
// Extracted from Ingest so other entry points (e.g. ChunkService.Parse)
// can reuse the same start-parse flow.
func (s *DocumentService) StartParseDocuments(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, userID string, opts StartParseOptions) error {
func (s *DocumentService) StartParseDocuments(ctx context.Context, doc *entity.Document, kb *entity.Knowledgebase, opts StartParseOptions) error {
if permission.UserIDFromContext(ctx) == "" {
return permission.ErrUnauthenticated
}
// Validate storage first so we don't clear prior results and then fail
// because the document can't be read, leaving the document with neither
// old nor new parse results.
@@ -134,7 +138,7 @@ func (s *DocumentService) StartParseDocuments(ctx context.Context, doc *entity.D
}
}
responses, err := s.IngestDocuments(ctx, doc.KbID, userID, []string{doc.ID})
responses, err := s.IngestDocuments(ctx, doc.KbID, []string{doc.ID})
if err != nil {
return err
}
@@ -431,7 +435,7 @@ func (s *DocumentService) clearKBChunkNumWhenRerun(doc *entity.Document) error {
}).Error
}
func (s *DocumentService) ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
func (s *DocumentService) ParseDocuments(ctx context.Context, datasetID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
// deduplicate the document id
uniqueDocIDs := common.Deduplicate(docIDs)
if uniqueDocIDs == nil || len(uniqueDocIDs) == 0 {
internal/service/document/document_test.go
+88 -65
View File
@@ -1210,7 +1210,7 @@ func TestSyncDocumentUpsertRemovesStagedBlobWhenInsertFails(t *testing.T) {
svc := testDocumentService(t)
_, err := svc.Upsert(ctx, service.DocumentUpsertInput{
TaskContext: service.SyncTaskContext{
Connector: entity.Connector{TenantID: "tenant-1"},
Connector: entity.Connector{UserID: "tenant-1", TenantID: "tenant-1"},
Knowledgebase: entity.Knowledgebase{
ID: "missing-kb",
TenantID: "tenant-1",
@@ -1264,7 +1264,7 @@ func TestSyncDocumentUpsertWithNilDatasetParserConfig(t *testing.T) {
result, err := testDocumentService(t).Upsert(ctx, service.DocumentUpsertInput{
TaskContext: service.SyncTaskContext{
Connector: entity.Connector{TenantID: "tenant-1"},
Connector: entity.Connector{UserID: "tenant-1", TenantID: "tenant-1"},
Knowledgebase: kb,
},
SourceType: "github",
@@ -1326,7 +1326,7 @@ func TestSyncDocumentUpsertRemovesStagedBlobWhenUpdateFails(t *testing.T) {
mockStorage.afterPut = cancel
_, err := svc.Upsert(ctx, service.DocumentUpsertInput{
TaskContext: service.SyncTaskContext{
Connector: entity.Connector{TenantID: "tenant-1"},
Connector: entity.Connector{UserID: "tenant-1", TenantID: "tenant-1"},
Knowledgebase: entity.Knowledgebase{
ID: "kb-sync",
TenantID: "tenant-1",
@@ -1583,7 +1583,7 @@ func TestDeleteDocuments_DeleteAll(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
deleted, err := svc.DeleteDocuments(ctx, nil, true, "kb-1", "user-1")
deleted, err := svc.DeleteDocuments(ctx, nil, true, "kb-1")
if err != nil {
t.Fatalf("DeleteDocuments failed: %v", err)
}
@@ -1613,7 +1613,7 @@ func TestDeleteDocuments_ByIDs(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
deleted, err := svc.DeleteDocuments(ctx, []string{"doc-1", "doc-2"}, false, "kb-1", "user-1")
deleted, err := svc.DeleteDocuments(ctx, []string{"doc-1", "doc-2"}, false, "kb-1")
if err != nil {
t.Fatalf("DeleteDocuments failed: %v", err)
}
@@ -1641,7 +1641,7 @@ func TestDeleteDocuments_WrongDataset(t *testing.T) {
svc := testDocumentService(t)
ctx := t.Context()
_, err := svc.DeleteDocuments(ctx, []string{"doc-1"}, false, "kb-1", "user-1")
_, err := svc.DeleteDocuments(ctx, []string{"doc-1"}, false, "kb-1")
if err == nil {
t.Fatal("expected error for doc not belonging to dataset")
}
@@ -1657,7 +1657,7 @@ func TestDeleteDocuments_NotAccessible(t *testing.T) {
ctx := t.Context()
// user-1 has no user_tenant entry → accessible returns false
_, err := svc.DeleteDocuments(ctx, []string{"doc-1"}, false, "kb-1", "user-1")
_, err := svc.DeleteDocuments(ctx, []string{"doc-1"}, false, "kb-1")
if err == nil {
t.Fatal("expected error for inaccessible dataset")
}
@@ -1673,7 +1673,7 @@ func TestDeleteDocuments_EmptyIDs(t *testing.T) {
// Empty ids, no deleteAll → returns 0, no error
ctx := t.Context()
deleted, err := svc.DeleteDocuments(ctx, []string{}, false, "kb-1", "user-1")
deleted, err := svc.DeleteDocuments(ctx, []string{}, false, "kb-1")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -1694,7 +1694,7 @@ func TestDeleteDocuments_Deduplicate(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
deleted, err := svc.DeleteDocuments(ctx, []string{"doc-1", "doc-1", "doc-1"}, false, "kb-1", "user-1")
deleted, err := svc.DeleteDocuments(ctx, []string{"doc-1", "doc-1", "doc-1"}, false, "kb-1")
if err != nil {
t.Fatalf("DeleteDocuments failed: %v", err)
}
@@ -1773,7 +1773,7 @@ func TestStartParseDocuments_ApplyKBMergesDatasetParserConfig(t *testing.T) {
publisher := &recordingTaskPublisher{}
svc := testDocumentService(t)
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := t.Context()
ctx := documentTestContext(t.Context())
kb, err := svc.kbDAO.GetByID(ctx, db, "kb-1")
if err != nil {
@@ -1784,7 +1784,7 @@ func TestStartParseDocuments_ApplyKBMergesDatasetParserConfig(t *testing.T) {
t.Fatalf("load doc: %v", err)
}
if err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{ApplyKB: true}); err != nil {
if err = svc.StartParseDocuments(ctx, doc, kb, StartParseOptions{ApplyKB: true}); err != nil {
t.Fatalf("StartParseDocuments: %v", err)
}
@@ -1885,7 +1885,7 @@ func TestIngest_ApplyKBPersistsDatasetParserConfig(t *testing.T) {
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusRunning),
Delete: true,
@@ -1974,7 +1974,7 @@ func TestStartParseDocuments_ApplyKBUsesDatasetDefaults(t *testing.T) {
publisher := &recordingTaskPublisher{}
svc := testDocumentService(t)
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := t.Context()
ctx := documentTestContext(t.Context())
kb, err := svc.kbDAO.GetByID(ctx, db, "kb-1")
if err != nil {
@@ -1985,7 +1985,7 @@ func TestStartParseDocuments_ApplyKBUsesDatasetDefaults(t *testing.T) {
t.Fatalf("load doc: %v", err)
}
if err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{ApplyKB: true}); err != nil {
if err = svc.StartParseDocuments(ctx, doc, kb, StartParseOptions{ApplyKB: true}); err != nil {
t.Fatalf("StartParseDocuments: %v", err)
}
@@ -2080,7 +2080,7 @@ func TestIngest_ApplyKBDisabledLeavesParserConfigUnchanged(t *testing.T) {
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusRunning),
Delete: true,
@@ -2165,7 +2165,7 @@ func TestStartParseDocuments_ApplyKBPersistFailureNamesDocument(t *testing.T) {
publisher := &recordingTaskPublisher{}
svc := testDocumentService(t)
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := t.Context()
ctx := documentTestContext(t.Context())
kb, err := svc.kbDAO.GetByID(ctx, db, "kb-1")
if err != nil {
@@ -2176,7 +2176,7 @@ func TestStartParseDocuments_ApplyKBPersistFailureNamesDocument(t *testing.T) {
t.Fatalf("load doc: %v", err)
}
err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{ApplyKB: true})
err = svc.StartParseDocuments(ctx, doc, kb, StartParseOptions{ApplyKB: true})
if err == nil {
t.Fatal("expected an error when the merged parser config cannot be persisted")
}
@@ -2261,7 +2261,7 @@ func TestStartParseDocuments_EnqueuesIngestionTask(t *testing.T) {
publisher := &recordingTaskPublisher{}
svc := testDocumentService(t)
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := t.Context()
ctx := documentTestContext(t.Context())
kb, err := svc.kbDAO.GetByID(ctx, db, "kb-1")
if err != nil {
@@ -2272,7 +2272,7 @@ func TestStartParseDocuments_EnqueuesIngestionTask(t *testing.T) {
t.Fatalf("load doc: %v", err)
}
if err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{}); err != nil {
if err = svc.StartParseDocuments(ctx, doc, kb, StartParseOptions{}); err != nil {
t.Fatalf("StartParseDocuments: %v", err)
}
@@ -2317,6 +2317,28 @@ func TestStartParseDocuments_EnqueuesIngestionTask(t *testing.T) {
}
}
func TestIngestDocumentsRejectsMissingUserIdentity(t *testing.T) {
db := setupServiceTestDB(t)
pushServiceDB(t, db)
svc := testDocumentService(t)
publisher := &recordingTaskPublisher{}
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
if _, err := svc.IngestDocuments(t.Context(), "kb-1", []string{"doc-1"}); !errors.Is(err, permission.ErrUnauthenticated) {
t.Fatalf("IngestDocuments error = %v, want unauthenticated", err)
}
if len(publisher.messages) != 0 {
t.Fatalf("parse publications = %d, want none", len(publisher.messages))
}
var taskCount int64
if err := db.Model(&entity.IngestionTask{}).Count(&taskCount).Error; err != nil {
t.Fatal(err)
}
if taskCount != 0 {
t.Fatalf("ingestion tasks = %d, want none", taskCount)
}
}
func TestStartParseDocumentsSerializesConcurrentReruns(t *testing.T) {
db := setupConcurrentServiceTestDB(t)
pushServiceDB(t, db)
@@ -2347,7 +2369,7 @@ func TestStartParseDocumentsSerializesConcurrentReruns(t *testing.T) {
secondPublisher := &runningTaskPublisher{db: db, published: make(chan struct{})}
secondService := testDocumentService(t)
secondService.ingestionTaskSvc.SetTaskPublisher(secondPublisher)
ctx := t.Context()
ctx := documentTestContext(t.Context())
kb, err := firstService.kbDAO.GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("load knowledgebase: %v", err)
@@ -2359,13 +2381,13 @@ func TestStartParseDocumentsSerializesConcurrentReruns(t *testing.T) {
firstDone := make(chan error, 1)
go func() {
firstDone <- firstService.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{RerunWithDelete: true})
firstDone <- firstService.StartParseDocuments(ctx, doc, kb, StartParseOptions{RerunWithDelete: true})
}()
<-firstTaskLookup
secondDone := make(chan error, 1)
go func() {
secondDone <- secondService.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{RerunWithDelete: true})
secondDone <- secondService.StartParseDocuments(ctx, doc, kb, StartParseOptions{RerunWithDelete: true})
}()
secondPublishedBeforeFirstCompletes := false
@@ -2398,7 +2420,7 @@ func TestStartParseDocumentsSupersedesQueuedRun(t *testing.T) {
svc := testDocumentService(t)
svc.ingestionTaskSvc.SetTaskPublisher(&recordingTaskPublisher{})
ctx := t.Context()
ctx := documentTestContext(t.Context())
kb, err := svc.kbDAO.GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("load knowledgebase: %v", err)
@@ -2408,10 +2430,10 @@ func TestStartParseDocumentsSupersedesQueuedRun(t *testing.T) {
t.Fatalf("load document: %v", err)
}
if err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{}); err != nil {
if err = svc.StartParseDocuments(ctx, doc, kb, StartParseOptions{}); err != nil {
t.Fatalf("start first parse: %v", err)
}
if err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{RerunWithDelete: true}); err != nil {
if err = svc.StartParseDocuments(ctx, doc, kb, StartParseOptions{RerunWithDelete: true}); err != nil {
t.Fatalf("rerun queued parse: %v", err)
}
@@ -2865,7 +2887,7 @@ func TestArtifactHelpers(t *testing.T) {
func TestGetDocumentArtifact_InvalidFilename(t *testing.T) {
svc := testDocumentService(t)
ctx := t.Context()
_, err := svc.GetDocumentArtifact(ctx, "../test.txt", "user-1")
_, err := svc.GetDocumentArtifact(ctx, "../test.txt")
if !errors.Is(err, ErrArtifactInvalidFilename) {
t.Errorf("expected ErrArtifactInvalidFilename, got %v", err)
}
@@ -2874,7 +2896,7 @@ func TestGetDocumentArtifact_InvalidFilename(t *testing.T) {
func TestGetDocumentArtifact_InvalidFileType(t *testing.T) {
svc := testDocumentService(t)
ctx := t.Context()
_, err := svc.GetDocumentArtifact(ctx, "test.exe", "user-1")
_, err := svc.GetDocumentArtifact(ctx, "test.exe")
if !errors.Is(err, ErrArtifactInvalidFileType) {
t.Errorf("expected ErrArtifactInvalidFileType, got %v", err)
}
@@ -3079,7 +3101,7 @@ func TestUpdateDatasetDocumentRejectsNonOwner(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestActorContext(t.Context(), "user-2", "tenant-1")
_, code, err := svc.UpdateDatasetDocument(ctx, "user-2", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{}, map[string]bool{})
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{}, map[string]bool{})
if err == nil {
t.Fatal("expected ownership error")
}
@@ -3100,7 +3122,7 @@ func TestUpdateDatasetDocumentRejectsCounterMutation(t *testing.T) {
chunkCount := int64(6)
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
_, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ChunkCount: &chunkCount,
}, map[string]bool{"chunk_count": true})
if err == nil {
@@ -3138,7 +3160,7 @@ func TestUpdateDatasetDocumentRejectsZeroImmutableFields(t *testing.T) {
t.Fatalf("prepare document: %v", err)
}
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &tt.request, tt.present)
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &tt.request, tt.present)
if err == nil {
t.Fatalf("expected %s mutation error", tt.name)
}
@@ -3287,7 +3309,7 @@ func TestUpdateDatasetDocumentSavesVisionModelAcrossFileTypes(t *testing.T) {
} else {
req.ChunkMethod = &file.parserID
}
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", req, present)
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", req, present)
if err != nil || code != common.CodeSuccess {
t.Fatalf("save vision model: code=%v err=%v", code, err)
}
@@ -3321,7 +3343,7 @@ func TestUpdateDatasetDocumentRejectedParserConfigPreservesNames(t *testing.T) {
insertTestFile(t, "file-1", "folder-1", "old.pdf", new("old.pdf"))
insertTestFile2Document(t, "f2d-1", "file-1", "doc-1")
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
Name: new("must_not_persist.pdf"),
ParserConfig: map[string]any{
"Parser:HipSignsRhyme": map[string]any{
@@ -3362,7 +3384,7 @@ func TestUpdateDatasetDocumentRenameUpdatesDocumentAndFile(t *testing.T) {
newName := "new.pdf"
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
resp, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
Name: &newName,
}, map[string]bool{"name": true})
if err != nil {
@@ -3428,7 +3450,7 @@ func TestUpdateDatasetDocumentParserConfigPreservesParseResults(t *testing.T) {
engine := &rerunDeleteDocEngine{}
svc := testDocumentService(t)
svc.docEngine = engine
resp, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &request, present)
resp, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &request, present)
if err != nil || code != common.CodeSuccess {
t.Fatalf("UpdateDatasetDocument failed: code=%v err=%v", code, err)
}
@@ -3863,7 +3885,7 @@ func TestIngest_RunningTask_SkippedWhenBatchStartingParse(t *testing.T) {
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1", "doc-2", "doc-3"},
Run: string(entity.TaskStatusRunning),
Delete: true,
@@ -3906,7 +3928,7 @@ func TestIngest_AllActiveTasks_ReturnsSuccessIdempotent(t *testing.T) {
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1", "doc-2"},
Run: string(entity.TaskStatusRunning),
Delete: true,
@@ -3939,7 +3961,7 @@ func TestIngest_CompletedTask_SkippedWhenDeleteIsFalse(t *testing.T) {
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1", "doc-2"},
Run: string(entity.TaskStatusRunning),
Delete: false,
@@ -3981,7 +4003,7 @@ func TestIngest_FailedAndStoppedTasks_Restarted(t *testing.T) {
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1", "doc-2"},
Run: string(entity.TaskStatusRunning),
Delete: false,
@@ -4010,7 +4032,7 @@ func TestUpdateDatasetDocumentPropagatesMetadataDeleteFailure(t *testing.T) {
svc.docEngine = engine
svc.metadataSvc = service.NewMetadataServiceForTest(nil, nil)
ctx := documentTestContext(t.Context())
_, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
MetaFields: map[string]any{"new": "value"},
}, map[string]bool{"meta_fields": true})
if err == nil {
@@ -4351,7 +4373,7 @@ func TestUpdateDatasetDocumentPipelineIDTakesPrecedenceOverParserID(t *testing.T
parseType := 2
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
resp, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
PipelineID: &pipelineID,
PipelineVersionID: sptr("v-pre-0001"),
ParseType: &parseType,
@@ -4387,7 +4409,7 @@ func TestUpdateDatasetDocumentParseTypeBuiltin(t *testing.T) {
pipelineIDEmpty := ""
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
resp, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParseType: &parseType,
ParserID: &parserID,
PipelineID: &pipelineIDEmpty,
@@ -4424,7 +4446,7 @@ func TestUpdateDatasetDocumentParseTypePipeline(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
seedDocVersion(t, db, pipelineID, "v-ptp-0001", 0)
resp, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParseType: &parseType,
PipelineID: &pipelineID,
PipelineVersionID: sptr("v-ptp-0001"),
@@ -4455,7 +4477,7 @@ func TestUpdateDatasetDocumentParseTypePipelineRejectsParserID(t *testing.T) {
pipelineID := "1234567890abcdef1234567890abcdef"
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
_, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParseType: &parseType,
ParserID: &parserID,
PipelineID: &pipelineID,
@@ -4485,7 +4507,7 @@ func TestUpdateDatasetDocumentParentChildConfigReachesCanvasChunker(t *testing.T
// submitted parent_child would be dropped on a switch; this test verifies the
// config reaches the canvas chunker when no switch occurs.
bindDocCanvas(t, db, "kb-1", "doc-1", pipelineID, "v-pcc-0001")
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParserConfig: map[string]any{
"GeneralChunker:SixApplesFall": map[string]any{
"parent_child": map[string]any{
@@ -4520,7 +4542,7 @@ func TestUpdateDatasetDocumentRejectsInvalidPages(t *testing.T) {
parserID := "manual"
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
_, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
_, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParseType: &parseType,
ParserID: &parserID,
ParserConfig: map[string]any{
@@ -4545,7 +4567,7 @@ func TestUpdateDatasetDocumentParentChildConfigReachesGeneralChunker(t *testing.
insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5)
insertNamedTestDoc(t, "doc-1", "kb-1", "doc.txt", 10, 5)
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParserConfig: map[string]any{
"GeneralChunker:SixApplesFall": map[string]any{
"parent_child": map[string]any{
@@ -4583,7 +4605,7 @@ func TestUpdateDatasetDocumentEnabledUpdatesStatus(t *testing.T) {
enabled := 0
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
resp, code, err := svc.UpdateDatasetDocument(ctx, "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
resp, code, err := svc.UpdateDatasetDocument(ctx, "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
Enabled: &enabled,
}, map[string]bool{"enabled": true})
if err != nil {
@@ -4754,7 +4776,7 @@ func TestGetDocumentArtifact_AuthGate(t *testing.T) {
}
// Case 0: the owner can fetch the artifact referenced by their session.
artifact, err := svc.GetDocumentArtifact(ctx, "result.png", "user-1")
artifact, err := svc.GetDocumentArtifact(ctx, "result.png")
if err != nil {
t.Fatalf("owner: want artifact, got %v", err)
}
@@ -4763,12 +4785,13 @@ func TestGetDocumentArtifact_AuthGate(t *testing.T) {
}
// Case 1: empty user -> not allowed.
if _, err := svc.GetDocumentArtifact(ctx, "result.png", ""); !errors.Is(err, ErrArtifactNotFound) {
if _, err := svc.GetDocumentArtifact(t.Context(), "result.png"); !errors.Is(err, ErrArtifactNotFound) {
t.Errorf("empty user: want ErrArtifactNotFound, got %v", err)
}
// Case 2: another user without any session reference -> not allowed.
if _, err := svc.GetDocumentArtifact(ctx, "result.png", "user-2"); !errors.Is(err, ErrArtifactNotFound) {
otherUserCtx := documentTestActorContext(t.Context(), "user-2", "tenant-1")
if _, err := svc.GetDocumentArtifact(otherUserCtx, "result.png"); !errors.Is(err, ErrArtifactNotFound) {
t.Errorf("unrelated user: want ErrArtifactNotFound, got %v", err)
}
@@ -4791,7 +4814,7 @@ func TestGetDocumentArtifact_AuthGate(t *testing.T) {
}); err != nil {
t.Fatalf("seed conv 2: %v", err)
}
if _, err := svc.GetDocumentArtifact(ctx, "result.png", "user-2"); !errors.Is(err, ErrArtifactNotFound) {
if _, err := svc.GetDocumentArtifact(otherUserCtx, "result.png"); !errors.Is(err, ErrArtifactNotFound) {
t.Errorf("user-2 with unrelated session: want ErrArtifactNotFound, got %v", err)
}
}
@@ -4861,7 +4884,7 @@ func TestGetThumbnails_AlignsWithPythonFormatting(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
got, err := svc.GetThumbnails(ctx, "user-1", []string{"doc-file", "doc-base64", "doc-other", "missing-doc"})
got, err := svc.GetThumbnails(ctx, []string{"doc-file", "doc-base64", "doc-other", "missing-doc"})
if err != nil {
t.Fatalf("GetThumbnails failed: %v", err)
}
@@ -4900,7 +4923,7 @@ func TestGetDocumentImageRejectsUnindexedObjectBeforeStorage(t *testing.T) {
svc := NewDocumentService()
svc.docEngine = engine
if _, err := svc.GetDocumentImage(ctx, "user-1", kbID+"-raw.pdf"); !errors.Is(err, ErrDocumentImageNotFound) {
if _, err := svc.GetDocumentImage(ctx, kbID+"-raw.pdf"); !errors.Is(err, ErrDocumentImageNotFound) {
t.Fatalf("GetDocumentImage() error = %v", err)
}
if store.getCalls != 0 {
@@ -4934,7 +4957,7 @@ func TestGetDocumentImageForDocumentAllowsIndexedSharedBucketImage(t *testing.T)
svc := NewDocumentService()
svc.docEngine = engine
got, err := svc.GetDocumentImageForDocument(ctx, "user-1", "doc-1", imageID)
got, err := svc.GetDocumentImageForDocument(ctx, "doc-1", imageID)
if err != nil {
t.Fatalf("GetDocumentImageForDocument() error = %v", err)
}
@@ -4988,7 +5011,7 @@ func TestStartParseDocuments_FailsBeforeClearing(t *testing.T) {
insertTestDoc(t, "doc-1", "kb-1", 10, 5)
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED)
ctx := t.Context()
ctx := documentTestContext(t.Context())
doc, err := dao.NewDocumentDAO().GetByID(ctx, db, "doc-1")
if err != nil {
t.Fatalf("get doc: %v", err)
@@ -5004,7 +5027,7 @@ func TestStartParseDocuments_FailsBeforeClearing(t *testing.T) {
}
svc := testDocumentService(t)
err = svc.StartParseDocuments(ctx, doc, kb, "user-1", StartParseOptions{RerunWithDelete: true})
err = svc.StartParseDocuments(ctx, doc, kb, StartParseOptions{RerunWithDelete: true})
if err == nil {
t.Fatal("expected error from GetDocumentStorageAddress, got nil")
}
@@ -5030,7 +5053,7 @@ func TestIngest_CancelDoesNotDeleteIngestionTask(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
_, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
_, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusCancel),
Delete: true,
@@ -5059,7 +5082,7 @@ func TestIngest_CancelUnstartedDocument(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusCancel),
})
@@ -5087,7 +5110,7 @@ func TestIngest_CancelCompletedDocument(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusCancel),
})
@@ -5115,7 +5138,7 @@ func TestIngest_CancelUnstartedWithInFlightTask(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusCancel),
})
@@ -5147,7 +5170,7 @@ func TestIngest_CancelAgainWhenAlreadyStopped(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusCancel),
})
@@ -5172,7 +5195,7 @@ func TestIngest_CancelScheduledTask_TransitionsToStopped(t *testing.T) {
svc := testDocumentService(t)
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusCancel),
})
@@ -5231,7 +5254,7 @@ func TestIngest_CancelRunningTask_LifecycleToStopped(t *testing.T) {
ctx := documentTestContext(t.Context())
// 1. First cancel call: transitions RUNNING -> STOPPING
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusCancel),
})
@@ -5259,7 +5282,7 @@ func TestIngest_CancelRunningTask_LifecycleToStopped(t *testing.T) {
}
// 2. Duplicate cancel while in STOPPING must succeed as an idempotent no-op
code, err = svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err = svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: string(entity.TaskStatusCancel),
})
@@ -5311,7 +5334,7 @@ func TestIngest_DeleteOnlyCleansTasks(t *testing.T) {
svc := testDocumentService(t)
svc.purgeTaskState = func(context.Context, string) error { return nil }
ctx := documentTestContext(t.Context())
code, err := svc.Ingest(ctx, "user-1", &IngestDocumentRequest{
code, err := svc.Ingest(ctx, &IngestDocumentRequest{
DocIDs: []string{"doc-1"},
Run: "0",
Delete: true,
internal/service/document/document_version_pin_list_test.go
+1 -1
View File
@@ -85,7 +85,7 @@ func TestUpdateDatasetDocumentListReturnsPinnedVersion(t *testing.T) {
pid := canvasID
pv := "v-pinned"
pt := 2
_, code, err := svc.UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
_, code, err := svc.UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParseType: &pt,
PipelineID: &pid,
PipelineVersionID: &pv,
internal/service/document/document_version_pin_pipeline_change_test.go
+1 -1
View File
@@ -42,7 +42,7 @@ func TestUpdateDatasetDocument_PipelineChangeWithoutVersionRejected(t *testing.T
// Switch to canvasB, omitting pipeline_version_id entirely.
parseType := 2 // Pipeline mode
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{ParseType: &parseType, PipelineID: sptr(canvasB)},
map[string]bool{"pipeline_id": true})
if err == nil || code == common.CodeSuccess {
internal/service/document/document_version_pin_test.go
+1 -1
View File
@@ -95,7 +95,7 @@ func TestUpdateDatasetDocumentUsesPinnedPipelineVersionForParserConfig(t *testin
}
svc := testDocumentService(t)
_, code, err := svc.UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
_, code, err := svc.UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1", &UpdateDatasetDocumentRequest{
ParserConfig: map[string]any{"Parser:Shared": map[string]any{"k": "override_val"}},
}, map[string]bool{"parser_config": true})
if err != nil || code != common.CodeSuccess {
internal/service/document/document_version_pin_update_test.go
+9 -9
View File
@@ -91,7 +91,7 @@ func TestUpdateDatasetDocument_PersistsPinnedPipelineVersion(t *testing.T) {
seedDocVersion(t, db, canvasID, "v-new", 200)
bindDocCanvas(t, db, "kb-1", "doc-1", canvasID, "v-old")
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{PipelineVersionID: sptr("v-new")},
map[string]bool{"pipeline_version_id": true})
if err != nil || code != common.CodeSuccess {
@@ -117,7 +117,7 @@ func TestUpdateDatasetDocument_RejectsForeignPipelineVersion(t *testing.T) {
seedDocVersion(t, db, canvasB, "v-foreign", 100)
bindDocCanvas(t, db, "kb-1", "doc-1", canvasA, "v-a")
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{PipelineVersionID: sptr("v-foreign")},
map[string]bool{"pipeline_version_id": true})
if err == nil {
@@ -145,7 +145,7 @@ func TestUpdateDatasetDocument_RejectsMissingPipelineVersion(t *testing.T) {
seedDocVersion(t, db, canvasID, "v-a", 100)
bindDocCanvas(t, db, "kb-1", "doc-1", canvasID, "v-a")
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{PipelineVersionID: sptr("does-not-exist")},
map[string]bool{"pipeline_version_id": true})
if err == nil || code != common.CodeDataError {
@@ -167,7 +167,7 @@ func TestUpdateDatasetDocument_PreservesPipelineVersionWhenAbsent(t *testing.T)
seedDocVersion(t, db, canvasID, "v-old", 100)
bindDocCanvas(t, db, "kb-1", "doc-1", canvasID, "v-old")
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{}, map[string]bool{})
if err != nil || code != common.CodeSuccess {
t.Fatalf("UpdateDatasetDocument err=%v code=%d", err, code)
@@ -191,7 +191,7 @@ func TestUpdateDatasetDocument_EmptyPipelineVersionTracksLatest(t *testing.T) {
// C5/G3: an explicit EMPTY version in Pipeline mode is REJECTED (design
// section 3: empty version is no longer auto-backfilled to latest). The pin
// must be preserved and the save must fail with ErrPipelineVersionEmpty.
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{PipelineVersionID: sptr("")},
map[string]bool{"pipeline_version_id": true})
if err == nil || code == common.CodeSuccess {
@@ -211,7 +211,7 @@ func TestUpdateDatasetDocument_BuiltinIgnoresPipelineVersion(t *testing.T) {
insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5)
insertNamedTestDoc(t, "doc-1", "kb-1", "doc.txt", 10, 5)
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{PipelineVersionID: sptr("v-whatever")},
map[string]bool{"pipeline_version_id": true})
if err != nil || code != common.CodeSuccess {
@@ -234,7 +234,7 @@ func TestUpdateDatasetDocument_CleansParserConfigWithNewVersion(t *testing.T) {
seedDocVersion(t, db, canvasID, "v-new", 200)
bindDocCanvas(t, db, "kb-1", "doc-1", canvasID, "v-old")
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{
PipelineVersionID: sptr("v-new"),
ParserConfig: map[string]any{"Parser:Shared": map[string]any{"k": "override"}},
@@ -278,7 +278,7 @@ func TestUpdateDatasetDocument_VersionSwitchResetsParserConfigToNewDefault(t *te
t.Fatalf("seed parser_config: %v", err)
}
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{PipelineVersionID: sptr("v-new")},
map[string]bool{"pipeline_version_id": true})
if err != nil || code != common.CodeSuccess {
@@ -320,7 +320,7 @@ func TestUpdateDatasetDocument_SameVersionKeepsParserConfig(t *testing.T) {
t.Fatalf("seed parser_config: %v", err)
}
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "user-1", "kb-1", "doc-1",
_, code, err := testDocumentService(t).UpdateDatasetDocument(documentTestContext(t.Context()), "kb-1", "doc-1",
&UpdateDatasetDocumentRequest{PipelineVersionID: sptr("v-old")},
map[string]bool{"pipeline_version_id": true})
if err != nil || code != common.CodeSuccess {
internal/service/document/sync_document.go
+23 -2
View File
@@ -27,6 +27,7 @@ import (
"ragflow/internal/dao"
"ragflow/internal/entity"
"ragflow/internal/permission"
"ragflow/internal/service"
"ragflow/internal/storage"
"ragflow/internal/utility"
@@ -36,6 +37,18 @@ import (
var syncFilenameUnsafeRE = regexp.MustCompile(`[\\/:*?"<>|]+`)
func syncDocumentParseContext(ctx context.Context, taskContext service.SyncTaskContext) (context.Context, error) {
userID := strings.TrimSpace(taskContext.Connector.UserID)
if userID == "" {
return nil, fmt.Errorf("connector %s has no owner user_id", taskContext.Connector.ID)
}
tenantID := strings.TrimSpace(taskContext.Connector.TenantID)
if tenantID == "" {
return nil, fmt.Errorf("connector %s has no tenant_id", taskContext.Connector.ID)
}
return permission.WithIdentity(ctx, userID, tenantID), nil
}
// Upsert stores one normalized source document.
func (s *DocumentService) Upsert(ctx context.Context, input service.DocumentUpsertInput) (service.DocumentUpsertResult, error) {
if input.DocumentID == "" {
@@ -223,7 +236,11 @@ func (s *DocumentService) afterSyncDocumentUpsert(ctx context.Context, input ser
}
// do auto_parse if enabled
if input.AutoParse {
if err := s.StartParseDocuments(ctx, doc, &input.TaskContext.Knowledgebase, input.TaskContext.Connector.TenantID, StartParseOptions{RerunWithDelete: rerun}); err != nil {
parseCtx, err := syncDocumentParseContext(ctx, input.TaskContext)
if err != nil {
return err
}
if err := s.StartParseDocuments(parseCtx, doc, &input.TaskContext.Knowledgebase, StartParseOptions{RerunWithDelete: rerun}); err != nil {
return err
}
}
@@ -264,7 +281,11 @@ func (s *DocumentService) ensureSyncDocumentPostWrite(ctx context.Context, input
if task != nil {
return nil
}
return s.StartParseDocuments(ctx, doc, &input.TaskContext.Knowledgebase, input.TaskContext.Connector.TenantID, StartParseOptions{})
parseCtx, err := syncDocumentParseContext(ctx, input.TaskContext)
if err != nil {
return err
}
return s.StartParseDocuments(parseCtx, doc, &input.TaskContext.Knowledgebase, StartParseOptions{})
}
// updateSyncDocumentFile updates the file-manager row linked to a synced document.
internal/service/document/sync_document_test.go
+72 -1
View File
@@ -11,6 +11,7 @@ import (
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/entity"
"ragflow/internal/permission"
"ragflow/internal/service"
"ragflow/internal/storage"
syncerconnector "ragflow/internal/syncer/connector"
@@ -47,7 +48,7 @@ func setupSyncDocumentTest(t *testing.T, existing bool) (*DocumentService, *gorm
t.Cleanup(func() { factory.SetStorage(original) })
svc := testDocumentService(t)
input := service.DocumentUpsertInput{
TaskContext: service.SyncTaskContext{Knowledgebase: *kb, Connector: entity.Connector{TenantID: "tenant-1"}},
TaskContext: service.SyncTaskContext{Knowledgebase: *kb, Connector: entity.Connector{UserID: "user-1", TenantID: "tenant-1"}},
SourceType: "github_connector-1",
DocumentID: "doc-sync",
SourceDocument: syncerconnector.SourceDocument{
@@ -395,8 +396,78 @@ func TestSyncDocumentAcknowledgesParseEnqueueAfterCancellation(t *testing.T) {
if len(publisher.messages) != 1 {
t.Fatalf("parse publications = %d, want 1", len(publisher.messages))
}
task, err := svc.ingestionTaskDAO.GetByID(t.Context(), db, publisher.messages[0].TaskID)
if err != nil {
t.Fatalf("load enqueued ingestion task: %v", err)
}
if task.UserID != input.TaskContext.Connector.UserID {
t.Fatalf("ingestion task user id = %q, want connector owner %q", task.UserID, input.TaskContext.Connector.UserID)
}
result, err = svc.Upsert(t.Context(), input)
if err != nil || result.Action != service.DocumentActionSkipped || len(publisher.messages) != 1 {
t.Fatalf("completed revision retried enqueue: result = %+v err = %v publications = %d", result, err, len(publisher.messages))
}
}
func TestSyncDocumentAutoParseRejectsMissingConnectorOwner(t *testing.T) {
svc, db, _, input := setupSyncDocumentTest(t, false)
input.AutoParse = true
input.TaskContext.Connector.UserID = ""
publisher := &recordingTaskPublisher{}
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
if _, err := svc.Upsert(t.Context(), input); err == nil || !strings.Contains(err.Error(), "no owner user_id") {
t.Fatalf("Upsert error = %v, want missing connector owner", err)
}
if len(publisher.messages) != 0 {
t.Fatalf("parse publications = %d, want none", len(publisher.messages))
}
var taskCount int64
if err := db.Model(&entity.IngestionTask{}).Where("document_id = ?", input.DocumentID).Count(&taskCount).Error; err != nil {
t.Fatal(err)
}
if taskCount != 0 {
t.Fatalf("ingestion tasks = %d, want none", taskCount)
}
}
func TestSyncDocumentReenqueuesParseWithConnectorOwner(t *testing.T) {
svc, db, _, input := setupSyncDocumentTest(t, true)
publisher := &recordingTaskPublisher{}
svc.ingestionTaskSvc.SetTaskPublisher(publisher)
if _, err := svc.Upsert(t.Context(), input); err != nil {
t.Fatalf("sync document before parse retry: %v", err)
}
input.AutoParse = true
result, err := svc.Upsert(t.Context(), input)
if err != nil || result.Action != service.DocumentActionSkipped {
t.Fatalf("parse retry = %+v, err = %v", result, err)
}
if len(publisher.messages) != 1 {
t.Fatalf("parse publications = %d, want 1", len(publisher.messages))
}
task, err := svc.ingestionTaskDAO.GetByID(t.Context(), db, publisher.messages[0].TaskID)
if err != nil {
t.Fatalf("load enqueued ingestion task: %v", err)
}
if task.UserID != input.TaskContext.Connector.UserID {
t.Fatalf("ingestion task user id = %q, want connector owner %q", task.UserID, input.TaskContext.Connector.UserID)
}
}
func TestSyncDocumentParseContextUsesConnectorIdentity(t *testing.T) {
ctx := permission.WithIdentity(t.Context(), "request-user", "request-tenant")
parseCtx, err := syncDocumentParseContext(ctx, service.SyncTaskContext{
Connector: entity.Connector{ID: "connector-1", UserID: "connector-owner", TenantID: "connector-tenant"},
})
if err != nil {
t.Fatal(err)
}
if got := permission.UserIDFromContext(parseCtx); got != "connector-owner" {
t.Fatalf("parse context user id = %q, want connector-owner", got)
}
if got := permission.TenantIDFromContext(parseCtx); got != "connector-tenant" {
t.Fatalf("parse context tenant id = %q, want connector-tenant", got)
}
}