cloud-ip-validator/internal/db/queries_sites.go

package db

import (
	"context"
	"database/sql"
	"fmt"
	"time"
)

// ListSites returns all configured prober sites, ordered by slot index.
func (d *DB) ListSites(ctx context.Context) ([]Site, error) {
	rows, err := d.QueryContext(ctx, `
		SELECT idx, site_id, hostname, state, last_heartbeat_at, created_at, updated_at FROM sites ORDER BY idx
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []Site
	for rows.Next() {
		s, err := scanSite(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, *s)
	}
	return out, rows.Err()
}

func scanSite(row rowScanner) (*Site, error) {
	var s Site
	var lastHeartbeat sql.NullString
	var createdAt, updatedAt string
	if err := row.Scan(&s.Index, &s.SiteID, &s.Hostname, &s.State, &lastHeartbeat, &createdAt, &updatedAt); err != nil {
		return nil, err
	}
	hb, err := nullStringToTimePtr(lastHeartbeat)
	if err != nil {
		return nil, err
	}
	s.LastHeartbeatAt = hb
	if s.CreatedAt, err = dbToTime(createdAt); err != nil {
		return nil, err
	}
	if s.UpdatedAt, err = dbToTime(updatedAt); err != nil {
		return nil, err
	}
	return &s, nil
}

// GetSiteIndex resolves a configured site_id to its slot index. It
// returns (0, nil) — not an error — if no site with that ID is configured,
// matching the "0 means unconfigured" convention used throughout the
// prober-facing handlers.
func (d *DB) GetSiteIndex(ctx context.Context, siteID string) (int, error) {
	var idx int
	err := d.QueryRowContext(ctx, `SELECT idx FROM sites WHERE site_id=?`, siteID).Scan(&idx)
	if err == sql.ErrNoRows {
		return 0, nil
	}
	if err != nil {
		return 0, err
	}
	return idx, nil
}

// UpsertSite assigns (or renames) the site occupying the given slot. idx
// has no upper bound — an admin can configure as many prober sites as they
// need — only idx >= 1 is enforced (slot numbering starts at 1, matching
// the existing admin-facing convention). site_id must be unique across
// slots.
//
// Re-pointing an existing slot to a (possibly different) site_id always
// resets hostname/state/last_heartbeat_at to their "never connected"
// defaults: a new site_id is a new prober identity, and even a same-value
// PUT is treated as idempotent-intent rather than a heartbeat, so the old
// prober's connection status must not silently linger. The prober
// re-establishes state on its next register/heartbeat within
// poll_interval_seconds.
func (d *DB) UpsertSite(ctx context.Context, idx int, siteID string) error {
	if idx < 1 {
		return fmt.Errorf("site index must be >= 1, got %d: %w", idx, ErrValidation)
	}
	if siteID == "" {
		return fmt.Errorf("site_id must not be empty: %w", ErrValidation)
	}

	var existingIdx int
	err := d.QueryRowContext(ctx, `SELECT idx FROM sites WHERE site_id=? AND idx!=?`, siteID, idx).Scan(&existingIdx)
	if err != nil && err != sql.ErrNoRows {
		return err
	}
	if err == nil {
		return fmt.Errorf("site_id %q already assigned to slot %d: %w", siteID, existingIdx, ErrConflict)
	}

	now := timeToDB(Now())
	_, err = d.ExecContext(ctx, `
		INSERT INTO sites (idx, site_id, created_at, updated_at)
		VALUES (?, ?, ?, ?)
		ON CONFLICT(idx) DO UPDATE SET
			site_id=excluded.site_id, updated_at=excluded.updated_at,
			hostname='', state=?, last_heartbeat_at=NULL
	`, idx, siteID, now, now, SiteUnregistered)
	if err != nil {
		return fmt.Errorf("upsert site: %w", err)
	}
	return nil
}

// DeleteSite frees the given slot.
func (d *DB) DeleteSite(ctx context.Context, idx int) error {
	res, err := d.ExecContext(ctx, `DELETE FROM sites WHERE idx=?`, idx)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return fmt.Errorf("site slot %d: %w", idx, ErrNotFound)
	}
	return nil
}

// RegisterSiteProber records that a prober has (re)connected for this
// site: sets hostname, stamps last_heartbeat_at, and reactivates state to
// idle if it was unregistered or unreachable — mirrors RegisterValidator's
// reactivation UPDATE. Called once at prober startup. Returns
// sql.ErrNoRows if site_id is unconfigured.
func (d *DB) RegisterSiteProber(ctx context.Context, siteID, hostname string) error {
	now := timeToDB(Now())
	res, err := d.ExecContext(ctx, `
		UPDATE sites SET hostname=?, last_heartbeat_at=?, updated_at=?,
			state = CASE WHEN state IN (?, ?) THEN ? ELSE state END
		WHERE site_id=?
	`, hostname, now, now, SiteUnregistered, SiteUnreachable, SiteIdle, siteID)
	if err != nil {
		return fmt.Errorf("register site prober: %w", err)
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return sql.ErrNoRows
	}
	return nil
}

// SiteHeartbeat is the periodic "I'm alive" touch — mirrors Heartbeat for
// validators exactly: stamps last_heartbeat_at and flips state from
// unreachable back to idle (does not touch hostname).
func (d *DB) SiteHeartbeat(ctx context.Context, siteID string) error {
	now := timeToDB(Now())
	res, err := d.ExecContext(ctx, `
		UPDATE sites SET last_heartbeat_at=?, updated_at=?,
			state = CASE WHEN state=? THEN ? ELSE state END
		WHERE site_id=?
	`, now, now, SiteUnreachable, SiteIdle, siteID)
	if err != nil {
		return fmt.Errorf("site heartbeat: %w", err)
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return sql.ErrNoRows
	}
	return nil
}

// ListStaleSiteHeartbeats returns sites whose last heartbeat predates the
// given cutoff and that aren't already marked unreachable.
func (d *DB) ListStaleSiteHeartbeats(ctx context.Context, cutoff time.Time) ([]Site, error) {
	rows, err := d.QueryContext(ctx, `
		SELECT idx, site_id, hostname, state, last_heartbeat_at, created_at, updated_at
		FROM sites
		WHERE state != ? AND last_heartbeat_at IS NOT NULL AND last_heartbeat_at < ?
	`, SiteUnreachable, timeToDB(cutoff))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []Site
	for rows.Next() {
		s, err := scanSite(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, *s)
	}
	return out, rows.Err()
}

func (d *DB) MarkSiteUnreachable(ctx context.Context, siteID string) error {
	_, err := d.ExecContext(ctx, `UPDATE sites SET state=?, updated_at=? WHERE site_id=?`,
		SiteUnreachable, timeToDB(Now()), siteID)
	return err
}