xinmi/internal/application/repository/knowledgebase_delete_postgres_test.go

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)
	})
}