241 lines
8.9 KiB
Go
241 lines
8.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/imageref"
|
|
"github.com/shcizo/package-updater/internal/master/model"
|
|
"github.com/shcizo/package-updater/internal/protocol"
|
|
)
|
|
|
|
type execer interface {
|
|
Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
|
|
}
|
|
|
|
// lockActiveAgent takes a share lock on the agent's row inside tx and fails
|
|
// with model.ErrRevoked if the agent is revoked or gone. RevokeAgent's
|
|
// UPDATE of the row waits for the lock, so an inventory write either commits
|
|
// before the revoke (whose service cleanup then sees and removes its rows)
|
|
// or runs after it and is refused. Without this, an in-flight write could
|
|
// bring a revoked agent's services back.
|
|
func lockActiveAgent(ctx context.Context, tx pgx.Tx, agentID int64) error {
|
|
var one int
|
|
err := tx.QueryRow(ctx,
|
|
`SELECT 1 FROM agents WHERE id = $1 AND revoked_at IS NULL FOR SHARE`, agentID).Scan(&one)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return model.ErrRevoked
|
|
}
|
|
if err != nil {
|
|
return fmt.Errorf("lock agent: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
const upsertServiceSQL = `
|
|
INSERT INTO services (
|
|
kind, location_key, agent_id, compose_project, compose_service, working_dir, config_files,
|
|
built, swarm_cluster_id, swarm_service, container_name, image_repo, actual_ref, actual_digest,
|
|
status, self, first_seen_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $17)
|
|
ON CONFLICT (location_key) DO UPDATE SET
|
|
agent_id = EXCLUDED.agent_id,
|
|
working_dir = EXCLUDED.working_dir,
|
|
config_files = EXCLUDED.config_files,
|
|
built = EXCLUDED.built,
|
|
image_repo = EXCLUDED.image_repo,
|
|
actual_ref = EXCLUDED.actual_ref,
|
|
actual_digest = EXCLUDED.actual_digest,
|
|
status = EXCLUDED.status,
|
|
self = EXCLUDED.self,
|
|
updated_at = EXCLUDED.updated_at,
|
|
removed_at = NULL`
|
|
|
|
func upsertService(ctx context.Context, db execer, agentID int64, clusterID string, svc protocol.Service, at time.Time) error {
|
|
clusterCol := ""
|
|
if svc.Kind == protocol.KindSwarm {
|
|
clusterCol = clusterID
|
|
}
|
|
_, err := db.Exec(ctx, upsertServiceSQL,
|
|
string(svc.Kind), model.LocationKey(agentID, clusterID, svc.Key()), agentID,
|
|
svc.ComposeProject, svc.ComposeService, svc.WorkingDir, svc.ConfigFiles, svc.Built,
|
|
clusterCol, svc.SwarmService, svc.ContainerName,
|
|
imageref.Parse(svc.Image).Repo, svc.Image, svc.Digest,
|
|
svc.Status, svc.Self, at)
|
|
if err != nil {
|
|
return fmt.Errorf("upsert service %s: %w", svc.Key(), err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ApplySnapshot makes the stored inventory for this agent's scope equal to
|
|
// services: everything listed is upserted, everything in scope but not
|
|
// listed is marked removed. Returns model.ErrRevoked without writing if the
|
|
// agent is revoked. Scope is the agent's own Compose services and
|
|
// containers, plus the cluster's Swarm services if the agent is a manager.
|
|
// A manager without a cluster ID is treated as a non-manager, since its
|
|
// Swarm rows could not be told apart from other agents' rows.
|
|
func (s *Store) ApplySnapshot(ctx context.Context, agentID int64, h protocol.Hello, services []protocol.Service, at time.Time) error {
|
|
tx, err := s.pool.Begin(ctx)
|
|
if err != nil {
|
|
return fmt.Errorf("begin snapshot: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
if err := lockActiveAgent(ctx, tx, agentID); err != nil {
|
|
return err
|
|
}
|
|
|
|
isManager := h.SwarmManager && h.SwarmClusterID != ""
|
|
keys := make([]string, 0, len(services))
|
|
for _, svc := range services {
|
|
if svc.Kind == protocol.KindSwarm && !isManager {
|
|
continue
|
|
}
|
|
if err := upsertService(ctx, tx, agentID, h.SwarmClusterID, svc, at); err != nil {
|
|
return err
|
|
}
|
|
keys = append(keys, model.LocationKey(agentID, h.SwarmClusterID, svc.Key()))
|
|
}
|
|
if _, err := tx.Exec(ctx, `
|
|
UPDATE services SET removed_at = $1
|
|
WHERE removed_at IS NULL
|
|
AND NOT (location_key = ANY($2))
|
|
AND ((kind IN ('compose', 'container') AND agent_id = $3)
|
|
OR (kind = 'swarm' AND $4 AND swarm_cluster_id = $5))`,
|
|
at, keys, agentID, isManager, h.SwarmClusterID); err != nil {
|
|
return fmt.Errorf("mark removed services: %w", err)
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return fmt.Errorf("commit snapshot: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// UpsertService stores one changed or new service. Returns
|
|
// model.ErrRevoked without writing if the agent is revoked.
|
|
func (s *Store) UpsertService(ctx context.Context, agentID int64, clusterID string, svc protocol.Service, at time.Time) error {
|
|
tx, err := s.pool.Begin(ctx)
|
|
if err != nil {
|
|
return fmt.Errorf("begin upsert: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
if err := lockActiveAgent(ctx, tx, agentID); err != nil {
|
|
return err
|
|
}
|
|
if err := upsertService(ctx, tx, agentID, clusterID, svc, at); err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return fmt.Errorf("commit upsert: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// RemoveService marks the service with the agent-relative key as removed.
|
|
func (s *Store) RemoveService(ctx context.Context, agentID int64, clusterID, agentKey string, at time.Time) error {
|
|
if _, err := s.pool.Exec(ctx,
|
|
`UPDATE services SET removed_at = $2 WHERE location_key = $1 AND removed_at IS NULL`,
|
|
model.LocationKey(agentID, clusterID, agentKey), at); err != nil {
|
|
return fmt.Errorf("remove service: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
const serviceColumns = `s.id, s.kind, s.location_key, s.agent_id, COALESCE(a.name, ''),
|
|
s.compose_project, s.compose_service, s.working_dir, s.config_files, s.built, s.swarm_cluster_id,
|
|
s.swarm_service, s.container_name, s.image_repo, s.desired_ref, s.actual_ref, s.actual_digest,
|
|
s.status, s.self, s.update_mode, s.available_ref, s.available_source, s.available_at,
|
|
s.first_seen_at, s.updated_at, s.removed_at`
|
|
|
|
func scanService(row pgx.Row) (model.Service, error) {
|
|
var m model.Service
|
|
var kind string
|
|
err := row.Scan(&m.ID, &kind, &m.LocationKey, &m.AgentID, &m.AgentName,
|
|
&m.ComposeProject, &m.ComposeService, &m.WorkingDir, &m.ConfigFiles, &m.Built, &m.SwarmClusterID,
|
|
&m.SwarmService, &m.ContainerName, &m.ImageRepo, &m.DesiredRef, &m.ActualRef, &m.ActualDigest,
|
|
&m.Status, &m.Self, &m.UpdateMode, &m.AvailableRef, &m.AvailableSource, &m.AvailableAt,
|
|
&m.FirstSeenAt, &m.UpdatedAt, &m.RemovedAt)
|
|
m.Kind = protocol.Kind(kind)
|
|
return m, err
|
|
}
|
|
|
|
// ListServices returns all services not marked removed: machine-bound
|
|
// services grouped by agent name first, Swarm services last.
|
|
func (s *Store) ListServices(ctx context.Context) ([]model.Service, error) {
|
|
return s.queryServices(ctx, "list services", "", nil)
|
|
}
|
|
|
|
// ServicesByRepo returns the services not marked removed whose image repo
|
|
// equals repo (case-sensitive), in the same order as ListServices.
|
|
func (s *Store) ServicesByRepo(ctx context.Context, repo string) ([]model.Service, error) {
|
|
return s.queryServices(ctx, "services by repo", "AND s.image_repo = $1", []any{repo})
|
|
}
|
|
|
|
func (s *Store) queryServices(ctx context.Context, what, extraWhere string, args []any) ([]model.Service, error) {
|
|
rows, err := s.pool.Query(ctx, `
|
|
SELECT `+serviceColumns+`
|
|
FROM services s LEFT JOIN agents a ON a.id = s.agent_id
|
|
WHERE s.removed_at IS NULL `+extraWhere+`
|
|
ORDER BY s.kind = 'swarm', COALESCE(a.name, ''), s.location_key`, args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%s: %w", what, err)
|
|
}
|
|
defer rows.Close()
|
|
var out []model.Service
|
|
for rows.Next() {
|
|
m, err := scanService(rows)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("scan service: %w", err)
|
|
}
|
|
out = append(out, m)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
// SetAvailable records a newer ref for the service, replacing any earlier
|
|
// suggestion.
|
|
func (s *Store) SetAvailable(ctx context.Context, id int64, a model.Available) error {
|
|
if _, err := s.pool.Exec(ctx, `
|
|
UPDATE services SET available_ref = $2, available_source = $3, available_at = $4
|
|
WHERE id = $1`, id, a.Ref, a.Source, a.At); err != nil {
|
|
return fmt.Errorf("set available: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// SetUpdateMode sets the service's update mode (model.ModeManual or
|
|
// model.ModeAuto). Returns model.ErrNotFound if the service does not exist
|
|
// or is marked removed.
|
|
func (s *Store) SetUpdateMode(ctx context.Context, id int64, mode string) error {
|
|
tag, err := s.pool.Exec(ctx,
|
|
`UPDATE services SET update_mode = $2 WHERE id = $1 AND removed_at IS NULL`, id, mode)
|
|
if err != nil {
|
|
return fmt.Errorf("set update mode: %w", err)
|
|
}
|
|
if tag.RowsAffected() == 0 {
|
|
return model.ErrNotFound
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetService returns one service, including a removed one (RemovedAt set).
|
|
// Returns model.ErrNotFound if the ID does not exist.
|
|
func (s *Store) GetService(ctx context.Context, id int64) (model.Service, error) {
|
|
m, err := scanService(s.pool.QueryRow(ctx, `
|
|
SELECT `+serviceColumns+`
|
|
FROM services s LEFT JOIN agents a ON a.id = s.agent_id
|
|
WHERE s.id = $1`, id))
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return model.Service{}, model.ErrNotFound
|
|
}
|
|
if err != nil {
|
|
return model.Service{}, fmt.Errorf("get service: %w", err)
|
|
}
|
|
return m, nil
|
|
}
|