WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/repository/knowledgebase_delete_postgres_test.go

@@ -0,0 +1,126 @@
+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)
+	})
+}