112 lines
3.9 KiB
Go
112 lines
3.9 KiB
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, 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.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 and its Compose and plain-container
|
|
// services removed and its queued and running jobs failed, in one transaction. 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) error {
|
|
tx, err := s.pool.Begin(ctx)
|
|
if err != nil {
|
|
return 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 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 fmt.Errorf("revoke agent: remove services: %w", err)
|
|
}
|
|
if _, err := tx.Exec(ctx, failActiveJobsSQL, id, "agent revoked", at); err != nil {
|
|
return fmt.Errorf("revoke agent: fail jobs: %w", err)
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return fmt.Errorf("revoke agent: commit: %w", err)
|
|
}
|
|
return 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, last_seen_at = $5 WHERE id = $1`,
|
|
id, h.Hostname, h.AgentVersion, h.SwarmManager, at); 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
|
|
}
|