ITWorx-Pulse-Public/internal/discovery/postgres_store.go

package discovery

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// DefaultJobType is the job_runs.job_type discovery claims are recorded under.
const DefaultJobType = "discovery"

const (
	defaultWindow   = time.Minute
	defaultLeaseTTL = 5 * time.Minute
	maxEventType    = 160
	maxSummary      = 500
	maxDedupKey     = 255
)

var (
	// ErrLeaseLost reports that the job run this process claimed was taken over
	// or completed elsewhere, so its result must not overwrite the newer one.
	ErrLeaseLost = errors.New("discovery job lease was lost")
	// ErrInvalidEvent reports an event that cannot be persisted safely.
	ErrInvalidEvent = errors.New("invalid discovery event")
	// ErrUnavailable reports a store without a usable database pool.
	ErrUnavailable = errors.New("discovery store is unavailable")

	uuidPattern     = regexp.MustCompile("^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-8][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$")
	allowedSeverity = map[string]struct{}{"info": {}, "attention": {}, "warning": {}, "critical": {}}
)

// PostgresStore persists discovery job runs and discovered events.
//
// It deliberately mirrors alertworker.PostgresLeaseStore: a claim is a row in
// job_runs guarded by (job_type, job_key, scheduled_at) with a bounded
// lease_owner/lease_until pair, so two workers never run the same discovery
// window twice and a worker that crashes mid-run has its lease reclaimed once
// it expires instead of blocking discovery forever. Events are inserted with
// ON CONFLICT DO NOTHING against the natural (source_id, dedup_key,
// occurred_at) key, which is what makes a repeated pass idempotent.
type PostgresStore struct {
	Pool *pgxpool.Pool
	// JobType overrides the job_runs.job_type value; DefaultJobType is used when empty.
	JobType string
	// Owner identifies this worker in job_runs.lease_owner.
	Owner string
	// Window is the schedule granularity a claim is truncated to. Two claims of
	// the same key inside one window are the same unit of work.
	Window time.Duration
	// LeaseTTL bounds how long a claim blocks another worker after a crash.
	LeaseTTL time.Duration
	// Now is injectable for tests; time.Now is used when nil.
	Now func() time.Time

	mu     sync.Mutex
	claims map[string]time.Time
}

// NewPostgresStore validates the configuration and returns a ready store.
func NewPostgresStore(pool *pgxpool.Pool, owner string) (*PostgresStore, error) {
	if pool == nil {
		return nil, ErrUnavailable
	}
	if strings.TrimSpace(owner) == "" || len(owner) > 120 {
		return nil, errors.New("discovery store requires a bounded owner")
	}
	return &PostgresStore{Pool: pool, JobType: DefaultJobType, Owner: owner, Window: defaultWindow, LeaseTTL: defaultLeaseTTL}, nil
}

func (s *PostgresStore) jobType() string {
	if strings.TrimSpace(s.JobType) == "" {
		return DefaultJobType
	}
	return s.JobType
}

func (s *PostgresStore) window() time.Duration {
	if s.Window <= 0 {
		return defaultWindow
	}
	return s.Window
}

func (s *PostgresStore) leaseTTL() time.Duration {
	if s.LeaseTTL <= 0 {
		return defaultLeaseTTL
	}
	return s.LeaseTTL
}

func (s *PostgresStore) now() time.Time {
	if s.Now != nil {
		return s.Now().UTC()
	}
	return time.Now().UTC()
}

// Claim reserves the discovery window that contains at for this worker. It
// returns false without an error when another worker holds a live lease or the
// window already completed, which is the normal "nothing to do" outcome.
func (s *PostgresStore) Claim(ctx context.Context, jobKey string, at time.Time) (bool, error) {
	if s == nil || s.Pool == nil {
		return false, ErrUnavailable
	}
	if strings.TrimSpace(jobKey) == "" || len(jobKey) > 255 {
		return false, errors.New("discovery job key is invalid")
	}
	if at.IsZero() {
		at = s.now()
	}
	now := s.now()
	scheduledAt := at.UTC().Truncate(s.window())
	leaseUntil := now.Add(s.leaseTTL())
	tx, err := s.Pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return false, fmt.Errorf("begin discovery claim: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	var insertedID string
	err = tx.QueryRow(ctx, `INSERT INTO job_runs (id,job_type,job_key,scheduled_at,started_at,status,lease_owner,lease_until) VALUES ($1,$2,$3,$4,$5,'running',$6,$7) ON CONFLICT (job_type,job_key,scheduled_at) DO NOTHING RETURNING id`,
		newID(), s.jobType(), jobKey, scheduledAt, now, s.Owner, leaseUntil).Scan(&insertedID)
	if err == nil {
		if err := tx.Commit(ctx); err != nil {
			return false, fmt.Errorf("commit discovery claim: %w", err)
		}
		s.rememberClaim(jobKey, scheduledAt)
		return true, nil
	}
	if !errors.Is(err, pgx.ErrNoRows) {
		return false, fmt.Errorf("insert discovery claim: %w", err)
	}
	var status string
	var existingUntil *time.Time
	if err := tx.QueryRow(ctx, `SELECT status,lease_until FROM job_runs WHERE job_type=$1 AND job_key=$2 AND scheduled_at=$3 FOR UPDATE`, s.jobType(), jobKey, scheduledAt).Scan(&status, &existingUntil); err != nil {
		return false, fmt.Errorf("read discovery claim: %w", err)
	}
	if status != "running" || (existingUntil != nil && existingUntil.After(now)) {
		if err := tx.Commit(ctx); err != nil {
			return false, fmt.Errorf("commit discovery claim contention: %w", err)
		}
		return false, nil
	}
	tag, err := tx.Exec(ctx, `UPDATE job_runs SET started_at=$1,completed_at=NULL,error_code=NULL,lease_owner=$2,lease_until=$3 WHERE job_type=$4 AND job_key=$5 AND scheduled_at=$6 AND status='running' AND (lease_until IS NULL OR lease_until <= $7)`,
		now, s.Owner, leaseUntil, s.jobType(), jobKey, scheduledAt, now)
	if err != nil {
		return false, fmt.Errorf("reclaim discovery lease: %w", err)
	}
	if tag.RowsAffected() != 1 {
		if err := tx.Commit(ctx); err != nil {
			return false, fmt.Errorf("commit discovery reclaim contention: %w", err)
		}
		return false, nil
	}
	if err := tx.Commit(ctx); err != nil {
		return false, fmt.Errorf("commit discovery reclaim: %w", err)
	}
	s.rememberClaim(jobKey, scheduledAt)
	return true, nil
}

// Finish records the outcome of a claimed run. It only updates the row this
// worker still owns, so a run whose lease expired cannot overwrite the result
// of the worker that took over.
func (s *PostgresStore) Finish(ctx context.Context, run JobRun) error {
	if s == nil || s.Pool == nil {
		return ErrUnavailable
	}
	if strings.TrimSpace(run.Key) == "" {
		return errors.New("discovery job key is required")
	}
	scheduledAt, ok := s.takeClaim(run.Key)
	if !ok {
		reference := run.StartedAt
		if reference.IsZero() {
			reference = s.now()
		}
		scheduledAt = reference.UTC().Truncate(s.window())
	}
	status := "completed"
	if run.Status != "succeeded" && run.Status != "completed" {
		status = "failed"
	}
	errorCode := boundedText(run.ErrorCode, 160)
	counts, err := json.Marshal(map[string]any{"status": status, "attempts": run.Attempts})
	if err != nil {
		return fmt.Errorf("encode discovery counts: %w", err)
	}
	completedAt := run.CompletedAt
	if completedAt.IsZero() {
		completedAt = s.now()
	}
	tag, err := s.Pool.Exec(ctx, `UPDATE job_runs SET status=$1,completed_at=$2,counts=$3::jsonb,error_code=NULLIF($4,''),lease_owner=NULL,lease_until=NULL WHERE job_type=$5 AND job_key=$6 AND scheduled_at=$7 AND lease_owner=$8 AND status='running'`,
		status, completedAt.UTC(), counts, errorCode, s.jobType(), run.Key, scheduledAt, s.Owner)
	if err != nil {
		return fmt.Errorf("complete discovery job: %w", err)
	}
	if tag.RowsAffected() != 1 {
		return ErrLeaseLost
	}
	return nil
}

// Emit persists one discovered event and reports whether it was new. A repeated
// pass emitting the same (source, dedup key, occurrence time) is a no-op.
func (s *PostgresStore) Emit(ctx context.Context, event Event) (bool, error) {
	if s == nil || s.Pool == nil {
		return false, ErrUnavailable
	}
	event, err := normalizeEvent(event)
	if err != nil {
		return false, err
	}
	attributes, err := json.Marshal(map[string]any{})
	if err != nil {
		return false, fmt.Errorf("encode discovery attributes: %w", err)
	}
	var inserted string
	err = s.Pool.QueryRow(ctx, `INSERT INTO events (id,event_type,severity,entity_id,source_id,occurred_at,received_at,dedup_key,summary,attributes) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10::jsonb) ON CONFLICT (source_id,dedup_key,occurred_at) DO NOTHING RETURNING id`,
		DeterministicEventID(event), event.Type, event.Severity, nullableUUID(event.EntityID), event.SourceID, event.OccurredAt, s.now(), event.DedupKey, event.Summary, attributes).Scan(&inserted)
	if errors.Is(err, pgx.ErrNoRows) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("emit discovery event: %w", err)
	}
	return true, nil
}

// DeterministicEventID derives a stable UUID from the event identity so a
// retried emit reuses the same primary key instead of racing on a new one.
func DeterministicEventID(event Event) string {
	sum := sha256.Sum256([]byte("itworx-pulse/discovery-event/v1/" + event.SourceID + "\x00" + event.DedupKey + "\x00" + event.OccurredAt.UTC().Format(time.RFC3339Nano)))
	return formatUUID(sum[:16], 0x50)
}

func normalizeEvent(event Event) (Event, error) {
	event.SourceID = strings.TrimSpace(event.SourceID)
	event.DedupKey = strings.TrimSpace(event.DedupKey)
	event.Type = strings.TrimSpace(event.Type)
	event.Summary = strings.TrimSpace(event.Summary)
	event.EntityID = strings.TrimSpace(event.EntityID)
	event.Severity = strings.ToLower(strings.TrimSpace(event.Severity))
	if event.Severity == "" {
		event.Severity = "info"
	}
	if !uuidPattern.MatchString(event.SourceID) {
		return Event{}, fmt.Errorf("%w: source id must be a registered data source UUID", ErrInvalidEvent)
	}
	if event.EntityID != "" && !uuidPattern.MatchString(event.EntityID) {
		return Event{}, fmt.Errorf("%w: entity id must be a UUID", ErrInvalidEvent)
	}
	if _, ok := allowedSeverity[event.Severity]; !ok {
		return Event{}, fmt.Errorf("%w: severity is unsupported", ErrInvalidEvent)
	}
	if event.DedupKey == "" || len(event.DedupKey) > maxDedupKey {
		return Event{}, fmt.Errorf("%w: dedup key must be 1-%d characters", ErrInvalidEvent, maxDedupKey)
	}
	if event.Type == "" || len(event.Type) > maxEventType {
		return Event{}, fmt.Errorf("%w: type must be 1-%d characters", ErrInvalidEvent, maxEventType)
	}
	if event.Summary == "" || len(event.Summary) > maxSummary {
		return Event{}, fmt.Errorf("%w: summary must be 1-%d characters", ErrInvalidEvent, maxSummary)
	}
	if event.OccurredAt.IsZero() {
		return Event{}, fmt.Errorf("%w: occurrence time is required", ErrInvalidEvent)
	}
	event.OccurredAt = event.OccurredAt.UTC()
	return event, nil
}

func (s *PostgresStore) rememberClaim(jobKey string, scheduledAt time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.claims == nil {
		s.claims = make(map[string]time.Time)
	}
	s.claims[jobKey] = scheduledAt
}

func (s *PostgresStore) takeClaim(jobKey string) (time.Time, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	scheduledAt, ok := s.claims[jobKey]
	if ok {
		delete(s.claims, jobKey)
	}
	return scheduledAt, ok
}

func nullableUUID(value string) any {
	if value == "" {
		return nil
	}
	return value
}

func boundedText(value string, max int) string {
	value = strings.TrimSpace(value)
	if len(value) > max {
		return value[:max]
	}
	return value
}

func newID() string {
	raw := make([]byte, 16)
	if _, err := rand.Read(raw); err != nil {
		sum := sha256.Sum256([]byte(fmt.Sprintf("itworx-pulse/discovery-run/v1/%d", time.Now().UnixNano())))
		copy(raw, sum[:16])
	}
	return formatUUID(raw, 0x40)
}

func formatUUID(raw []byte, version byte) string {
	b := make([]byte, 16)
	copy(b, raw)
	b[6] = (b[6] & 0x0f) | version
	b[8] = (b[8] & 0x3f) | 0x80
	return fmt.Sprintf("%s-%s-%s-%s-%s", hex.EncodeToString(b[0:4]), hex.EncodeToString(b[4:6]), hex.EncodeToString(b[6:8]), hex.EncodeToString(b[8:10]), hex.EncodeToString(b[10:16]))
}