diff --git a/internal/application/service/resumable_upload.go b/internal/application/service/resumable_upload.go index c103e4e..d84a708 100644 --- a/internal/application/service/resumable_upload.go +++ b/internal/application/service/resumable_upload.go @@ -161,18 +161,27 @@ func NewResumableUploadService(db *gorm.DB, bases interfaces.KnowledgeBaseServic } return accesspolicy.WithSnapshot(ctx, snapshot), nil } + var uploads *resumable.Service ingest := func(ctx context.Context, row *resumable.Upload, path string) (int64, types.JSON, bool, error) { if row.Kind == "digital-human" { source, selection, e := humans.readResumableMaterial(ctx, row.Name, path, row.Size) if e != nil { return 0, nil, false, e } - result, e := json.Marshal(uploadResult{Source: &source, Selection: selection}) + if e = uploads.SaveMaterial(row.ID, source); e != nil { + return 0, nil, false, e + } + // jsonb keeps only small metadata; full text lives in {id}.material. + result, e := json.Marshal(uploadResult{ + Selection: selection, + Source: &types.DigitalHumanSource{Name: source.Name}, + }) return row.Size, result, true, e } return ingestResumableKnowledge(ctx, knowledge, row, path) } - uploads, e := resumable.NewService(db, os.Getenv("RESUMABLE_UPLOAD_DIR"), authorize, workerContext, ingest) + var e error + uploads, e = resumable.NewService(db, os.Getenv("RESUMABLE_UPLOAD_DIR"), authorize, workerContext, ingest) if e != nil { return nil, e } @@ -192,19 +201,40 @@ func NewResumableUploadService(db *gorm.DB, bases interfaces.KnowledgeBaseServic return nil, digitalHumanInputError("材料尚未解析完成、已过期或不属于此数字人") } var result uploadResult - if json.Unmarshal(row.Result, &result) != nil || result.Source == nil { + if json.Unmarshal(row.Result, &result) != nil { return nil, digitalHumanInputError("材料解析结果无效") } if result.Selection != "full" { return nil, digitalHumanInputError("材料“" + row.Name + "”是旧版原文节选，请重新上传以完整蒸馏") } - sources = append(sources, *result.Source) + source, e := resolveDigitalHumanUploadSource(ctx, humans, uploads, row, result) + if e != nil { + return nil, e + } + sources = append(sources, source) } return sources, nil } return uploads, nil } +// resolveDigitalHumanUploadSource prefers the file-backed material, then legacy +// jsonb text, then re-parses the raw upload and backfills .material. +func resolveDigitalHumanUploadSource(ctx context.Context, humans *DigitalHumanService, uploads *resumable.Service, row *resumable.Upload, result uploadResult) (types.DigitalHumanSource, error) { + if source, e := uploads.LoadMaterial(row.ID); e == nil { + return source, nil + } + if result.Source != nil && strings.TrimSpace(result.Source.Text) != "" { + return *result.Source, nil + } + source, _, e := humans.readResumableMaterial(ctx, row.Name, uploads.FilePath(row.ID), row.Size) + if e != nil { + return types.DigitalHumanSource{}, digitalHumanInputError("材料解析结果无效") + } + _ = uploads.SaveMaterial(row.ID, source) + return source, nil +} + func textUpload(name string) bool { switch strings.ToLower(filepath.Ext(name)) { case ".txt", ".md", ".markdown", ".srt", ".vtt", ".csv", ".json", ".html", ".htm": diff --git a/internal/application/service/resumable_upload_material_test.go b/internal/application/service/resumable_upload_material_test.go new file mode 100644 index 0000000..d974d4d --- /dev/null +++ b/internal/application/service/resumable_upload_material_test.go @@ -0,0 +1,66 @@ +package service + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/resumable" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +func TestResolveDigitalHumanUploadSourcePrefersMaterialThenLegacyThenReparse(t *testing.T) { + db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "mat.db")), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + t.Cleanup(func() { _ = sqlDB.Close() }) + require.NoError(t, db.AutoMigrate(&resumable.Upload{})) + root := t.TempDir() + uploads, err := resumable.NewService(db, root, func(context.Context, string, string) error { return nil }, + func(ctx context.Context, _ *resumable.Upload) (context.Context, error) { return ctx, nil }, + func(context.Context, *resumable.Upload, string) (int64, types.JSON, bool, error) { + return 0, nil, true, nil + }) + require.NoError(t, err) + humans := &DigitalHumanService{} + + row := &resumable.Upload{ + ID: "u1", Name: "note.txt", Size: 11, Kind: "digital-human", + State: "completed", ExpiresAt: time.Now().Add(time.Hour), + Result: types.JSON(`{"selection":"full","source":{"name":"note.txt"}}`), + } + require.NoError(t, os.WriteFile(uploads.FilePath(row.ID), []byte("from-upload"), 0600)) + + // 1) material file wins + require.NoError(t, uploads.SaveMaterial(row.ID, types.DigitalHumanSource{Name: "note.txt", Text: "from-material"})) + source, err := resolveDigitalHumanUploadSource(context.Background(), humans, uploads, row, uploadResult{Selection: "full", Source: &types.DigitalHumanSource{Name: "note.txt"}}) + require.NoError(t, err) + require.Equal(t, "from-material", source.Text) + + // 2) legacy jsonb text when material missing + require.NoError(t, os.Remove(filepath.Join(root, row.ID+".material"))) + legacy := uploadResult{Selection: "full", Source: &types.DigitalHumanSource{Name: "note.txt", Text: "from-jsonb"}} + source, err = resolveDigitalHumanUploadSource(context.Background(), humans, uploads, row, legacy) + require.NoError(t, err) + require.Equal(t, "from-jsonb", source.Text) + + // 3) reparse upload and backfill material + source, err = resolveDigitalHumanUploadSource(context.Background(), humans, uploads, row, uploadResult{Selection: "full", Source: &types.DigitalHumanSource{Name: "note.txt"}}) + require.NoError(t, err) + require.Equal(t, "from-upload", source.Text) + loaded, err := uploads.LoadMaterial(row.ID) + require.NoError(t, err) + require.Equal(t, "from-upload", loaded.Text) + + meta, err := json.Marshal(uploadResult{Selection: "full", Source: &types.DigitalHumanSource{Name: "note.txt"}}) + require.NoError(t, err) + require.NotContains(t, string(meta), "from-upload") +} diff --git a/internal/resumable/handler.go b/internal/resumable/handler.go index f48dd53..783be66 100644 --- a/internal/resumable/handler.go +++ b/internal/resumable/handler.go @@ -47,12 +47,8 @@ func respond(c *gin.Context, status int, value any, err error) { func (s *Service) Register(group *gin.RouterGroup) { r := group.Group("/uploads") r.GET("", func(c *gin.Context) { - rows, e := s.List(c.Request.Context(), c.Query("kind"), c.Query("target_id")) - if e == nil && c.Query("view") == "summary" { - for i := range rows { - rows[i].Result = uploadResultSummary(rows[i].Result) - } - } + summary := c.Query("view") == "summary" + rows, e := s.List(c.Request.Context(), c.Query("kind"), c.Query("target_id"), summary) respond(c, 200, rows, e) }) r.POST("", func(c *gin.Context) { diff --git a/internal/resumable/material_test.go b/internal/resumable/material_test.go new file mode 100644 index 0000000..5b258ba --- /dev/null +++ b/internal/resumable/material_test.go @@ -0,0 +1,75 @@ +package resumable + +import ( + "encoding/json" + "os" + "strings" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/require" +) + +func TestMaterialSaveLoadAndCancelRemovesBothFiles(t *testing.T) { + s := uploadService(t) + ctx := uploadContext(1, "alice") + data := []byte("material-body") + u := create(t, s, data) + _, err := s.Put(ctx, u.ID, 0, int64(len(data)), sum(data), strings.NewReader(string(data))) + require.NoError(t, err) + + source := types.DigitalHumanSource{Name: "fixture.txt", Text: strings.Repeat("正文内容-", 200)} + require.NoError(t, s.SaveMaterial(u.ID, source)) + loaded, err := s.LoadMaterial(u.ID) + require.NoError(t, err) + require.Equal(t, source, loaded) + _, err = os.Stat(s.materialPath(u.ID)) + require.NoError(t, err) + + require.NoError(t, s.Cancel(ctx, u.ID)) + _, err = os.Stat(s.path(u.ID)) + require.True(t, os.IsNotExist(err)) + _, err = os.Stat(s.materialPath(u.ID)) + require.True(t, os.IsNotExist(err)) +} + +func TestListSummaryDoesNotTransferEmbeddedSourceText(t *testing.T) { + s := uploadService(t) + secret := strings.Repeat("huge-source-payload-", 8000) + fat, err := json.Marshal(map[string]any{ + "selection": "full", + "source": map[string]string{"name": "big.txt", "text": secret}, + }) + require.NoError(t, err) + require.NoError(t, s.db.Create(&Upload{ + ID: "fat-1", TenantID: 1, OwnerID: "alice", Kind: "digital-human", TargetID: "dh-1", + Name: "big.txt", Size: 10, Offset: 10, IngestOffset: 10, State: "completed", + Options: types.JSON(`{}`), Result: fat, ExpiresAt: time.Now().Add(time.Hour), + }).Error) + + rows, err := s.List(uploadContext(1, "alice"), "", "", true) + require.NoError(t, err) + require.Len(t, rows, 1) + require.JSONEq(t, `{"selection":"full"}`, string(rows[0].Result)) + require.NotContains(t, string(rows[0].Result), "huge-source-payload") + require.Less(t, len(rows[0].Result), 64) + + full, err := s.List(uploadContext(1, "alice"), "", "", false) + require.NoError(t, err) + require.Contains(t, string(full[0].Result), secret) +} + +func TestCleanupRemovesMaterialFile(t *testing.T) { + s := uploadService(t) + ctx := uploadContext(1, "alice") + data := []byte("x") + u := create(t, s, data) + _, err := s.Put(ctx, u.ID, 0, 1, sum(data), strings.NewReader("x")) + require.NoError(t, err) + require.NoError(t, s.SaveMaterial(u.ID, types.DigitalHumanSource{Name: "x.txt", Text: "hello"})) + require.NoError(t, s.db.Model(&Upload{}).Where("id=?", u.ID).Update("expires_at", time.Now().Add(-time.Hour)).Error) + require.NoError(t, s.Cleanup(ctx)) + _, err = os.Stat(s.materialPath(u.ID)) + require.True(t, os.IsNotExist(err)) +} diff --git a/internal/resumable/service.go b/internal/resumable/service.go index 57cbbb0..da30892 100644 --- a/internal/resumable/service.go +++ b/internal/resumable/service.go @@ -4,6 +4,7 @@ import ( "context" "crypto/sha256" "encoding/hex" + "encoding/json" "errors" "fmt" "io" @@ -61,7 +62,68 @@ func validHash(v string) bool { b, e := hex.DecodeString(v) return e == nil && len(b) == 32 && strings.ToLower(v) == v } -func (s *Service) path(id string) string { return filepath.Join(s.root, id+".upload") } +func (s *Service) path(id string) string { return filepath.Join(s.root, id+".upload") } +func (s *Service) materialPath(id string) string { return filepath.Join(s.root, id+".material") } + +// FilePath is the durable raw upload bytes on disk. +func (s *Service) FilePath(id string) string { return s.path(id) } + +// SaveMaterial stores parsed digital-human source text beside the upload file +// so result jsonb only needs small metadata. +func (s *Service) SaveMaterial(id string, source types.DigitalHumanSource) error { + if id == "" || strings.TrimSpace(source.Name) == "" || source.Text == "" { + return ErrInvalid + } + data, err := json.Marshal(source) + if err != nil { + return err + } + tmp := s.materialPath(id) + ".tmp" + if err = os.WriteFile(tmp, data, 0600); err != nil { + return err + } + return os.Rename(tmp, s.materialPath(id)) +} + +// LoadMaterial reads the file-backed digital-human source for an upload. +func (s *Service) LoadMaterial(id string) (types.DigitalHumanSource, error) { + data, err := os.ReadFile(s.materialPath(id)) + if err != nil { + return types.DigitalHumanSource{}, err + } + var source types.DigitalHumanSource + if err = json.Unmarshal(data, &source); err != nil { + return types.DigitalHumanSource{}, err + } + if strings.TrimSpace(source.Name) == "" || source.Text == "" { + return types.DigitalHumanSource{}, ErrInvalid + } + return source, nil +} + +func (s *Service) removeStorage(id string) error { + for _, p := range []string{s.path(id), s.materialPath(id), s.materialPath(id) + ".tmp"} { + if e := os.Remove(p); e != nil && !errors.Is(e, os.ErrNotExist) { + return e + } + } + return nil +} + +// listSelect projects columns for List. Summary mode never loads full result +// jsonb payloads — only a tiny selection object when present. +func (s *Service) listSelect(summary bool) string { + cols := `id, tenant_id, owner_id, kind, target_id, name, size, "offset", chunk_size, sha256, fingerprint, state, error, options, ingest_offset, part, expires_at, lease_until, lease_token, created_at, updated_at` + if !summary { + return cols + ", result" + } + switch s.db.Dialector.Name() { + case "postgres": + return cols + `, CASE WHEN result->>'selection' IN ('full','sampled') THEN jsonb_build_object('selection', result->'selection') ELSE NULL END AS result` + default: + return cols + `, CASE WHEN json_extract(result, '$.selection') IN ('full','sampled') THEN json_object('selection', json_extract(result, '$.selection')) ELSE NULL END AS result` + } +} func (s *Service) Create(ctx context.Context, in CreateInput) (*Upload, error) { t, u, e := identity(ctx) @@ -143,13 +205,13 @@ func (s *Service) Chunks(ctx context.Context, id string, offset int64) ([]Chunk, return chunks, err } -func (s *Service) List(ctx context.Context, kind, target string) ([]Upload, error) { +func (s *Service) List(ctx context.Context, kind, target string, summary bool) ([]Upload, error) { q, e := s.scoped(ctx) if e != nil { return nil, e } rows := []Upload{} - q = q.Where("state NOT IN ('cancelled','expired')") + q = q.Select(s.listSelect(summary)).Where("state NOT IN ('cancelled','expired')") if kind != "" { q = q.Where("kind=?", kind) } @@ -161,8 +223,17 @@ func (s *Service) List(ctx context.Context, kind, target string) ([]Upload, erro return nil, e } visible := make([]Upload, 0, len(rows)) + type authKey struct{ kind, target string } + authCache := map[authKey]error{} for _, row := range rows { - if s.authorize != nil && s.authorize(ctx, row.Kind, row.TargetID) == nil { + if s.authorize == nil { + continue + } + key := authKey{row.Kind, row.TargetID} + if _, ok := authCache[key]; !ok { + authCache[key] = s.authorize(ctx, row.Kind, row.TargetID) + } + if authCache[key] == nil { visible = append(visible, row) } } @@ -296,7 +367,7 @@ func (s *Service) Cancel(ctx context.Context, id string) error { if row.State == "verifying" || row.State == "processing" { return ErrConflict } - if e := os.Remove(s.path(id)); e != nil && !errors.Is(e, os.ErrNotExist) { + if e := s.removeStorage(id); e != nil { return e } if e := tx.Where("upload_id=?", id).Delete(&Chunk{}).Error; e != nil { @@ -320,7 +391,7 @@ func (s *Service) Cleanup(ctx context.Context) error { if current.ExpiresAt.After(time.Now()) || (current.LeaseUntil != nil && current.LeaseUntil.After(time.Now())) { return nil } - if e := os.Remove(s.path(row.ID)); e != nil && !errors.Is(e, os.ErrNotExist) { + if e := s.removeStorage(row.ID); e != nil { return e } if e := tx.Where("upload_id=?", row.ID).Delete(&Chunk{}).Error; e != nil { diff --git a/internal/resumable/service_test.go b/internal/resumable/service_test.go index fbffdcd..87cb892 100644 --- a/internal/resumable/service_test.go +++ b/internal/resumable/service_test.go @@ -196,7 +196,7 @@ func TestUploadRevocationHidesEvidenceAndStopsWritesButAllowsCleanup(t *testing. require.ErrorIs(t, e, ErrForbidden) _, e = s.Finalize(ctx, u.ID) require.ErrorIs(t, e, ErrForbidden) - rows, e := s.List(ctx, "", "") + rows, e := s.List(ctx, "", "", false) require.NoError(t, e) require.Empty(t, rows) require.NoError(t, s.Cancel(ctx, u.ID))