xinmi/internal/application/repository/knowledge_version.go

package repository

import (
	"context"
	"fmt"

	"github.com/Tencent/WeKnora/internal/common"
	"github.com/Tencent/WeKnora/internal/knowledgeversion"
	"github.com/Tencent/WeKnora/internal/types"
	"gorm.io/gorm"
)

// CreateVersionKnowledge keeps staging documents and their accounted bytes
// atomic so an interrupted copy can later reconcile storage exactly once.
func (r *knowledgeRepository) CreateVersionKnowledge(ctx context.Context, doc *types.Knowledge) error {
	tenant, ok := types.TenantIDFromContext(ctx)
	if !knowledgeversion.IsWorker(ctx) || !ok || doc.TenantID != tenant || doc.StorageSize < 0 {
		return fmt.Errorf("version document requires its owning workspace worker")
	}
	doc.ErrorMessage = common.CleanInvalidUTF8(doc.ErrorMessage)
	return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Create(doc).Error; err != nil {
			return err
		}
		result := tx.Model(&types.Tenant{}).Where("id = ?", tenant).UpdateColumn("storage_used", gorm.Expr("storage_used + ?", doc.StorageSize))
		if result.Error != nil {
			return result.Error
		}
		if result.RowsAffected != 1 {
			return fmt.Errorf("version workspace not found")
		}
		return nil
	})
}