package repository import ( "context" "fmt" "os" "os/exec" "path/filepath" "testing" "time" "github.com/Tencent/WeKnora/internal/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/driver/postgres" "gorm.io/gorm" ) func setupKBDeletePostgres(t *testing.T) *gorm.DB { t.Helper() bin := os.Getenv("POSTGRES_TEST_BIN") if bin == "" { bin = "/opt/homebrew/opt/postgresql@16/bin" } if _, err := os.Stat(filepath.Join(bin, "initdb")); err != nil { t.Skip("isolated PostgreSQL binaries unavailable; set POSTGRES_TEST_BIN to run lock regression") } dir, err := os.MkdirTemp("/tmp", "kb-delete-pg-") require.NoError(t, err) data := filepath.Join(dir, "data") running := false run := func(name string, args ...string) { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() output, err := exec.CommandContext(ctx, filepath.Join(bin, name), args...).CombinedOutput() require.NoError(t, err, "%s: %s", name, output) } t.Cleanup(func() { if running { run("pg_ctl", "-D", data, "-m", "fast", "-w", "stop") } assert.NoError(t, os.RemoveAll(dir)) }) run("initdb", "-D", data, "-U", "fixture", "--auth=trust", "--no-locale", "--encoding=UTF8") run("pg_ctl", "-D", data, "-l", filepath.Join(dir, "postgres.log"), "-o", fmt.Sprintf("-F -c listen_addresses='' -k %s -p 55460", dir), "-w", "start") running = true db, err := gorm.Open(postgres.Open(fmt.Sprintf("host=%s port=55460 user=fixture dbname=postgres sslmode=disable application_name=kb-delete-lock-test", dir)), &gorm.Config{}) require.NoError(t, err) sqlDB, err := db.DB() require.NoError(t, err) t.Cleanup(func() { assert.NoError(t, sqlDB.Close()) }) require.NoError(t, db.Exec(` CREATE TABLE knowledge_bases (id text PRIMARY KEY, tenant_id bigint NOT NULL, deleted_at timestamptz); CREATE TABLE knowledges (id text PRIMARY KEY, tenant_id bigint NOT NULL, knowledge_base_id text NOT NULL, parse_status text, deleted_at timestamptz); CREATE TABLE knowledge_tags (id text PRIMARY KEY, seq_id bigserial UNIQUE, tenant_id bigint NOT NULL, knowledge_base_id text NOT NULL, name text NOT NULL, color text, sort_order int, created_at timestamptz, updated_at timestamptz); CREATE TABLE knowledge_tag_relations (knowledge_id text NOT NULL, tag_id text NOT NULL, created_at timestamptz, PRIMARY KEY (knowledge_id, tag_id)); `).Error) return db } func awaitKBDeleteBlockedWriter(t *testing.T, db *gorm.DB) { t.Helper() require.Eventually(t, func() bool { var count int64 err := db.Raw("SELECT count(*) FROM pg_stat_activity WHERE application_name = 'kb-delete-lock-test' AND wait_event_type = 'Lock'").Scan(&count).Error return err == nil && count > 0 }, 5*time.Second, 10*time.Millisecond, "concurrent operation must wait on the KB row lock") } func TestKnowledgeBaseDeletePostgresSerializesCategoryWriters(t *testing.T) { db := setupKBDeletePostgres(t) ctx, cancel := context.WithTimeout(kbDeleteContext(1), 30*time.Second) defer cancel() tagRepo := &knowledgeTagRepository{db: db} knowledgeRepo := &knowledgeRepository{db: db} for _, method := range []string{"create", "update", "set relations", "add relations"} { t.Run(method, func(t *testing.T) { require.NoError(t, db.Exec("TRUNCATE knowledge_tag_relations, knowledge_tags, knowledges, knowledge_bases").Error) require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id, tenant_id) VALUES ('kb', 1)").Error) require.NoError(t, db.Exec("INSERT INTO knowledges (id, tenant_id, knowledge_base_id, parse_status) VALUES ('doc', 1, 'kb', 'completed')").Error) tag := &types.KnowledgeTag{ID: "tag", TenantID: 1, KnowledgeBaseID: "kb", Name: "category"} require.NoError(t, tagRepo.Create(ctx, tag)) require.NoError(t, knowledgeRepo.SetKnowledgeTags(ctx, "doc", []string{"tag"})) tx := db.WithContext(ctx).Begin() require.NoError(t, tx.Error) defer tx.Rollback() require.NoError(t, (&knowledgeBaseRepository{db: tx}).DeleteKnowledgeBase(ctx, "kb")) done := make(chan error, 1) go func() { switch method { case "create": done <- tagRepo.Create(ctx, &types.KnowledgeTag{ID: "late", TenantID: 1, KnowledgeBaseID: "kb", Name: "late"}) case "update": done <- tagRepo.Update(ctx, tag) case "set relations": done <- knowledgeRepo.SetKnowledgeTags(ctx, "doc", []string{"tag"}) case "add relations": done <- knowledgeRepo.AddKnowledgeTagRelations(ctx, 1, "kb", "doc", []string{"tag"}) } }() awaitKBDeleteBlockedWriter(t, db) require.NoError(t, tx.Commit().Error) require.ErrorIs(t, <-done, ErrKnowledgeBaseNotFound) assertKBDeleteTableCount(t, db, "knowledge_tags", 0) assertKBDeleteTableCount(t, db, "knowledge_tag_relations", 0) }) } t.Run("delete waits for category writer then removes its category", func(t *testing.T) { require.NoError(t, db.Exec("TRUNCATE knowledge_tag_relations, knowledge_tags, knowledges, knowledge_bases").Error) require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id, tenant_id) VALUES ('kb', 1)").Error) tx := db.WithContext(ctx).Begin() require.NoError(t, tx.Error) defer tx.Rollback() require.NoError(t, (&knowledgeTagRepository{db: tx}).Create(ctx, &types.KnowledgeTag{ID: "early", TenantID: 1, KnowledgeBaseID: "kb", Name: "early"})) done := make(chan error, 1) go func() { done <- (&knowledgeBaseRepository{db: db}).DeleteKnowledgeBase(ctx, "kb") }() awaitKBDeleteBlockedWriter(t, db) require.NoError(t, tx.Commit().Error) require.NoError(t, <-done) assertKBDeleteTableCount(t, db, "knowledge_tags", 0) }) }