From 2a54bb2d45dd41b9569aa2dad753f9cecc8d36cc Mon Sep 17 00:00:00 2001 From: wolfkill <27876473+wolfkill@users.noreply.github.com> Date: Fri, 29 May 2026 11:26:25 +0800 Subject: [PATCH] fix: avoid stale knowledge records on upload failure --- .../application/service/knowledge_create.go | 26 +- .../service/knowledge_create_test.go | 231 ++++++++++++++++++ 2 files changed, 244 insertions(+), 13 deletions(-) create mode 100644 internal/application/service/knowledge_create_test.go diff --git a/internal/application/service/knowledge_create.go b/internal/application/service/knowledge_create.go index 2f1c628f..49102e15 100644 --- a/internal/application/service/knowledge_create.go +++ b/internal/application/service/knowledge_create.go @@ -179,9 +179,10 @@ func (s *knowledgeService) CreateKnowledgeFromFile(ctx context.Context, return nil, werrors.NewValidationError("文件名包含非法字符") } - // Create knowledge record - logger.Info(ctx, "Creating knowledge record") + // Prepare knowledge record + logger.Info(ctx, "Preparing knowledge record") knowledge := &types.Knowledge{ + ID: uuid.New().String(), TenantID: tenantID, KnowledgeBaseID: kbID, TagID: tagID, // 设置分类ID，用于知识分类管理 @@ -199,25 +200,24 @@ func (s *knowledgeService) CreateKnowledgeFromFile(ctx context.Context, EmbeddingModelID: kb.EmbeddingModelID, Metadata: metadataJSON, } - // Save knowledge record to database - logger.Info(ctx, "Saving knowledge record to database") - if err := s.repo.CreateKnowledge(ctx, knowledge); err != nil { - logger.Errorf(ctx, "Failed to create knowledge record, ID: %s, error: %v", knowledge.ID, err) - return nil, err - } + // Save the file to storage (use KB-level storage engine if configured) logger.Infof(ctx, "Saving file, knowledge ID: %s", knowledge.ID) - filePath, err := s.resolveFileService(ctx, kb).SaveFile(ctx, file, knowledge.TenantID, knowledge.ID) + fileSvc := s.resolveFileService(ctx, kb) + filePath, err := fileSvc.SaveFile(ctx, file, knowledge.TenantID, knowledge.ID) if err != nil { logger.Errorf(ctx, "Failed to save file, knowledge ID: %s, error: %v", knowledge.ID, err) return nil, err } knowledge.FilePath = filePath - // Update knowledge record with file path - logger.Info(ctx, "Updating knowledge record with file path") - if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { - logger.Errorf(ctx, "Failed to update knowledge with file path, ID: %s, error: %v", knowledge.ID, err) + // Save knowledge record to database after the file is safely stored. + logger.Info(ctx, "Saving knowledge record to database") + if err := s.repo.CreateKnowledge(ctx, knowledge); err != nil { + logger.Errorf(ctx, "Failed to create knowledge record, ID: %s, error: %v", knowledge.ID, err) + if deleteErr := fileSvc.DeleteFile(ctx, filePath); deleteErr != nil { + logger.Errorf(ctx, "Failed to delete saved file after knowledge creation failed, path: %s, error: %v", filePath, deleteErr) + } return nil, err } diff --git a/internal/application/service/knowledge_create_test.go b/internal/application/service/knowledge_create_test.go new file mode 100644 index 00000000..7f6be154 --- /dev/null +++ b/internal/application/service/knowledge_create_test.go @@ -0,0 +1,231 @@ +package service + +import ( + "bytes" + "context" + "errors" + "io" + "mime/multipart" + "net/http/httptest" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/hibiken/asynq" + "github.com/stretchr/testify/require" +) + +type createKnowledgeFileRepoStub struct { + interfaces.KnowledgeRepository + + createCalls int + createErr error + createdKnowledge *types.Knowledge +} + +func (r *createKnowledgeFileRepoStub) CheckKnowledgeExists( + ctx context.Context, + tenantID uint64, + kbID string, + params *types.KnowledgeCheckParams, +) (bool, *types.Knowledge, error) { + return false, nil, nil +} + +func (r *createKnowledgeFileRepoStub) CreateKnowledge(ctx context.Context, knowledge *types.Knowledge) error { + r.createCalls++ + copied := *knowledge + r.createdKnowledge = &copied + return r.createErr +} + +type createKnowledgeFileKBServiceStub struct { + interfaces.KnowledgeBaseService + + kb *types.KnowledgeBase +} + +func (s *createKnowledgeFileKBServiceStub) GetKnowledgeBaseByID( + ctx context.Context, + id string, +) (*types.KnowledgeBase, error) { + return s.kb, nil +} + +type createKnowledgeFileServiceStub struct { + saveErr error + saveCalls int + savedWithKnowledgeID string + deleteCalls int + deletedPath string +} + +func (s *createKnowledgeFileServiceStub) CheckConnectivity(ctx context.Context) error { + return nil +} + +func (s *createKnowledgeFileServiceStub) SaveFile( + ctx context.Context, + file *multipart.FileHeader, + tenantID uint64, + knowledgeID string, +) (string, error) { + s.saveCalls++ + s.savedWithKnowledgeID = knowledgeID + if s.saveErr != nil { + return "", s.saveErr + } + return "stored/" + knowledgeID, nil +} + +func (s *createKnowledgeFileServiceStub) SaveBytes( + ctx context.Context, + data []byte, + tenantID uint64, + fileName string, + temp bool, +) (string, error) { + return "", errors.New("not implemented") +} + +func (s *createKnowledgeFileServiceStub) GetFile(ctx context.Context, filePath string) (io.ReadCloser, error) { + return nil, errors.New("not implemented") +} + +func (s *createKnowledgeFileServiceStub) GetFileURL(ctx context.Context, filePath string) (string, error) { + return "", errors.New("not implemented") +} + +func (s *createKnowledgeFileServiceStub) DeleteFile(ctx context.Context, filePath string) error { + s.deleteCalls++ + s.deletedPath = filePath + return nil +} + +type createKnowledgeTaskEnqueuerStub struct { + calls int +} + +func (s *createKnowledgeTaskEnqueuerStub) Enqueue( + task *asynq.Task, + opts ...asynq.Option, +) (*asynq.TaskInfo, error) { + s.calls++ + return &asynq.TaskInfo{ID: "task-1", Queue: "default"}, nil +} + +func TestCreateKnowledgeFromFileDoesNotPersistWhenStorageSaveFails(t *testing.T) { + t.Parallel() + + repo := &createKnowledgeFileRepoStub{} + fileSvc := &createKnowledgeFileServiceStub{saveErr: errors.New("storage unavailable")} + svc := &knowledgeService{ + repo: repo, + kbService: &createKnowledgeFileKBServiceStub{kb: &types.KnowledgeBase{ID: "kb-1"}}, + fileSvc: fileSvc, + } + + knowledge, err := svc.CreateKnowledgeFromFile( + newCreateKnowledgeFileContext(), + "kb-1", + newMultipartFileHeader(t, "doc.txt", "hello"), + nil, + nil, + "", + "", + "", + ) + + require.Error(t, err) + require.Nil(t, knowledge) + require.Equal(t, 1, fileSvc.saveCalls) + require.Zero(t, repo.createCalls) +} + +func TestCreateKnowledgeFromFilePersistsStoredFilePathOnCreate(t *testing.T) { + t.Parallel() + + repo := &createKnowledgeFileRepoStub{} + fileSvc := &createKnowledgeFileServiceStub{} + task := &createKnowledgeTaskEnqueuerStub{} + svc := &knowledgeService{ + repo: repo, + kbService: &createKnowledgeFileKBServiceStub{kb: &types.KnowledgeBase{ID: "kb-1"}}, + fileSvc: fileSvc, + task: task, + } + + knowledge, err := svc.CreateKnowledgeFromFile( + newCreateKnowledgeFileContext(), + "kb-1", + newMultipartFileHeader(t, "doc.txt", "hello"), + nil, + nil, + "", + "", + "", + ) + + require.NoError(t, err) + require.NotNil(t, knowledge) + require.Equal(t, 1, fileSvc.saveCalls) + require.NotEmpty(t, fileSvc.savedWithKnowledgeID) + require.Equal(t, fileSvc.savedWithKnowledgeID, knowledge.ID) + require.Equal(t, 1, repo.createCalls) + require.NotNil(t, repo.createdKnowledge) + require.Equal(t, "stored/"+knowledge.ID, repo.createdKnowledge.FilePath) + require.Equal(t, 1, task.calls) +} + +func TestCreateKnowledgeFromFileDeletesStoredFileWhenCreateFails(t *testing.T) { + t.Parallel() + + repo := &createKnowledgeFileRepoStub{createErr: errors.New("database unavailable")} + fileSvc := &createKnowledgeFileServiceStub{} + svc := &knowledgeService{ + repo: repo, + kbService: &createKnowledgeFileKBServiceStub{kb: &types.KnowledgeBase{ID: "kb-1"}}, + fileSvc: fileSvc, + } + + knowledge, err := svc.CreateKnowledgeFromFile( + newCreateKnowledgeFileContext(), + "kb-1", + newMultipartFileHeader(t, "doc.txt", "hello"), + nil, + nil, + "", + "", + "", + ) + + require.EqualError(t, err, "database unavailable") + require.Nil(t, knowledge) + require.Equal(t, 1, fileSvc.saveCalls) + require.Equal(t, 1, repo.createCalls) + require.Equal(t, 1, fileSvc.deleteCalls) + require.Equal(t, "stored/"+fileSvc.savedWithKnowledgeID, fileSvc.deletedPath) +} + +func newCreateKnowledgeFileContext() context.Context { + ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1)) + ctx = context.WithValue(ctx, types.TenantInfoContextKey, &types.Tenant{}) + return ctx +} + +func newMultipartFileHeader(t *testing.T, filename string, content string) *multipart.FileHeader { + t.Helper() + + var body bytes.Buffer + writer := multipart.NewWriter(&body) + part, err := writer.CreateFormFile("file", filename) + require.NoError(t, err) + _, err = part.Write([]byte(content)) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + req := httptest.NewRequest("POST", "/", &body) + req.Header.Set("Content-Type", writer.FormDataContentType()) + require.NoError(t, req.ParseMultipartForm(1024)) + return req.MultipartForm.File["file"][0] +}