package-updater/internal/master/store/agents.go

package store

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"

	"github.com/shcizo/package-updater/internal/master/model"
	"github.com/shcizo/package-updater/internal/protocol"
)

const agentColumns = `id, name, hostname, agent_version, swarm_manager, swarm_cluster_id, last_seen_at, revoked_at, created_at`

func scanAgent(row pgx.Row) (model.Agent, error) {
	var a model.Agent
	err := row.Scan(&a.ID, &a.Name, &a.Hostname, &a.AgentVersion, &a.SwarmManager, &a.SwarmClusterID, &a.LastSeenAt, &a.RevokedAt, &a.CreatedAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return model.Agent{}, model.ErrNotFound
	}
	return a, err
}

// CreateAgent registers an agent. Returns model.ErrConflict if an active
// (non-revoked) agent already has the name.
func (s *Store) CreateAgent(ctx context.Context, name, tokenHash string) (model.Agent, error) {
	a, err := scanAgent(s.pool.QueryRow(ctx,
		`INSERT INTO agents (name, token_hash) VALUES ($1, $2) RETURNING `+agentColumns, name, tokenHash))
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) && pgErr.Code == "23505" {
		return model.Agent{}, model.ErrConflict
	}
	if err != nil {
		return model.Agent{}, fmt.Errorf("create agent: %w", err)
	}
	return a, nil
}

// ListAgents returns all agents, including revoked ones, ordered by name.
func (s *Store) ListAgents(ctx context.Context) ([]model.Agent, error) {
	rows, err := s.pool.Query(ctx, `SELECT `+agentColumns+` FROM agents ORDER BY name`)
	if err != nil {
		return nil, fmt.Errorf("list agents: %w", err)
	}
	defer rows.Close()
	var out []model.Agent
	for rows.Next() {
		a, err := scanAgent(rows)
		if err != nil {
			return nil, fmt.Errorf("scan agent: %w", err)
		}
		out = append(out, a)
	}
	return out, rows.Err()
}

// AgentByTokenHash finds the agent a token belongs to. Revoked agents are
// returned too; the caller decides. Returns model.ErrNotFound if none.
func (s *Store) AgentByTokenHash(ctx context.Context, hash string) (model.Agent, error) {
	return scanAgent(s.pool.QueryRow(ctx, `SELECT `+agentColumns+` FROM agents WHERE token_hash = $1`, hash))
}

// RevokeAgent marks the agent revoked, its Compose and plain-container
// services removed and its queued and running jobs failed, in one
// transaction, and returns the IDs of the jobs it failed. Swarm rows are
// shared with other managers and left alone. Its token stops working
// immediately; closing a live connection is the hub's job.
func (s *Store) RevokeAgent(ctx context.Context, id int64, at time.Time) ([]int64, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, fmt.Errorf("revoke agent: begin: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()

	if _, err := tx.Exec(ctx,
		`UPDATE agents SET revoked_at = $2 WHERE id = $1 AND revoked_at IS NULL`, id, at); err != nil {
		return nil, fmt.Errorf("revoke agent: %w", err)
	}
	if _, err := tx.Exec(ctx,
		`UPDATE services SET removed_at = $2
		 WHERE agent_id = $1 AND kind IN ('compose', 'container') AND removed_at IS NULL`, id, at); err != nil {
		return nil, fmt.Errorf("revoke agent: remove services: %w", err)
	}
	rows, err := tx.Query(ctx, failActiveJobsSQL, id, "agent revoked", at)
	if err != nil {
		return nil, fmt.Errorf("revoke agent: fail jobs: %w", err)
	}
	failed, err := pgx.CollectRows(rows, pgx.RowTo[int64])
	if err != nil {
		return nil, fmt.Errorf("revoke agent: fail jobs: %w", err)
	}
	if err := tx.Commit(ctx); err != nil {
		return nil, fmt.Errorf("revoke agent: commit: %w", err)
	}
	return failed, nil
}

// RecordHello stores what the agent reported about itself on connect.
func (s *Store) RecordHello(ctx context.Context, id int64, h protocol.Hello, at time.Time) error {
	if _, err := s.pool.Exec(ctx,
		`UPDATE agents SET hostname = $2, agent_version = $3, swarm_manager = $4, swarm_cluster_id = $6, last_seen_at = $5 WHERE id = $1`,
		id, h.Hostname, h.AgentVersion, h.SwarmManager, at, h.SwarmClusterID); err != nil {
		return fmt.Errorf("record hello: %w", err)
	}
	return nil
}

// MarkSeen updates last_seen_at, called when an agent disconnects.
func (s *Store) MarkSeen(ctx context.Context, id int64, at time.Time) error {
	if _, err := s.pool.Exec(ctx, `UPDATE agents SET last_seen_at = $2 WHERE id = $1`, id, at); err != nil {
		return fmt.Errorf("mark agent seen: %w", err)
	}
	return nil
}