fix: avoid stale knowledge records on upload failure
This commit is contained in:
wolfkill
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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]
|
||||
}
|
||||
Reference in New Issue
Block a user