yawan
f7fc992f57
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
308 lines
15 KiB
Go
308 lines
15 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Tencent/WeKnora/internal/models/requestbudget"
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
)
|
|
|
|
func openRecoveryBudgetSQLite(t *testing.T, path string) *gorm.DB {
|
|
t.Helper()
|
|
db, err := gorm.Open(sqlite.Open(path), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
|
|
require.NoError(t, err)
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err)
|
|
sqlDB.SetMaxOpenConns(1)
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
|
return db
|
|
}
|
|
|
|
func seedRecoveryBudget(t *testing.T, db *gorm.DB, max int64) types.WikiRecoveryBudget {
|
|
t.Helper()
|
|
require.NoError(t, db.AutoMigrate(&types.WikiRecoveryBudget{}, &types.TaskPendingOp{}))
|
|
for _, ddl := range []string{
|
|
`CREATE TABLE IF NOT EXISTS knowledge_bases (id text PRIMARY KEY, tenant_id bigint, deleted_at timestamp)`,
|
|
`CREATE TABLE IF NOT EXISTS knowledges (id text PRIMARY KEY, tenant_id bigint, knowledge_base_id text, deleted_at timestamp)`,
|
|
`ALTER TABLE knowledge_bases ADD COLUMN summary_model_id text DEFAULT ''`,
|
|
`ALTER TABLE knowledge_bases ADD COLUMN embedding_model_id text DEFAULT ''`,
|
|
`ALTER TABLE knowledge_bases ADD COLUMN wiki_config text DEFAULT '{}'`,
|
|
`ALTER TABLE knowledge_bases ADD COLUMN indexing_strategy text DEFAULT '{}'`,
|
|
`ALTER TABLE knowledges ADD COLUMN title text DEFAULT ''`,
|
|
`CREATE TABLE chunks (id text PRIMARY KEY, tenant_id bigint, knowledge_base_id text, knowledge_id text, content text, content_revision integer, chunk_index integer, start_at integer, end_at integer, chunk_type text, parent_chunk_id text, image_info text, is_enabled boolean, deleted_at timestamp)`,
|
|
`CREATE TABLE models (id text PRIMARY KEY, tenant_id bigint, name text, type text, source text, status text, parameters text, deleted_at timestamp)`,
|
|
} {
|
|
require.NoError(t, db.Exec(ddl).Error)
|
|
}
|
|
require.NoError(t, db.Exec(`INSERT INTO knowledge_bases (id, tenant_id, summary_model_id, embedding_model_id, wiki_config, indexing_strategy) VALUES ('kb',7,'model','embed','{"synthesis_model_id":"model"}','{"wiki_enabled":true}')`).Error)
|
|
require.NoError(t, db.Exec(`INSERT INTO knowledges (id,tenant_id,knowledge_base_id,title) VALUES ('doc',7,'kb','title')`).Error)
|
|
require.NoError(t, db.Exec(`INSERT INTO chunks (id,tenant_id,knowledge_base_id,knowledge_id,content,content_revision,chunk_index,start_at,end_at,chunk_type,parent_chunk_id,image_info,is_enabled) VALUES ('chunk',7,'kb','doc','source',0,0,0,6,'text','','',?)`, true).Error)
|
|
for _, id := range []string{"model", "embed"} {
|
|
require.NoError(t, db.Exec(`INSERT INTO models (id,tenant_id,name,type,source,status,parameters) VALUES (?,7,?,'KnowledgeQA','remote','active','{}')`, id, id).Error)
|
|
}
|
|
if db.Dialector.Name() == "postgres" {
|
|
for _, ddl := range []string{
|
|
`ALTER TABLE knowledge_bases ALTER COLUMN wiki_config DROP DEFAULT`,
|
|
`ALTER TABLE knowledge_bases ALTER COLUMN wiki_config TYPE jsonb USING wiki_config::jsonb`,
|
|
`ALTER TABLE knowledge_bases ALTER COLUMN indexing_strategy DROP DEFAULT`,
|
|
`ALTER TABLE knowledge_bases ALTER COLUMN indexing_strategy TYPE jsonb USING indexing_strategy::jsonb`,
|
|
`ALTER TABLE models ALTER COLUMN parameters TYPE jsonb USING parameters::jsonb`,
|
|
} {
|
|
require.NoError(t, db.Exec(ddl).Error)
|
|
}
|
|
}
|
|
row := types.WikiRecoveryBudget{TenantID: 7, DeadLetterID: 10, AttemptNumber: 5, PendingOpID: 42, KnowledgeBaseID: "kb", KnowledgeID: "doc", ModelID: "model", EmbeddingModelID: "embed", ReviewedRelease: "fixture", ReviewedAPIImage: "sha256:fixture", MaxProviderRequests: max, DeadlineAt: time.Now().UTC().Add(time.Hour).Truncate(time.Microsecond), Status: "active", PausedPayload: "[]", VerifiedProgress: "[]"}
|
|
var err error
|
|
row.SourceFingerprint, row.ModelFingerprint, err = WikiRecoveryFingerprints(db, row)
|
|
require.NoError(t, err)
|
|
require.NoError(t, db.Create(&row).Error)
|
|
require.NoError(t, db.Exec(`CREATE TABLE wiki_dead_letter_recoveries (tenant_id bigint, dead_letter_id bigint, attempt_number integer, pending_op_id bigint, knowledge_base_id text, knowledge_id text, model_id text, deployed_fix_review text)`).Error)
|
|
review, err := json.Marshal(map[string]any{"tenant_id": row.TenantID, "dead_letter_id": row.DeadLetterID, "max_attempt": row.AttemptNumber, "budget_version": 1,
|
|
"knowledge_base_id": row.KnowledgeBaseID, "knowledge_id": row.KnowledgeID, "model_id": row.ModelID, "embedding_model_id": row.EmbeddingModelID,
|
|
"source_fingerprint": row.SourceFingerprint, "model_fingerprint": row.ModelFingerprint, "fix_release": row.ReviewedRelease, "api_image_id": row.ReviewedAPIImage,
|
|
"max_provider_requests": row.MaxProviderRequests, "deadline_at": row.DeadlineAt})
|
|
require.NoError(t, err)
|
|
require.NoError(t, db.Exec(`INSERT INTO wiki_dead_letter_recoveries VALUES (?,?,?,?,?,?,?,?)`, row.TenantID, row.DeadLetterID, row.AttemptNumber, row.PendingOpID, row.KnowledgeBaseID, row.KnowledgeID, row.ModelID, string(review)).Error)
|
|
return row
|
|
}
|
|
|
|
func TestWikiRecoveryBudgetLastSlotPostgresIsAtomic(t *testing.T) {
|
|
db := setupKBDeletePostgres(t)
|
|
row := seedRecoveryBudget(t, db, 1)
|
|
repo := WikiRecoveryBudgetRepository{DB: db}
|
|
var wg sync.WaitGroup
|
|
var accepted atomic.Int32
|
|
for i := 0; i < 24; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
if repo.Reserve(context.Background(), row, "model", "candidate") == nil {
|
|
accepted.Add(1)
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
require.Equal(t, int32(1), accepted.Load())
|
|
saved, err := repo.Load(context.Background(), 7, row.Identity())
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), saved.ReservedProviderRequests)
|
|
require.Equal(t, "stopped_budget", saved.Status)
|
|
require.Equal(t, "request_limit", saved.StopReason)
|
|
require.Equal(t, row.DeadlineAt, saved.DeadlineAt.UTC())
|
|
}
|
|
|
|
func TestWikiRecoveryBudgetRestartSharesAllPhasesAndDeadline(t *testing.T) {
|
|
path := filepath.Join(t.TempDir(), "budget.sqlite")
|
|
db := openRecoveryBudgetSQLite(t, path)
|
|
row := seedRecoveryBudget(t, db, 3)
|
|
repo := WikiRecoveryBudgetRepository{DB: db}
|
|
require.NoError(t, repo.Reserve(context.Background(), row, "model", "candidate"))
|
|
require.NoError(t, repo.Reserve(context.Background(), row, "embed", "taxonomy"))
|
|
sqlDB, _ := db.DB()
|
|
require.NoError(t, sqlDB.Close())
|
|
repo.DB = openRecoveryBudgetSQLite(t, path)
|
|
saved, err := repo.Load(context.Background(), 7, row.Identity())
|
|
require.NoError(t, err)
|
|
require.Equal(t, row.DeadlineAt, saved.DeadlineAt)
|
|
require.Equal(t, int64(2), saved.ReservedProviderRequests)
|
|
require.NoError(t, repo.Reserve(context.Background(), *saved, "model", "finalize"))
|
|
require.ErrorIs(t, repo.Reserve(context.Background(), *saved, "model", "retry"), ErrWikiRecoveryBudgetStopped)
|
|
reloaded, err := repo.Load(context.Background(), 7, row.Identity())
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(3), reloaded.ReservedProviderRequests)
|
|
require.Equal(t, row.DeadlineAt, reloaded.DeadlineAt)
|
|
require.Equal(t, "request_limit", reloaded.StopReason)
|
|
}
|
|
|
|
func TestWikiRecoveryBudgetScopeAndLiveConfigurationFailClosed(t *testing.T) {
|
|
for _, tc := range []string{"tenant", "pending_op", "model", "request_model", "source", "model_config", "expired", "extended_limit", "extended_deadline", "edited_persisted_limit", "edited_persisted_deadline", "missing_receipt", "receipt_scope"} {
|
|
t.Run(tc, func(t *testing.T) {
|
|
db := openRecoveryBudgetSQLite(t, filepath.Join(t.TempDir(), "budget.sqlite"))
|
|
row := seedRecoveryBudget(t, db, 3)
|
|
original := row
|
|
model := "model"
|
|
switch tc {
|
|
case "tenant":
|
|
row.TenantID++
|
|
case "pending_op":
|
|
row.PendingOpID++
|
|
case "model":
|
|
row.ModelID = "foreign"
|
|
case "request_model":
|
|
model = "foreign"
|
|
case "source":
|
|
require.NoError(t, db.Exec("UPDATE chunks SET content = 'changed'").Error)
|
|
case "model_config":
|
|
require.NoError(t, db.Exec("UPDATE models SET parameters = '{\"changed\":true}'").Error)
|
|
case "expired":
|
|
row.DeadlineAt = time.Now().Add(-time.Second)
|
|
require.NoError(t, db.Model(&types.WikiRecoveryBudget{}).Where("tenant_id = 7").Update("deadline_at", row.DeadlineAt).Error)
|
|
case "extended_limit":
|
|
row.MaxProviderRequests++
|
|
case "extended_deadline":
|
|
row.DeadlineAt = row.DeadlineAt.Add(time.Hour)
|
|
case "edited_persisted_limit":
|
|
row.MaxProviderRequests++
|
|
require.NoError(t, db.Model(&types.WikiRecoveryBudget{}).Where("tenant_id = 7").Update("max_provider_requests", row.MaxProviderRequests).Error)
|
|
case "edited_persisted_deadline":
|
|
row.DeadlineAt = row.DeadlineAt.Add(time.Hour)
|
|
require.NoError(t, db.Model(&types.WikiRecoveryBudget{}).Where("tenant_id = 7").Update("deadline_at", row.DeadlineAt).Error)
|
|
case "missing_receipt":
|
|
require.NoError(t, db.Exec("DELETE FROM wiki_dead_letter_recoveries").Error)
|
|
case "receipt_scope":
|
|
require.NoError(t, db.Exec("UPDATE wiki_dead_letter_recoveries SET knowledge_id = 'foreign'").Error)
|
|
}
|
|
repo := WikiRecoveryBudgetRepository{DB: db}
|
|
require.ErrorIs(t, repo.Reserve(context.Background(), row, model, "candidate"), ErrWikiRecoveryBudgetStopped)
|
|
saved, err := repo.Load(context.Background(), 7, original.Identity())
|
|
require.NoError(t, err)
|
|
require.Zero(t, saved.ReservedProviderRequests)
|
|
if tc == "expired" {
|
|
require.Equal(t, "deadline", saved.StopReason)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestWikiRecoveryBudgetArchiveExactRowsAndKeepEvidence(t *testing.T) {
|
|
db := openRecoveryBudgetSQLite(t, filepath.Join(t.TempDir(), "budget.sqlite"))
|
|
row := seedRecoveryBudget(t, db, 1)
|
|
repo := WikiRecoveryBudgetRepository{DB: db}
|
|
identity := row.Identity()
|
|
for _, item := range []struct {
|
|
id int64
|
|
tenant uint64
|
|
kind string
|
|
identity types.WikiRecoveryIdentity
|
|
}{
|
|
{42, 7, types.TypeWikiIngest, identity}, {43, 7, types.TypeWikiFinalize, identity},
|
|
{44, 7, types.TypeWikiFinalize, types.WikiRecoveryIdentity{}}, {45, 8, types.TypeWikiFinalize, identity},
|
|
{46, 7, types.TypeWikiIngest, identity},
|
|
} {
|
|
body, _ := json.Marshal(item.identity)
|
|
require.NoError(t, db.Create(&types.TaskPendingOp{ID: item.id, TenantID: item.tenant, TaskType: item.kind, ScopeID: "kb", Payload: body}).Error)
|
|
}
|
|
require.NoError(t, repo.Reserve(context.Background(), row, "model", "candidate"))
|
|
require.ErrorIs(t, repo.Reserve(context.Background(), row, "model", "summary"), ErrWikiRecoveryBudgetStopped)
|
|
require.NoError(t, repo.StopAndArchive(context.Background(), row, "other", "stop"))
|
|
require.NoError(t, repo.StopAndArchive(context.Background(), row, "other", "stop"))
|
|
var ids []int64
|
|
require.NoError(t, db.Model(&types.TaskPendingOp{}).Order("id").Pluck("id", &ids).Error)
|
|
require.Equal(t, []int64{44, 45, 46}, ids)
|
|
saved, err := repo.Load(context.Background(), 7, identity)
|
|
require.NoError(t, err)
|
|
var paused []types.TaskPendingOp
|
|
require.NoError(t, json.Unmarshal([]byte(saved.PausedPayload), &paused))
|
|
require.Len(t, paused, 2)
|
|
require.Equal(t, "request_limit", saved.StopReason)
|
|
for _, table := range []string{"knowledges", "chunks"} {
|
|
var count int64
|
|
require.NoError(t, db.Table(table).Count(&count).Error)
|
|
require.Equal(t, int64(1), count)
|
|
}
|
|
}
|
|
|
|
type repositoryHTTPGuard struct {
|
|
repo WikiRecoveryBudgetRepository
|
|
row types.WikiRecoveryBudget
|
|
}
|
|
|
|
func (g repositoryHTTPGuard) Key() string { return g.row.Identity().Key() }
|
|
func (g repositoryHTTPGuard) Check(ctx context.Context) error { return g.repo.Check(ctx, g.row) }
|
|
func (g repositoryHTTPGuard) Reserve(ctx context.Context, model, phase string) error {
|
|
return g.repo.Reserve(ctx, g.row, model, phase)
|
|
}
|
|
|
|
func TestWikiRecoveryBudgetDatabaseFailureSendsNothingAndFailedSendConsumes(t *testing.T) {
|
|
db := openRecoveryBudgetSQLite(t, filepath.Join(t.TempDir(), "budget.sqlite"))
|
|
row := seedRecoveryBudget(t, db, 2)
|
|
guard := repositoryHTTPGuard{WikiRecoveryBudgetRepository{DB: db}, row}
|
|
var sends atomic.Int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
sends.Add(1)
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
}))
|
|
defer server.Close()
|
|
client := &http.Client{Transport: requestbudget.Transport{Base: http.DefaultTransport}}
|
|
send := func() error {
|
|
ctx := requestbudget.WithModel(requestbudget.WithGuard(context.Background(), guard), "model")
|
|
req, _ := http.NewRequestWithContext(ctx, "POST", server.URL, nil)
|
|
resp, err := client.Do(req)
|
|
if resp != nil {
|
|
resp.Body.Close()
|
|
}
|
|
return err
|
|
}
|
|
require.NoError(t, send())
|
|
require.NoError(t, send())
|
|
require.Error(t, send())
|
|
require.Equal(t, int32(2), sends.Load())
|
|
saved, err := guard.repo.Load(context.Background(), 7, row.Identity())
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(2), saved.ReservedProviderRequests)
|
|
sqlDB, _ := db.DB()
|
|
require.NoError(t, sqlDB.Close())
|
|
require.Error(t, send())
|
|
require.Equal(t, int32(2), sends.Load())
|
|
}
|
|
|
|
func TestWikiRecoveryBudgetCompleteRejectsForeignFinalizeAndWaitsForOwnRows(t *testing.T) {
|
|
db := openRecoveryBudgetSQLite(t, filepath.Join(t.TempDir(), "budget.sqlite"))
|
|
row := seedRecoveryBudget(t, db, 3)
|
|
repo := WikiRecoveryBudgetRepository{DB: db}
|
|
controlled, _ := json.Marshal(row.Identity())
|
|
for _, op := range []types.TaskPendingOp{
|
|
{ID: 50, TenantID: 7, TaskType: types.TypeWikiFinalize, ScopeID: "kb", Payload: controlled},
|
|
{ID: 51, TenantID: 7, TaskType: types.TypeWikiFinalize, ScopeID: "kb", Payload: []byte(`{}`)},
|
|
{ID: 52, TenantID: 7, TaskType: types.TypeWikiFinalize, ScopeID: "kb", Payload: controlled},
|
|
} {
|
|
require.NoError(t, db.Create(&op).Error)
|
|
}
|
|
completed, err := repo.Complete(context.Background(), row, []int64{51})
|
|
require.ErrorIs(t, err, ErrWikiRecoveryBudgetStopped)
|
|
require.False(t, completed)
|
|
var count int64
|
|
require.NoError(t, db.Model(&types.TaskPendingOp{}).Count(&count).Error)
|
|
require.Equal(t, int64(3), count)
|
|
completed, err = repo.Complete(context.Background(), row, []int64{50})
|
|
require.NoError(t, err)
|
|
require.False(t, completed)
|
|
completed, err = repo.Complete(context.Background(), row, []int64{52})
|
|
require.NoError(t, err)
|
|
require.True(t, completed)
|
|
require.NoError(t, db.Model(&types.TaskPendingOp{}).Count(&count).Error)
|
|
require.Equal(t, int64(1), count)
|
|
}
|
|
|
|
func TestWikiRecoveryBudgetFingerprintMatchesScriptPostgres(t *testing.T) {
|
|
node, err := exec.LookPath("node")
|
|
if err != nil {
|
|
t.Skip("Node is required to verify the recovery script fingerprint contract")
|
|
}
|
|
db := setupKBDeletePostgres(t)
|
|
row := seedRecoveryBudget(t, db, 3)
|
|
goSource, goModel, err := WikiRecoveryFingerprints(db, row)
|
|
require.NoError(t, err)
|
|
output, err := exec.Command(node, "-e", `process.stdout.write(JSON.stringify(require('./../../../scripts/wiki_recovery_budget_sql.cjs').recoveryFingerprintsSQL()))`).Output()
|
|
require.NoError(t, err)
|
|
var sql struct{ Source, Model string }
|
|
require.NoError(t, json.Unmarshal(output, &sql))
|
|
var got struct{ Source, Model string }
|
|
require.NoError(t, db.Raw("SELECT "+sql.Source+" AS source, "+sql.Model+" AS model FROM (SELECT 7::bigint AS tenant_id, 'kb' AS scope_id, 'doc' AS related_id) d JOIN knowledge_bases k ON k.id=d.scope_id JOIN knowledges n ON n.id=d.related_id").Scan(&got).Error)
|
|
require.Equal(t, goSource, got.Source)
|
|
require.Equal(t, goModel, got.Model)
|
|
}
|