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'
	// comes back: to idle when it holds no address, to assigned when it still
	// does (its current_ip_id is kept, so it must not be handed another).
	_, err = d.ExecContext(ctx, `
		UPDATE validators SET
			state = CASE WHEN current_ip_id IS NULL THEN ? ELSE ? END,
			updated_at=?
		WHERE validator_id=? AND state IN (?, ?)
	`, ValidatorIdle, ValidatorAssigned, now, validatorID, ValidatorUnregistered, ValidatorUnreachable)
	if err != nil {
		return fmt.Errorf("register validator (reactivate): %w", err)
	}
	return nil
}

// Heartbeat records a sign of life. A validator that was marked unreachable
// returns to idle if it holds no address, but to assigned if it still does:
// it was only silent (for example busy with slow checks), and handing it a
// second address while it works on the first would leave the second one
// without an owner that can ever pick it up.
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 (CASE WHEN current_ip_id IS NULL THEN ? ELSE ? END) ELSE state END
		WHERE validator_id=?
	`, now, now, ValidatorUnreachable, ValidatorIdle, ValidatorAssigned, 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
}

// freeValidatorSQL releases a validator from the address it holds. It only
// applies when the validator's current address is the one being released
// (args: now, validator id, ip id): a late release of an old address must not
// free a validator that has already moved on to another one. An unreachable
// validator stays unreachable until its next heartbeat, so a dead validator
// is not handed new addresses just because its lease was reclaimed.
var freeValidatorSQL = fmt.Sprintf(`
	UPDATE validators SET current_ip_id=NULL,
		state = CASE WHEN state='%s' THEN state ELSE '%s' END,
		updated_at=?
	WHERE validator_id=? AND current_ip_id=?`, ValidatorUnreachable, ValidatorIdle)

// FreeValidator releases a validator from the given address (see
// freeValidatorSQL). Used after an IP finishes (success or failure) or is
// reclaimed by the lease sweep.
func (d *DB) FreeValidator(ctx context.Context, validatorID string, ipID int64) error {
	_, err := d.ExecContext(ctx, freeValidatorSQL, timeToDB(Now()), validatorID, ipID)
	return err
}

// ReconcileValidators repairs validators whose state disagrees with the
// queue: a validator pointing at an address that no longer exists, is
// finished, or belongs to another validator is released; an "assigned"
// validator that holds nothing goes back to idle. The invariants normally
// hold by construction (every change is one transaction); this heals what a
// crash or an older version left behind. It returns the number of validators
// repaired.
func (d *DB) ReconcileValidators(ctx context.Context) (int64, error) {
	now := timeToDB(Now())
	res, err := d.ExecContext(ctx, fmt.Sprintf(`
		UPDATE validators SET current_ip_id=NULL,
			state = CASE WHEN state='%s' THEN state ELSE '%s' END,
			updated_at=?
		WHERE current_ip_id IS NOT NULL AND NOT EXISTS (
			SELECT 1 FROM ip_queue q
			WHERE q.id = validators.current_ip_id
			  AND q.owner_validator_id = validators.validator_id
			  AND q.state NOT IN ('%s','%s','%s'))`,
		ValidatorUnreachable, ValidatorIdle, IPDone, IPFailed, IPOccupied), now)
	if err != nil {
		return 0, fmt.Errorf("reconcile validators: %w", err)
	}
	n, _ := res.RowsAffected()
	res, err = d.ExecContext(ctx, `
		UPDATE validators SET state=?, updated_at=?
		WHERE state=? AND current_ip_id IS NULL`, ValidatorIdle, now, ValidatorAssigned)
	if err != nil {
		return n, fmt.Errorf("reconcile validators: %w", err)
	}
	m, _ := res.RowsAffected()
	return n + m, nil
}

// 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
}