ITWorx-Pulse-Public/internal/agentstore/postgres.go

package agentstore

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"strings"
	"time"

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

// MaxAgentIDBytes bounds the reporting agent identifier. It matches the column check in
// migration 0016 so a rejection surfaces as a terse store error rather than a constraint
// violation from PostgreSQL.
const MaxAgentIDBytes = 128

// MaxClockSkew is the only tolerance granted to an agent clock that runs ahead of the
// database. The agent and the database share one host in the supported topology, so real
// skew is sub-millisecond; one second absorbs scheduling jitter while staying far below
// the tightest freshness window, so a skewed clock cannot mask meaningful staleness.
const MaxClockSkew = time.Second

// ErrUnavailable is returned when the store has no database pool. Callers translate it
// into an Unknown status exactly as they do ErrNoSnapshot.
var ErrUnavailable = errors.New("agent snapshot store is unavailable")

// PostgresStore is the only implementation of Store. It owns the agent_snapshots table
// created by migration 0016 and keeps exactly one row per (agent, capability).
//
// Transaction boundary: Put commits the latest snapshot and any derived capacity samples
// atomically. Latest remains a single read statement.
type PostgresStore struct {
	// Pool is the shared pgx pool. A nil pool makes every call return ErrUnavailable.
	Pool *pgxpool.Pool
	// Clock supplies ReceivedAt and the future-observation check. It exists for tests;
	// production leaves it nil and the store uses the wall clock in UTC.
	Clock func() time.Time
}

var _ Store = PostgresStore{}

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

// Put records the newest snapshot for one capability, replacing any previous one for the
// same agent. It rejects an unknown capability, a missing or oversized payload, a payload
// that is not a JSON object, and a zero or future ObservedAt.
//
// ReceivedAt is always taken from the store clock: the caller's value is ignored so a
// skewed or hostile agent clock cannot make stale data look freshly received.
//
// A snapshot that is older than the row already stored is accepted but does not overwrite
// it, so a delayed retry cannot resurrect superseded telemetry.
func (s PostgresStore) Put(ctx context.Context, snapshot Snapshot) error {
	row, err := prepare(snapshot, s.now())
	if err != nil {
		return err
	}
	if s.Pool == nil {
		return ErrUnavailable
	}
	tx, err := s.Pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("begin agent snapshot write: %w", err)
	}
	defer func() {
		rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = tx.Rollback(rollbackCtx)
	}()
	if _, err := tx.Exec(ctx, `INSERT INTO agent_snapshots (agent_id, capability, observed_at, received_at, payload)
        VALUES ($1, $2, $3, $4, $5::jsonb)
        ON CONFLICT (agent_id, capability) DO UPDATE
        SET observed_at = EXCLUDED.observed_at, received_at = EXCLUDED.received_at, payload = EXCLUDED.payload
        WHERE agent_snapshots.observed_at <= EXCLUDED.observed_at`,
		row.AgentID, string(row.Capability), row.ObservedAt, row.ReceivedAt, []byte(row.Payload)); err != nil {
		return fmt.Errorf("write agent snapshot: %w", err)
	}
	if err := persistCapacitySamples(ctx, tx, row); err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit agent snapshot write: %w", err)
	}
	return nil
}

type capacitySample struct {
	Kind          string    `json:"kind"`
	ID            string    `json:"id"`
	Name          string    `json:"name"`
	ObservedAt    time.Time `json:"observed_at"`
	UsedBytes     uint64    `json:"used_bytes"`
	CapacityBytes uint64    `json:"capacity_bytes"`
}

func persistCapacitySamples(ctx context.Context, tx pgx.Tx, snapshot Snapshot) error {
	samples, err := capacitySamples(snapshot)
	if err != nil {
		return err
	}
	if len(samples) == 0 {
		return nil
	}
	payload, err := json.Marshal(samples)
	if err != nil {
		return fmt.Errorf("encode capacity samples: %w", err)
	}
	_, err = tx.Exec(ctx, `INSERT INTO capacity_samples (entity_kind,entity_id,entity_name,source_id,sampled_at,observed_at,used_bytes,capacity_bytes)
SELECT sample.kind,sample.id,sample.name,$1,
       date_trunc('day',sample.observed_at) + floor(extract(hour FROM sample.observed_at)/6)*interval '6 hours',
       sample.observed_at,sample.used_bytes,sample.capacity_bytes
FROM jsonb_to_recordset($2::jsonb) AS sample(kind text,id text,name text,observed_at timestamptz,used_bytes bigint,capacity_bytes bigint)
ON CONFLICT (entity_kind,entity_id,source_id,sampled_at) DO UPDATE
SET entity_name=EXCLUDED.entity_name,observed_at=EXCLUDED.observed_at,used_bytes=EXCLUDED.used_bytes,capacity_bytes=EXCLUDED.capacity_bytes
WHERE capacity_samples.observed_at <= EXCLUDED.observed_at`, snapshot.AgentID, payload)
	if err != nil {
		return fmt.Errorf("persist capacity samples: %w", err)
	}
	return nil
}

func capacitySamples(snapshot Snapshot) ([]capacitySample, error) {
	kind := ""
	collection := ""
	capacityField := ""
	switch snapshot.Capability {
	case CapabilityShares:
		kind, collection = "share", "shares"
	case CapabilityPools:
		kind, collection, capacityField = "pool", "pools", "usableBytes"
	case CapabilityDisks:
		kind, collection, capacityField = "disk", "disks", "sizeBytes"
	default:
		return nil, nil
	}
	var document map[string]json.RawMessage
	if err := json.Unmarshal(snapshot.Payload, &document); err != nil {
		return nil, errors.New("decode capacity snapshot payload")
	}
	var items []map[string]json.RawMessage
	if err := json.Unmarshal(document[collection], &items); err != nil {
		return nil, fmt.Errorf("decode %s capacity collection", kind)
	}
	samples := make([]capacitySample, 0, len(items))
	for _, item := range items {
		var id, name string
		var used, capacity uint64
		if json.Unmarshal(item["id"], &id) != nil || json.Unmarshal(item["name"], &name) != nil || json.Unmarshal(item["usedBytes"], &used) != nil {
			continue
		}
		if capacityField != "" {
			if json.Unmarshal(item[capacityField], &capacity) != nil {
				continue
			}
		}
		observed := snapshot.ObservedAt
		if kind == "share" {
			var sizeObserved time.Time
			if json.Unmarshal(item["sizeObservedAt"], &sizeObserved) == nil && !sizeObserved.IsZero() {
				observed = sizeObserved.UTC()
			}
		}
		id, name = strings.TrimSpace(id), strings.TrimSpace(name)
		if id == "" || name == "" || len(id) > 128 || len(name) > 255 || used > math.MaxInt64 || capacity > math.MaxInt64 || observed.IsZero() || observed.After(snapshot.ReceivedAt.Add(MaxClockSkew)) {
			continue
		}
		samples = append(samples, capacitySample{Kind: kind, ID: id, Name: name, ObservedAt: observed, UsedBytes: used, CapacityBytes: capacity})
	}
	return samples, nil
}

// Latest returns the most recent snapshot for the capability across every reporting
// agent, or ErrNoSnapshot when none has been recorded.
func (s PostgresStore) Latest(ctx context.Context, capability Capability) (Snapshot, error) {
	if !capability.Valid() {
		return Snapshot{}, fmt.Errorf("unknown agent capability %q", capability)
	}
	if s.Pool == nil {
		return Snapshot{}, ErrUnavailable
	}
	return scanSnapshot(s.Pool.QueryRow(ctx, `SELECT agent_id, capability, observed_at, received_at, payload
        FROM agent_snapshots WHERE capability = $1 ORDER BY observed_at DESC, received_at DESC LIMIT 1`,
		string(capability)))
}

// scanSnapshot reads one row and maps the absence of a row onto ErrNoSnapshot, which is an
// expected condition on a fresh install rather than a failure.
func scanSnapshot(row pgx.Row) (Snapshot, error) {
	var (
		snapshot Snapshot
		name     string
		payload  []byte
	)
	err := row.Scan(&snapshot.AgentID, &name, &snapshot.ObservedAt, &snapshot.ReceivedAt, &payload)
	if errors.Is(err, pgx.ErrNoRows) {
		return Snapshot{}, ErrNoSnapshot
	}
	if err != nil {
		return Snapshot{}, fmt.Errorf("read agent snapshot: %w", err)
	}
	snapshot.Capability = Capability(name)
	snapshot.ObservedAt = snapshot.ObservedAt.UTC()
	snapshot.ReceivedAt = snapshot.ReceivedAt.UTC()
	snapshot.Payload = json.RawMessage(payload)
	return snapshot, nil
}

// prepare enforces every bound the Writer contract documents and returns the row the
// store persists. It is the single place where an inbound snapshot is trusted, and it
// overwrites ReceivedAt with the store clock so the caller cannot influence freshness.
func prepare(snapshot Snapshot, now time.Time) (Snapshot, error) {
	if snapshot.AgentID == "" || len(snapshot.AgentID) > MaxAgentIDBytes {
		return Snapshot{}, errors.New("agent id is required and bounded")
	}
	if !snapshot.Capability.Valid() {
		return Snapshot{}, fmt.Errorf("unknown agent capability %q", snapshot.Capability)
	}
	if len(snapshot.Payload) == 0 {
		return Snapshot{}, errors.New("agent snapshot payload is required")
	}
	if len(snapshot.Payload) > MaxPayloadBytes {
		return Snapshot{}, fmt.Errorf("agent snapshot payload exceeds %d bytes", MaxPayloadBytes)
	}
	if snapshot.ObservedAt.IsZero() {
		return Snapshot{}, errors.New("agent snapshot observed time is required")
	}
	if snapshot.ObservedAt.After(now.Add(MaxClockSkew)) {
		return Snapshot{}, errors.New("agent snapshot observed time is in the future")
	}
	var object map[string]json.RawMessage
	if err := json.Unmarshal(snapshot.Payload, &object); err != nil || object == nil {
		return Snapshot{}, errors.New("agent snapshot payload must be a JSON object")
	}
	snapshot.ObservedAt = snapshot.ObservedAt.UTC()
	snapshot.ReceivedAt = now
	return snapshot, nil
}