30 lines
1.3 KiB
Go
30 lines
1.3 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/requestbudget"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Reserving the whole execution in the existing dispatch ledger fences both
|
||
|
|
// concurrent delivery and worker restart. A crash never restores its three sends.
|
||
|
|
func (s *ChunkExtractService) claimGraphExecution(ctx context.Context, p types.ExtractChunkPayload, taskID string) error {
|
||
|
|
e := p.ExpectedSource
|
||
|
|
if s.executionDB == nil || e == nil || taskID == "" || p.ModelID == "" || p.TenantID == 0 ||
|
||
|
|
p.ChunkID == "" || e.KnowledgeBaseID == "" || e.KnowledgeID == "" || e.ContentSHA256 == "" {
|
||
|
|
return fmt.Errorf("%w: missing graph execution receipt", requestbudget.ErrStopped)
|
||
|
|
}
|
||
|
|
result := s.executionDB.WithContext(ctx).Exec(`UPDATE native_graph_backfill_failures
|
||
|
|
SET state='execution_reserved',updated_at=?
|
||
|
|
WHERE task_id=? AND tenant_id=? AND knowledge_base_id=? AND chunk_id=? AND knowledge_id=? AND content_hash=?
|
||
|
|
AND state='dispatch_requested' AND manual_retries=0 AND reason=''`,
|
||
|
|
time.Now().UTC(), taskID, p.TenantID, e.KnowledgeBaseID, p.ChunkID, e.KnowledgeID, e.ContentSHA256)
|
||
|
|
if result.Error != nil || result.RowsAffected != 1 {
|
||
|
|
return fmt.Errorf("%w: graph execution receipt unavailable or already spent", requestbudget.ErrStopped)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|