knowledge-governance/internal/workflows/testweknora/poll.go

package testweknora

import (
	"context"
	"errors"
	"time"

	"knowledge-governance/internal/governance"
	"knowledge-governance/internal/jobs"
	postgresstore "knowledge-governance/internal/platform/postgres"
	"knowledge-governance/internal/platform/weknora"
)

var (
	ErrInvalidPollJobPayload = errors.New("test_poll_invalid_payload")
	ErrPollRemoteFailed      = errors.New("test_poll_remote_failed")
	ErrPollRemoteIdentity    = errors.New("test_poll_remote_identity_invalid")
	ErrPollStateConflict     = errors.New("test_poll_state_conflict")
)

type PollStore interface {
	LoadPublicationForPoll(context.Context, string) (postgresstore.PublicationForPoll, error)
	RecordRemoteStatusAndScheduleNextPoll(context.Context, string, int, governance.ProcessingState, *time.Time, time.Time, time.Time) (int, error)
	RecordCompletedAndScheduleValidation(context.Context, string, int, *time.Time, time.Time) error
	RecordProcessingFailure(context.Context, string, int, string, time.Time) error
	MarkStalled(context.Context, string, int, time.Time, time.Time) error
}

type PollClient interface {
	GetKnowledge(context.Context, string) (weknora.Knowledge, error)
}

type PollHandlerConfig struct {
	Store      PollStore
	Client     PollClient
	PollDelay  time.Duration
	StallAfter time.Duration
	Now        func() time.Time
}

type PollHandler struct {
	store      PollStore
	client     PollClient
	pollDelay  time.Duration
	stallAfter time.Duration
	now        func() time.Time
}

type pollPayload struct {
	PublicationID string `json:"publication_id"`
	Sequence      int    `json:"sequence"`
}

func NewPollHandler(config PollHandlerConfig) (*PollHandler, error) {
	if config.Store == nil || config.Client == nil || config.PollDelay <= 0 || config.StallAfter <= 0 {
		return nil, errors.New("invalid test poll handler configuration")
	}
	if config.Now == nil {
		config.Now = time.Now
	}
	return &PollHandler{
		store: config.Store, client: config.Client, pollDelay: config.PollDelay,
		stallAfter: config.StallAfter, now: config.Now,
	}, nil
}

func (handler *PollHandler) Handle(ctx context.Context, job jobs.Job) error {
	var payload pollPayload
	if err := decodeStrictJSON(job.Payload, &payload); err != nil || payload.PublicationID == "" || payload.Sequence < 1 {
		return jobs.Permanent(jobs.ErrCodeHandlerFallback, ErrInvalidPollJobPayload)
	}
	publication, err := handler.store.LoadPublicationForPoll(ctx, payload.PublicationID)
	if err != nil {
		return ErrPollStateConflict
	}
	if publication.KnowledgeID == "" || publication.PollSequence != payload.Sequence {
		return ErrPollStateConflict
	}
	now := handler.now()
	knowledge, err := handler.client.GetKnowledge(ctx, publication.KnowledgeID)
	if err != nil {
		return ErrPollRemoteFailed
	}
	if knowledge.ID != publication.KnowledgeID || knowledge.KnowledgeBaseID != publication.TestKBID {
		return ErrPollRemoteIdentity
	}
	state := mapParseStatus(knowledge.ParseStatus)
	if state == governance.ProcessingCompleted {
		if err := handler.store.RecordCompletedAndScheduleValidation(ctx, publication.ID, publication.PollSequence, &knowledge.UpdatedAt, now); err != nil {
			return ErrPollStateConflict
		}
		return nil
	}
	if state == governance.ProcessingFailed {
		if err := handler.store.RecordProcessingFailure(ctx, publication.ID, publication.PollSequence, "remote_parse_failed", now); err != nil {
			return ErrPollStateConflict
		}
		return nil
	}
	if !state.Valid() {
		if err := handler.store.RecordProcessingFailure(ctx, publication.ID, publication.PollSequence, "remote_unknown_status", now); err != nil {
			return ErrPollStateConflict
		}
		return nil
	}
	cutoff := now.Add(-handler.stallAfter)
	lastProgress := publication.LastProgressAt
	if lastProgress == nil {
		lastProgress = publication.ProcessingStartedAt
	}
	if lastProgress == nil {
		lastProgress = publication.LastObservedAt
	}
	if lastProgress != nil && !lastProgress.After(cutoff) {
		if err := handler.store.MarkStalled(ctx, publication.ID, publication.PollSequence, cutoff, now); err != nil {
			return ErrPollStateConflict
		}
		return nil
	}
	if _, err := handler.store.RecordRemoteStatusAndScheduleNextPoll(ctx, publication.ID, publication.PollSequence, state, &knowledge.UpdatedAt, now, now.Add(handler.pollDelay)); err != nil {
		return ErrPollStateConflict
	}
	return nil
}

func mapParseStatus(status string) governance.ProcessingState {
	switch status {
	case "pending":
		return governance.ProcessingPending
	case "processing":
		return governance.ProcessingProcessing
	case "finalizing":
		return governance.ProcessingFinalizing
	case "completed":
		return governance.ProcessingCompleted
	case "failed":
		return governance.ProcessingFailed
	default:
		return governance.ProcessingState(status)
	}
}