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 }