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

package db

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

// RegisterValidator inserts a new validator or updates an existing one's
// hostname/agent_version on re-registration. It never overwrites the state
// of a validator that's mid-assignment, so an agent restarting while it
// owns an IP doesn't silently lose that ownership.
func (d *DB) RegisterValidator(ctx context.Context, validatorID, hostname, osPortID, agentVersion string) error {
	now := timeToDB(Now())
	_, err := d.ExecContext(ctx, `
		INSERT INTO validators (validator_id, hostname, os_port_id, agent_version, state, created_at, updated_at)
		VALUES (?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT(validator_id) DO UPDATE SET
			hostname=excluded.hostname,
			os_port_id=excluded.os_port_id,
			agent_version=excluded.agent_version,
			updated_at=excluded.updated_at
	`, validatorID, hostname, osPortID, agentVersion, ValidatorIdle, now, now)
	if err != nil {
		return fmt.Errorf("register validator: %w", err)
	}
	// A brand-new row already lands in ValidatorIdle via the INSERT branch;
	// a re-registering validator that was 'unregistered' or 'unreachable'
	// (but not mid-assignment) should also come back to idle.
	_, err = d.ExecContext(ctx, `
		UPDATE validators SET state=?, updated_at=?
		WHERE validator_id=? AND state IN (?, ?)
	`, ValidatorIdle, now, validatorID, ValidatorUnregistered, ValidatorUnreachable)
	if err != nil {
		return fmt.Errorf("register validator (reactivate): %w", err)
	}
	return nil
}

func (d *DB) Heartbeat(ctx context.Context, validatorID string) error {
	now := timeToDB(Now())
	res, err := d.ExecContext(ctx, `
		UPDATE validators SET last_heartbeat_at=?, updated_at=?,
			state = CASE WHEN state=? THEN ? ELSE state END
		WHERE validator_id=?
	`, now, now, ValidatorUnreachable, ValidatorIdle, validatorID)
	if err != nil {
		return fmt.Errorf("heartbeat: %w", err)
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return sql.ErrNoRows
	}
	return nil
}

func (d *DB) GetValidator(ctx context.Context, validatorID string) (*Validator, error) {
	row := d.QueryRowContext(ctx, `
		SELECT validator_id, hostname, os_port_id, state, current_ip_id, agent_version,
			last_heartbeat_at, created_at, updated_at
		FROM validators WHERE validator_id=?
	`, validatorID)
	return scanValidator(row)
}

func (d *DB) ListValidators(ctx context.Context) ([]Validator, error) {
	rows, err := d.QueryContext(ctx, `
		SELECT validator_id, hostname, os_port_id, state, current_ip_id, agent_version,
			last_heartbeat_at, created_at, updated_at
		FROM validators ORDER BY validator_id
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Validator
	for rows.Next() {
		v, err := scanValidator(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, *v)
	}
	return out, rows.Err()
}

// ListIdleValidators returns validators currently eligible to be handed a
// new IP to work on.
func (d *DB) ListIdleValidators(ctx context.Context) ([]Validator, error) {
	rows, err := d.QueryContext(ctx, `
		SELECT validator_id, hostname, os_port_id, state, current_ip_id, agent_version,
			last_heartbeat_at, created_at, updated_at
		FROM validators WHERE state=? ORDER BY validator_id
	`, ValidatorIdle)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Validator
	for rows.Next() {
		v, err := scanValidator(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, *v)
	}
	return out, rows.Err()
}

// ListStaleHeartbeats returns validators whose last heartbeat predates the
// given cutoff and that aren't already marked unreachable.
func (d *DB) ListStaleHeartbeats(ctx context.Context, cutoff time.Time) ([]Validator, error) {
	rows, err := d.QueryContext(ctx, `
		SELECT validator_id, hostname, os_port_id, state, current_ip_id, agent_version,
			last_heartbeat_at, created_at, updated_at
		FROM validators
		WHERE state != ? AND last_heartbeat_at IS NOT NULL AND last_heartbeat_at < ?
	`, ValidatorUnreachable, timeToDB(cutoff))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Validator
	for rows.Next() {
		v, err := scanValidator(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, *v)
	}
	return out, rows.Err()
}

func (d *DB) MarkValidatorUnreachable(ctx context.Context, validatorID string) error {
	_, err := d.ExecContext(ctx, `UPDATE validators SET state=?, updated_at=? WHERE validator_id=?`,
		ValidatorUnreachable, timeToDB(Now()), validatorID)
	return err
}

// FreeValidator returns a validator to idle with no assigned IP. Used after
// an IP finishes (success or failure) or is reclaimed by the lease sweep.
func (d *DB) FreeValidator(ctx context.Context, validatorID string) error {
	_, err := d.ExecContext(ctx, `
		UPDATE validators SET state=?, current_ip_id=NULL, updated_at=?
		WHERE validator_id=?
	`, ValidatorIdle, timeToDB(Now()), validatorID)
	return err
}

// AdminCreateValidator registers a brand-new validator via the admin API.
// Unlike RegisterValidator (used by the agent's self-registration call),
// this refuses to upsert over an existing row.
func (d *DB) AdminCreateValidator(ctx context.Context, validatorID, osPortID string) error {
	var exists int
	err := d.QueryRowContext(ctx, `SELECT 1 FROM validators WHERE validator_id=?`, validatorID).Scan(&exists)
	if err != nil && err != sql.ErrNoRows {
		return err
	}
	if err == nil {
		return fmt.Errorf("validator %q: %w", validatorID, ErrConflict)
	}

	now := timeToDB(Now())
	_, err = d.ExecContext(ctx, `
		INSERT INTO validators (validator_id, hostname, os_port_id, agent_version, state, created_at, updated_at)
		VALUES (?, '', ?, '', ?, ?, ?)
	`, validatorID, osPortID, ValidatorIdle, now, now)
	if err != nil {
		return fmt.Errorf("create validator: %w", err)
	}
	return nil
}

// AdminUpdateValidatorPort updates an existing validator's Neutron port ID.
func (d *DB) AdminUpdateValidatorPort(ctx context.Context, validatorID, osPortID string) error {
	res, err := d.ExecContext(ctx, `
		UPDATE validators SET os_port_id=?, updated_at=? WHERE validator_id=?
	`, osPortID, timeToDB(Now()), validatorID)
	if err != nil {
		return fmt.Errorf("update validator port: %w", err)
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return fmt.Errorf("validator %q: %w", validatorID, ErrNotFound)
	}
	return nil
}

// DeleteValidator removes a validator, refusing if it currently owns an IP.
// ip_queue.owner_validator_id is a permanent historical record (set once an
// IP is claimed, never cleared on completion — see RequeueOrFail/
// ReleaseFIP), so any validator that has ever processed an IP would
// otherwise violate the FK constraint on delete; those historical
// references are cleared in the same transaction once we've confirmed the
// validator isn't currently busy.
func (d *DB) DeleteValidator(ctx context.Context, validatorID string) error {
	tx, err := d.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	var currentIPID sql.NullInt64
	err = tx.QueryRowContext(ctx, `SELECT current_ip_id FROM validators WHERE validator_id=?`, validatorID).Scan(&currentIPID)
	if err == sql.ErrNoRows {
		return fmt.Errorf("validator %q: %w", validatorID, ErrNotFound)
	}
	if err != nil {
		return err
	}
	if currentIPID.Valid {
		return fmt.Errorf("validator %q owns ip_id %d: %w", validatorID, currentIPID.Int64, ErrBusy)
	}

	if _, err := tx.ExecContext(ctx, `UPDATE ip_queue SET owner_validator_id=NULL WHERE owner_validator_id=?`, validatorID); err != nil {
		return fmt.Errorf("clear historical ip_queue references: %w", err)
	}
	if _, err := tx.ExecContext(ctx, `DELETE FROM validators WHERE validator_id=?`, validatorID); err != nil {
		return fmt.Errorf("delete validator: %w", err)
	}
	return tx.Commit()
}

type rowScanner interface {
	Scan(dest ...interface{}) error
}

func scanValidator(row rowScanner) (*Validator, error) {
	var v Validator
	var currentIPID sql.NullInt64
	var lastHeartbeat sql.NullString
	var createdAt, updatedAt string
	if err := row.Scan(&v.ValidatorID, &v.Hostname, &v.OSPortID, &v.State, &currentIPID,
		&v.AgentVersion, &lastHeartbeat, &createdAt, &updatedAt); err != nil {
		return nil, err
	}
	if currentIPID.Valid {
		v.CurrentIPID = &currentIPID.Int64
	}
	hb, err := nullStringToTimePtr(lastHeartbeat)
	if err != nil {
		return nil, err
	}
	v.LastHeartbeatAt = hb
	if v.CreatedAt, err = dbToTime(createdAt); err != nil {
		return nil, err
	}
	if v.UpdatedAt, err = dbToTime(updatedAt); err != nil {
		return nil, err
	}
	return &v, nil
}