461 lines
23 KiB
Go
461 lines
23 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"crypto/sha256"
|
||
|
|
"encoding/hex"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"io"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||
|
|
"github.com/Tencent/WeKnora/internal/application/repository"
|
||
|
|
filesvc "github.com/Tencent/WeKnora/internal/application/service/file"
|
||
|
|
"github.com/Tencent/WeKnora/internal/knowledgeversion"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
"gorm.io/driver/sqlite"
|
||
|
|
"gorm.io/gorm"
|
||
|
|
)
|
||
|
|
|
||
|
|
type versionCopyKBService struct {
|
||
|
|
interfaces.KnowledgeBaseService
|
||
|
|
bases map[string]*types.KnowledgeBase
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s versionCopyKBService) CopyKnowledgeBase(_ context.Context, source, target string) (*types.KnowledgeBase, *types.KnowledgeBase, error) {
|
||
|
|
return s.bases[source], s.bases[target], nil
|
||
|
|
}
|
||
|
|
func (s versionCopyKBService) GetKnowledgeBaseByID(_ context.Context, id string) (*types.KnowledgeBase, error) {
|
||
|
|
return s.bases[id], nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type versionCopyTenantRepo struct {
|
||
|
|
interfaces.TenantRepository
|
||
|
|
tenant *types.Tenant
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r versionCopyTenantRepo) GetTenantByID(context.Context, uint64) (*types.Tenant, error) {
|
||
|
|
return r.tenant, nil
|
||
|
|
}
|
||
|
|
func (r versionCopyTenantRepo) AdjustStorageUsed(context.Context, uint64, int64) error { return nil }
|
||
|
|
|
||
|
|
type versionCopyCreateFailureRepo struct{ interfaces.KnowledgeRepository }
|
||
|
|
|
||
|
|
func (r versionCopyCreateFailureRepo) CreateVersionKnowledge(context.Context, *types.Knowledge) error {
|
||
|
|
return errors.New("document insert failed")
|
||
|
|
}
|
||
|
|
|
||
|
|
type versionCopyIndex struct {
|
||
|
|
knowledgeID, kbID, tagID string
|
||
|
|
enabled bool
|
||
|
|
}
|
||
|
|
type versionCopyEngine struct {
|
||
|
|
parentChildRetrieveEngine
|
||
|
|
indices map[string]versionCopyIndex
|
||
|
|
copyErr error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (e *versionCopyEngine) CopyIndices(_ context.Context, source string, docs, chunks map[string]string, target string, _ int, _ string) error {
|
||
|
|
if e.copyErr != nil {
|
||
|
|
return e.copyErr
|
||
|
|
}
|
||
|
|
for sourceID, targetID := range chunks {
|
||
|
|
index := e.indices[sourceID]
|
||
|
|
if index.kbID != source {
|
||
|
|
return errors.New("wrong source index")
|
||
|
|
}
|
||
|
|
index.knowledgeID, index.kbID = docs[index.knowledgeID], target
|
||
|
|
e.indices[targetID] = index
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
func (e *versionCopyEngine) BatchUpdateChunkTagID(_ context.Context, tags map[string]string) error {
|
||
|
|
for id, tag := range tags {
|
||
|
|
index := e.indices[id]
|
||
|
|
index.tagID = tag
|
||
|
|
e.indices[id] = index
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
func (e *versionCopyEngine) BatchUpdateChunkEnabledStatus(_ context.Context, enabled map[string]bool) error {
|
||
|
|
for id, status := range enabled {
|
||
|
|
index := e.indices[id]
|
||
|
|
index.enabled = status
|
||
|
|
e.indices[id] = index
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type versionCopyFiles struct {
|
||
|
|
interfaces.FileService
|
||
|
|
created, deleted []string
|
||
|
|
failSaveAt int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (f *versionCopyFiles) SaveBytes(ctx context.Context, data []byte, tenantID uint64, name string, temporary bool) (string, error) {
|
||
|
|
if f.failSaveAt > 0 && len(f.created)+1 == f.failSaveAt {
|
||
|
|
return "", errors.New("image copy failed")
|
||
|
|
}
|
||
|
|
path, err := f.FileService.SaveBytes(ctx, data, tenantID, name, temporary)
|
||
|
|
if err == nil {
|
||
|
|
f.created = append(f.created, path)
|
||
|
|
}
|
||
|
|
return path, err
|
||
|
|
}
|
||
|
|
func (f *versionCopyFiles) CopyFile(ctx context.Context, source string, tenantID uint64, id string) (string, error) {
|
||
|
|
path, err := f.FileService.CopyFile(ctx, source, tenantID, id)
|
||
|
|
if err == nil {
|
||
|
|
f.created = append(f.created, path)
|
||
|
|
}
|
||
|
|
return path, err
|
||
|
|
}
|
||
|
|
func (f *versionCopyFiles) DeleteFile(ctx context.Context, path string) error {
|
||
|
|
if err := ctx.Err(); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
f.deleted = append(f.deleted, path)
|
||
|
|
return f.FileService.DeleteFile(ctx, path)
|
||
|
|
}
|
||
|
|
|
||
|
|
type versionCopyFixture struct {
|
||
|
|
service *knowledgeService
|
||
|
|
db *gorm.DB
|
||
|
|
ctx context.Context
|
||
|
|
files *versionCopyFiles
|
||
|
|
engine *versionCopyEngine
|
||
|
|
source *types.Knowledge
|
||
|
|
sourceChunk *types.Chunk
|
||
|
|
raw, image, metadataImage string
|
||
|
|
}
|
||
|
|
|
||
|
|
func newVersionCopyFixture(t *testing.T) *versionCopyFixture {
|
||
|
|
t.Helper()
|
||
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||
|
|
require.NoError(t, err)
|
||
|
|
sqlDB, err := db.DB()
|
||
|
|
require.NoError(t, err)
|
||
|
|
sqlDB.SetMaxOpenConns(1)
|
||
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
||
|
|
require.NoError(t, db.AutoMigrate(&types.Tenant{}, &types.Knowledge{}, &types.Chunk{}, &types.KnowledgeTag{}, &types.KnowledgeTagRelation{}))
|
||
|
|
require.NoError(t, db.Exec("CREATE TABLE knowledge_bases (id TEXT PRIMARY KEY, tenant_id INTEGER NOT NULL, deleted_at DATETIME)").Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id, tenant_id) VALUES ('source-kb', 1), ('target-kb', 1)").Error)
|
||
|
|
require.NoError(t, db.Create(&types.Tenant{ID: 1, Name: "version-copy-test"}).Error)
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
|
||
|
|
ctx = knowledgeversion.WithWorker(ctx)
|
||
|
|
files := &versionCopyFiles{FileService: filesvc.NewLocalFileService(t.TempDir(), "")}
|
||
|
|
raw, err := files.FileService.SaveBytes(ctx, []byte("source raw bytes"), 1, "raw.txt", false)
|
||
|
|
require.NoError(t, err)
|
||
|
|
image, err := files.FileService.SaveBytes(ctx, []byte("chunk image bytes"), 1, "image.png", false)
|
||
|
|
require.NoError(t, err)
|
||
|
|
metadataImage, err := files.FileService.SaveBytes(ctx, []byte("metadata image bytes"), 1, "metadata.png", false)
|
||
|
|
require.NoError(t, err)
|
||
|
|
metadata, _ := json.Marshal(map[string]any{"content": "![manual](" + image + ")", "preview": map[string]any{"url": metadataImage}, "number": json.Number("9007199254740993"), "source_file": raw})
|
||
|
|
source := &types.Knowledge{ID: "source-doc", TenantID: 1, KnowledgeBaseID: "source-kb", Type: types.KnowledgeTypeManual, Title: "source", ParseStatus: types.ParseStatusCompleted, EnableStatus: "enabled", EmbeddingModelID: "embedding", FilePath: raw, Metadata: metadata, CustomMetadata: types.JSON(`{"description":"keep me"}`), FolderPath: "folder/subfolder"}
|
||
|
|
require.NoError(t, db.Create(source).Error)
|
||
|
|
tag := &types.KnowledgeTag{ID: "source-tag", SeqID: 1, TenantID: 1, KnowledgeBaseID: "source-kb", Name: "classification"}
|
||
|
|
require.NoError(t, db.Create(tag).Error)
|
||
|
|
repo := repository.NewKnowledgeRepository(db)
|
||
|
|
require.NoError(t, repo.SetKnowledgeTags(ctx, source.ID, []string{tag.ID}))
|
||
|
|
imageInfo, _ := json.Marshal([]*types.ImageInfo{{URL: image, OriginalURL: image}})
|
||
|
|
chunk := &types.Chunk{ID: "source-chunk", TenantID: 1, KnowledgeID: source.ID, KnowledgeBaseID: source.KnowledgeBaseID, TagID: tag.ID, ChunkType: types.ChunkTypeText, Content: "edited ![image](" + image + ")", SourceContent: "original ![image](" + image + ")", ImageInfo: string(imageInfo), ContentRevision: 3, ContextHeader: "# Header", IndexStatus: "ready", IsEnabled: false, Metadata: types.JSON(`{"custom":123}`)}
|
||
|
|
chunkRepo := repository.NewChunkRepository(db)
|
||
|
|
require.NoError(t, chunkRepo.CreateChunks(ctx, []*types.Chunk{chunk}))
|
||
|
|
chunk.IsEnabled = false
|
||
|
|
chunk.Flags = 0
|
||
|
|
require.NoError(t, chunkRepo.UpdateChunk(ctx, chunk))
|
||
|
|
bases := map[string]*types.KnowledgeBase{"source-kb": {ID: "source-kb", TenantID: 1, Type: types.KnowledgeBaseTypeDocument}, "target-kb": {ID: "target-kb", TenantID: 1, Type: types.KnowledgeBaseTypeDocument, VersionState: "copying"}}
|
||
|
|
engine := &versionCopyEngine{indices: map[string]versionCopyIndex{chunk.ID: {knowledgeID: source.ID, kbID: source.KnowledgeBaseID, tagID: tag.ID, enabled: false}}}
|
||
|
|
tenant := &types.Tenant{ID: 1, RetrieverEngines: types.RetrieverEngines{Engines: []types.RetrieverEngineParams{{RetrieverType: types.VectorRetrieverType, RetrieverEngineType: types.PostgresRetrieverEngineType}}}}
|
||
|
|
svc := &knowledgeService{repo: repo, chunkRepo: chunkRepo, tagRepo: repository.NewKnowledgeTagRepository(db), kbService: versionCopyKBService{bases: bases}, tenantRepo: versionCopyTenantRepo{tenant: tenant}, fileSvc: files, modelService: parentChildModelService{embedder: parentChildEmbedder{}}, retrieveEngine: parentChildRetrieveRegistry{engine: engine}}
|
||
|
|
return &versionCopyFixture{service: svc, db: db, ctx: ctx, files: files, engine: engine, source: source, sourceChunk: chunk, raw: raw, image: image, metadataImage: metadataImage}
|
||
|
|
}
|
||
|
|
|
||
|
|
func readVersionCopyFile(t *testing.T, file interfaces.FileService, path string) string {
|
||
|
|
t.Helper()
|
||
|
|
r, err := file.GetFile(context.Background(), path)
|
||
|
|
require.NoError(t, err)
|
||
|
|
defer r.Close()
|
||
|
|
data, err := io.ReadAll(r)
|
||
|
|
require.NoError(t, err)
|
||
|
|
return string(data)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyOwnsFilesMetadataChunksTagsAndIndices(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
var progress [][2]int
|
||
|
|
require.NoError(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", func(done, total int) error { progress = append(progress, [2]int{done, total}); return nil }))
|
||
|
|
require.Equal(t, [][2]int{{0, 1}, {1, 1}}, progress)
|
||
|
|
docs, err := f.service.repo.ListKnowledgeByKnowledgeBaseID(f.ctx, 1, "target-kb")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, docs, 1)
|
||
|
|
doc := docs[0]
|
||
|
|
require.NotEqual(t, f.source.ID, doc.ID)
|
||
|
|
require.NotEqual(t, f.raw, doc.FilePath)
|
||
|
|
require.Equal(t, f.source.FolderPath, doc.FolderPath)
|
||
|
|
require.Equal(t, types.ParseStatusCompleted, doc.ParseStatus)
|
||
|
|
require.Equal(t, "source raw bytes", readVersionCopyFile(t, f.files, doc.FilePath))
|
||
|
|
require.NotContains(t, string(doc.Metadata), f.image)
|
||
|
|
require.NotContains(t, string(doc.Metadata), f.metadataImage)
|
||
|
|
require.NotContains(t, string(doc.Metadata), f.raw)
|
||
|
|
require.Contains(t, string(doc.Metadata), "9007199254740993")
|
||
|
|
tags, err := f.service.repo.GetKnowledgeTags(f.ctx, []string{doc.ID})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, tags[doc.ID], 1)
|
||
|
|
require.NotEqual(t, "source-tag", tags[doc.ID][0].ID)
|
||
|
|
var chunks []*types.Chunk
|
||
|
|
require.NoError(t, f.db.Where("knowledge_base_id = ?", "target-kb").Find(&chunks).Error)
|
||
|
|
require.Len(t, chunks, 1)
|
||
|
|
chunk := chunks[0]
|
||
|
|
require.NotEqual(t, f.sourceChunk.ID, chunk.ID)
|
||
|
|
require.NotContains(t, chunk.Content, f.image)
|
||
|
|
require.NotContains(t, chunk.SourceContent, f.image)
|
||
|
|
require.Equal(t, 3, chunk.ContentRevision)
|
||
|
|
require.Equal(t, "# Header", chunk.ContextHeader)
|
||
|
|
require.False(t, chunk.IsEnabled)
|
||
|
|
require.Zero(t, chunk.Flags)
|
||
|
|
require.Equal(t, tags[doc.ID][0].ID, chunk.TagID)
|
||
|
|
require.Equal(t, versionCopyIndex{knowledgeID: doc.ID, kbID: doc.KnowledgeBaseID, tagID: chunk.TagID, enabled: false}, f.engine.indices[chunk.ID])
|
||
|
|
original, err := f.service.repo.GetKnowledgeByID(f.ctx, 1, f.source.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.JSONEq(t, string(f.source.Metadata), string(original.Metadata))
|
||
|
|
require.Equal(t, f.raw, original.FilePath)
|
||
|
|
// Destroy the source's real files and rows. Every destination object and
|
||
|
|
// independently keyed retrieval record must remain usable.
|
||
|
|
for _, path := range []string{f.raw, f.image, f.metadataImage} {
|
||
|
|
require.NoError(t, f.files.FileService.DeleteFile(f.ctx, path))
|
||
|
|
}
|
||
|
|
require.NoError(t, f.db.Delete(&types.Chunk{}, "id = ?", f.sourceChunk.ID).Error)
|
||
|
|
require.NoError(t, f.db.Delete(&types.Knowledge{}, "id = ?", f.source.ID).Error)
|
||
|
|
delete(f.engine.indices, f.sourceChunk.ID)
|
||
|
|
require.Equal(t, "source raw bytes", readVersionCopyFile(t, f.files, doc.FilePath))
|
||
|
|
var images []*types.ImageInfo
|
||
|
|
require.NoError(t, json.Unmarshal([]byte(chunk.ImageInfo), &images))
|
||
|
|
require.Len(t, images, 2)
|
||
|
|
for _, img := range images {
|
||
|
|
require.NotEmpty(t, readVersionCopyFile(t, f.files, img.URL))
|
||
|
|
}
|
||
|
|
require.Contains(t, f.engine.indices, chunk.ID)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyFailureCleansOnlyNewImages(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
f.engine.copyErr = errors.New("index unavailable")
|
||
|
|
err := f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil)
|
||
|
|
require.ErrorContains(t, err, "index unavailable")
|
||
|
|
require.Len(t, f.files.deleted, 2)
|
||
|
|
for _, source := range []string{f.raw, f.image, f.metadataImage} {
|
||
|
|
require.NotContains(t, f.files.deleted, source)
|
||
|
|
require.NotEmpty(t, readVersionCopyFile(t, f.files, source))
|
||
|
|
}
|
||
|
|
for _, path := range f.files.deleted {
|
||
|
|
_, err := f.files.GetFile(f.ctx, path)
|
||
|
|
require.Error(t, err)
|
||
|
|
}
|
||
|
|
docs, err := f.service.repo.ListKnowledgeByKnowledgeBaseID(f.ctx, 1, "target-kb")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, docs, 1)
|
||
|
|
require.Equal(t, types.ParseStatusProcessing, docs[0].ParseStatus)
|
||
|
|
fields, inventory, err := readKnowledgeVersionInventory(docs[0].Metadata, docs[0].ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, fields, 1)
|
||
|
|
require.Len(t, inventory.Objects, 2)
|
||
|
|
for _, object := range inventory.Objects {
|
||
|
|
require.Contains(t, f.files.deleted, object.Reference)
|
||
|
|
}
|
||
|
|
require.Equal(t, "source raw bytes", readVersionCopyFile(t, f.files, docs[0].FilePath))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyPartialImageFailureCleansEarlierObjects(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
f.files.failSaveAt = 3
|
||
|
|
require.ErrorContains(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil), "image copy failed")
|
||
|
|
require.Len(t, f.files.deleted, 1)
|
||
|
|
for _, source := range []string{f.raw, f.image, f.metadataImage} {
|
||
|
|
require.NotContains(t, f.files.deleted, source)
|
||
|
|
require.NotEmpty(t, readVersionCopyFile(t, f.files, source))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyAccountsRetainedFailedDocument(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
require.NoError(t, f.db.Model(f.source).Update("storage_size", 64).Error)
|
||
|
|
f.engine.copyErr = errors.New("index unavailable")
|
||
|
|
require.ErrorContains(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil), "index unavailable")
|
||
|
|
var tenant types.Tenant
|
||
|
|
require.NoError(t, f.db.First(&tenant, 1).Error)
|
||
|
|
require.EqualValues(t, 64, tenant.StorageUsed)
|
||
|
|
docs, err := f.service.repo.ListKnowledgeByKnowledgeBaseID(f.ctx, 1, "target-kb")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, docs, 1)
|
||
|
|
require.EqualValues(t, 64, docs[0].StorageSize)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyRollsBackUnaccountedDocument(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
require.NoError(t, f.db.Delete(&types.Tenant{}, 1).Error)
|
||
|
|
require.ErrorContains(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil), "workspace not found")
|
||
|
|
docs, err := f.service.repo.ListKnowledgeByKnowledgeBaseID(f.ctx, 1, "target-kb")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Empty(t, docs)
|
||
|
|
require.Len(t, f.files.deleted, 1)
|
||
|
|
require.Equal(t, "source raw bytes", readVersionCopyFile(t, f.files, f.raw))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyIncludesFAQAndRejectsUnsupportedChunks(t *testing.T) {
|
||
|
|
for _, kind := range []string{types.ChunkTypeFAQ, types.ChunkTypeEntity} {
|
||
|
|
t.Run(kind, func(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
require.NoError(t, f.db.Model(&types.Chunk{}).Where("id = ?", f.sourceChunk.ID).Update("chunk_type", kind).Error)
|
||
|
|
err := f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil)
|
||
|
|
if kind == types.ChunkTypeEntity {
|
||
|
|
require.ErrorContains(t, err, "尚不支持")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
require.NoError(t, err)
|
||
|
|
var count int64
|
||
|
|
require.NoError(t, f.db.Model(&types.Chunk{}).Where("knowledge_base_id = ? AND chunk_type = ?", "target-kb", kind).Count(&count).Error)
|
||
|
|
require.EqualValues(t, 1, count)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyPreservesStoredParentAndChildLinks(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
parent := &types.Chunk{ID: "source-parent", TenantID: 1, KnowledgeID: f.source.ID, KnowledgeBaseID: f.source.KnowledgeBaseID, ChunkType: types.ChunkTypeParentText, Status: int(types.ChunkStatusStored), Content: "Full parent context", IndexStatus: "ready", NextChunkID: f.sourceChunk.ID}
|
||
|
|
require.NoError(t, f.service.chunkRepo.CreateChunks(f.ctx, []*types.Chunk{parent}))
|
||
|
|
f.sourceChunk.ParentChunkID, f.sourceChunk.PreChunkID = parent.ID, parent.ID
|
||
|
|
require.NoError(t, f.service.chunkRepo.UpdateChunk(f.ctx, f.sourceChunk))
|
||
|
|
f.engine.indices[parent.ID] = versionCopyIndex{knowledgeID: f.source.ID, kbID: f.source.KnowledgeBaseID, enabled: true}
|
||
|
|
require.NoError(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil))
|
||
|
|
var chunks []*types.Chunk
|
||
|
|
require.NoError(t, f.db.Where("knowledge_base_id = ?", "target-kb").Find(&chunks).Error)
|
||
|
|
require.Len(t, chunks, 2)
|
||
|
|
var copiedParent, copiedChild *types.Chunk
|
||
|
|
for _, chunk := range chunks {
|
||
|
|
if chunk.ChunkType == types.ChunkTypeParentText {
|
||
|
|
copiedParent = chunk
|
||
|
|
} else {
|
||
|
|
copiedChild = chunk
|
||
|
|
}
|
||
|
|
}
|
||
|
|
require.NotNil(t, copiedParent)
|
||
|
|
require.NotNil(t, copiedChild)
|
||
|
|
require.NotEqual(t, parent.ID, copiedParent.ID)
|
||
|
|
require.Equal(t, int(types.ChunkStatusStored), copiedParent.Status)
|
||
|
|
require.Equal(t, copiedParent.ID, copiedChild.ParentChunkID)
|
||
|
|
require.Equal(t, copiedParent.ID, copiedChild.PreChunkID)
|
||
|
|
require.Equal(t, copiedChild.ID, copiedParent.NextChunkID)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyFailsInsteadOfSilentlyOmittingUnknownChunk(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
require.NoError(t, f.db.Model(&types.Chunk{}).Where("id = ?", f.sourceChunk.ID).Update("chunk_type", "future_chunk").Error)
|
||
|
|
// Keep this document free of image attachments so the final count guard
|
||
|
|
// verifies omissions independently of the attachment inventory guard.
|
||
|
|
require.NoError(t, f.db.Model(&types.Knowledge{}).Where("id = ?", f.source.ID).Update("metadata", nil).Error)
|
||
|
|
require.ErrorContains(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil), "未完整复制")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyRawFileCleanedWhenDocumentInsertFails(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
f.service.repo = versionCopyCreateFailureRepo{KnowledgeRepository: f.service.repo}
|
||
|
|
require.ErrorContains(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil), "document insert failed")
|
||
|
|
require.Len(t, f.files.created, 1)
|
||
|
|
require.Equal(t, f.files.created, f.files.deleted)
|
||
|
|
_, err := f.files.GetFile(f.ctx, f.files.created[0])
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Equal(t, "source raw bytes", readVersionCopyFile(t, f.files, f.raw))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestCleanupCopiedObjectsSurvivesCancellation(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
path, err := f.files.SaveBytes(f.ctx, []byte("discard copy"), 1, "discard.png", false)
|
||
|
|
require.NoError(t, err)
|
||
|
|
ctx, cancel := context.WithCancel(f.ctx)
|
||
|
|
cancel()
|
||
|
|
cleanupCopiedObjects(ctx, f.files, []string{path})
|
||
|
|
require.Equal(t, []string{path}, f.files.deleted)
|
||
|
|
_, err = f.files.GetFile(f.ctx, path)
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Equal(t, "chunk image bytes", readVersionCopyFile(t, f.files, f.image))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopyGuards(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
require.ErrorContains(t, f.service.CloneKnowledgeVersion(context.Background(), "source-kb", "target-kb", nil), "isolated worker")
|
||
|
|
require.ErrorContains(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "source-kb", nil), "isolated worker")
|
||
|
|
require.NoError(t, f.db.Model(&types.Knowledge{}).Where("id = ?", f.source.ID).Update("parse_status", types.ParseStatusProcessing).Error)
|
||
|
|
require.ErrorContains(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil), "尚未完成索引")
|
||
|
|
require.Empty(t, f.files.created)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeCopyLineagePreventsExportLaundering(t *testing.T) {
|
||
|
|
for _, global := range []bool{false, true} {
|
||
|
|
snapshot := &accesspolicy.Snapshot{NoExport: map[string]bool{}, Publications: map[string]bool{}}
|
||
|
|
if global {
|
||
|
|
snapshot.Publications["knowledge:restricted"] = true
|
||
|
|
} else {
|
||
|
|
snapshot.NoExport["knowledge:restricted"] = true
|
||
|
|
}
|
||
|
|
ctx := accesspolicy.WithSnapshot(context.Background(), snapshot)
|
||
|
|
require.ErrorContains(t, checkKnowledgeCopyLineage(ctx, "restricted"), "不能复制或移动")
|
||
|
|
require.NoError(t, checkKnowledgeCopyLineage(ctx, "ordinary"))
|
||
|
|
require.NoError(t, checkKnowledgeCopyLineage(knowledgeversion.WithWorker(ctx), "restricted"))
|
||
|
|
require.Error(t, (&knowledgeService{}).CloneKnowledgeBase(ctx, "restricted", "ordinary"))
|
||
|
|
require.Error(t, (&knowledgeService{}).moveOneKnowledge(ctx, "document", &types.KnowledgeBase{ID: "restricted"}, nil, "reparse"))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeVersionCopiesGraphEvidenceAndRehashesOwnedImages(t *testing.T) {
|
||
|
|
f := newVersionCopyFixture(t)
|
||
|
|
require.NoError(t, f.db.AutoMigrate(&types.KnowledgeBase{}, &knowledgeversion.GraphRecord{}, &knowledgeversion.GraphSource{}))
|
||
|
|
bases := f.service.kbService.(versionCopyKBService).bases
|
||
|
|
bases["target-kb"].IsTemporary = true
|
||
|
|
for _, kb := range bases {
|
||
|
|
require.NoError(t, f.db.Save(kb).Error)
|
||
|
|
}
|
||
|
|
record := knowledgeversion.GraphRecord{TenantID: 1, KnowledgeBaseID: "source-kb", SourceTable: "knowledgeentity", LegacyID: 7, GraphVersionID: 2, Payload: types.JSON(`{"id":7,"canonical_name":"Original evidence"}`), SourceSnapshotSHA256: "archive-digest"}
|
||
|
|
require.NoError(t, f.db.Create(&record).Error)
|
||
|
|
oldHash := sha256.Sum256([]byte(f.sourceChunk.Content))
|
||
|
|
hash := hex.EncodeToString(oldHash[:])
|
||
|
|
badHash := "edited-after-graph-extraction"
|
||
|
|
evidence := []knowledgeversion.GraphSource{
|
||
|
|
{TenantID: 1, KnowledgeBaseID: "source-kb", SourceChunkID: "mapped", SourceFile: "original.jsonl", Excerpt: "preserved excerpt", NativeKnowledgeID: &f.source.ID, NativeChunkID: &f.sourceChunk.ID, NativeContentSHA256: &hash, MappingState: "exact_content_hash"},
|
||
|
|
{TenantID: 1, KnowledgeBaseID: "source-kb", SourceChunkID: "stale", Excerpt: "old evidence", NativeKnowledgeID: &f.source.ID, NativeChunkID: &f.sourceChunk.ID, NativeContentSHA256: &badHash, MappingState: "exact_content_hash"},
|
||
|
|
{TenantID: 1, KnowledgeBaseID: "source-kb", SourceChunkID: "unmapped", SourceFile: "absent.jsonl", Excerpt: "unmapped evidence", MappingState: "unmapped"},
|
||
|
|
}
|
||
|
|
require.NoError(t, f.db.Create(&evidence).Error)
|
||
|
|
require.NoError(t, f.service.CloneKnowledgeVersion(f.ctx, "source-kb", "target-kb", nil))
|
||
|
|
var copied knowledgeversion.GraphRecord
|
||
|
|
require.NoError(t, f.db.Where("knowledge_base_id=?", "target-kb").First(&copied).Error)
|
||
|
|
require.Equal(t, record.Payload, copied.Payload)
|
||
|
|
require.Equal(t, record.LegacyID, copied.LegacyID)
|
||
|
|
require.Equal(t, record.SourceSnapshotSHA256, copied.SourceSnapshotSHA256)
|
||
|
|
var sources []knowledgeversion.GraphSource
|
||
|
|
require.NoError(t, f.db.Where("knowledge_base_id=?", "target-kb").Order("source_chunk_id").Find(&sources).Error)
|
||
|
|
require.Len(t, sources, 3)
|
||
|
|
require.NotNil(t, sources[0].NativeChunkID)
|
||
|
|
require.NotEqual(t, f.sourceChunk.ID, *sources[0].NativeChunkID)
|
||
|
|
var copiedChunk types.Chunk
|
||
|
|
require.NoError(t, f.db.First(&copiedChunk, "id=?", *sources[0].NativeChunkID).Error)
|
||
|
|
newHash := sha256.Sum256([]byte(copiedChunk.Content))
|
||
|
|
require.Equal(t, hex.EncodeToString(newHash[:]), *sources[0].NativeContentSHA256)
|
||
|
|
require.NotEqual(t, hash, *sources[0].NativeContentSHA256, "owned image reference must be rehashed")
|
||
|
|
require.Nil(t, sources[1].NativeChunkID)
|
||
|
|
require.Equal(t, "stale", sources[1].MappingState)
|
||
|
|
require.Nil(t, sources[2].NativeChunkID)
|
||
|
|
require.Equal(t, "absent.jsonl", sources[2].SourceFile)
|
||
|
|
require.Equal(t, "unmapped evidence", sources[2].Excerpt)
|
||
|
|
// The source graph survives independently and cross-workspace copies fail.
|
||
|
|
var count int64
|
||
|
|
require.NoError(t, f.db.Model(&knowledgeversion.GraphRecord{}).Where("knowledge_base_id=?", "source-kb").Count(&count).Error)
|
||
|
|
require.EqualValues(t, 1, count)
|
||
|
|
repo := f.service.repo.(interface {
|
||
|
|
CopyVersionGraph(context.Context, string, string, map[string]string, map[string]string) error
|
||
|
|
})
|
||
|
|
require.Error(t, repo.CopyVersionGraph(context.WithValue(f.ctx, types.TenantIDContextKey, uint64(2)), "source-kb", "target-kb", nil, nil))
|
||
|
|
}
|