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

internal/accesscontrol/resource_version_knowledge_cleanup.go

@@ -0,0 +1,153 @@
+package accesscontrol
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/knowledgeversion"
+	"github.com/Tencent/WeKnora/internal/types"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+)
+
+// A copy can run for two hours. Wait longer before deleting abandoned targets,
+// including targets left by a process whose heartbeat stopped before its work.
+const knowledgeVersionCleanupDelay = 3 * time.Hour
+
+func (s *Service) CleanupAbandonedKnowledgeVersions(ctx context.Context) error {
+	cleaner, ok := s.versionKnowledge.(interface {
+		CleanupKnowledgeVersion(context.Context, string) error
+	})
+	if !ok {
+		return nil
+	}
+	now := time.Now()
+	var job KnowledgeVersionJob
+	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		result := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).
+			Where("status IN ? AND updated_at < ?", []string{"completed", "failed"}, now.Add(-knowledgeVersionCleanupDelay)).
+			Where(`(abandoned_ids IS NOT NULL AND CAST(abandoned_ids AS TEXT) <> '[]' AND CAST(abandoned_ids AS TEXT) <> 'null') OR (status='failed' AND EXISTS (SELECT 1 FROM knowledge_bases kb WHERE kb.id=workbench_knowledge_version_jobs.target_id AND kb.tenant_id=workbench_knowledge_version_jobs.tenant_id AND kb.is_temporary=true AND kb.version_state='copying' AND kb.deleted_at IS NULL))`).
+			Order("updated_at").Limit(1).Find(&job)
+		if result.Error != nil {
+			return result.Error
+		}
+		if result.RowsAffected == 0 {
+			return gorm.ErrRecordNotFound
+		}
+		return tx.Model(&job).UpdateColumn("updated_at", now).Error
+	})
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+	worker, cancel := context.WithTimeout(knowledgeversion.WithWorker(ctx), 5*time.Minute)
+	defer cancel()
+	worker = context.WithValue(worker, types.TenantIDContextKey, job.TenantID)
+	worker = context.WithValue(worker, types.UserIDContextKey, job.OwnerID)
+	worker = context.WithValue(worker, types.TenantRoleContextKey, types.TenantRoleAdmin)
+	worker = context.WithValue(worker, types.SystemAdminContextKey, false)
+	targets := append([]string{}, job.AbandonedIDs...)
+	if job.Status == "failed" {
+		targets = append(targets, job.TargetID)
+	}
+	seen := map[string]bool{}
+	for _, target := range targets {
+		if target == "" || seen[target] {
+			continue
+		}
+		seen[target] = true
+		var kb types.KnowledgeBase
+		result := s.db.WithContext(worker).Where("tenant_id=? AND id=?", job.TenantID, target).Limit(1).Find(&kb)
+		if result.Error != nil {
+			return result.Error
+		}
+		if result.RowsAffected == 0 {
+			if err := s.finishKnowledgeVersionCleanup(worker, job.ID, job.TenantID, target); err != nil {
+				return err
+			}
+			continue
+		}
+		if !kb.IsTemporary || kb.VersionState != "copying" {
+			return fmt.Errorf("refusing to clean non-staging knowledge version %s", target)
+		}
+		if err := cleaner.CleanupKnowledgeVersion(worker, target); err != nil {
+			return err
+		}
+		if err := s.finishKnowledgeVersionCleanup(worker, job.ID, job.TenantID, target); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (s *Service) finishKnowledgeVersionCleanup(ctx context.Context, jobID string, tenant uint64, target string) error {
+	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var job KnowledgeVersionJob
+		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=?", jobID, tenant).First(&job).Error; err != nil {
+			return err
+		}
+		allowed := job.Status == "failed" && job.TargetID == target
+		remaining := make(types.StringArray, 0, len(job.AbandonedIDs))
+		for _, id := range job.AbandonedIDs {
+			if id == target {
+				allowed = true
+			} else {
+				remaining = append(remaining, id)
+			}
+		}
+		if !allowed || job.TargetID == target && job.Status != "failed" {
+			return fmt.Errorf("staging cleanup ownership changed")
+		}
+		var kb types.KnowledgeBase
+		result := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id=? AND id=?", tenant, target).Limit(1).Find(&kb)
+		if result.Error != nil {
+			return result.Error
+		}
+		if result.RowsAffected == 1 {
+			if !kb.IsTemporary || kb.VersionState != "copying" {
+				return fmt.Errorf("staging cleanup target became published")
+			}
+			var docs []types.Knowledge
+			if err := tx.Unscoped().Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Find(&docs).Error; err != nil {
+				return err
+			}
+			var size int64
+			for _, doc := range docs {
+				size += doc.StorageSize
+			}
+			var chunks int64
+			if err := tx.Model(&types.Chunk{}).Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Count(&chunks).Error; err != nil {
+				return err
+			}
+			if chunks != 0 {
+				return fmt.Errorf("staging external cleanup is incomplete")
+			}
+			if size > 0 {
+				r := tx.Model(&types.Tenant{}).Where("id=? AND storage_used>=?", tenant, size).UpdateColumn("storage_used", gorm.Expr("storage_used - ?", size))
+				if r.Error != nil {
+					return r.Error
+				}
+				if r.RowsAffected != 1 {
+					return fmt.Errorf("staging storage accounting mismatch")
+				}
+			}
+			if err := knowledgeversion.DeleteGraph(tx, tenant, target); err != nil {
+				return err
+			}
+			if err := tx.Unscoped().Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Delete(&types.Chunk{}).Error; err != nil {
+				return err
+			}
+			if err := tx.Unscoped().Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Delete(&types.Knowledge{}).Error; err != nil {
+				return err
+			}
+			if err := tx.Unscoped().Where("tenant_id=? AND id=?", tenant, target).Delete(&types.KnowledgeBase{}).Error; err != nil {
+				return err
+			}
+		}
+		return tx.Model(&job).Updates(map[string]any{"abandoned_ids": remaining}).Error
+	})
+}