feat(master): postgres store with migrations, master identity and agents
This commit is contained in:
@@ -0,0 +1,91 @@
|
||||
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 the name is taken.
|
||||
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 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 {
|
||||
if _, err := s.pool.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)
|
||||
}
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user