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
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
//go:build integration
|
||||
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/shcizo/package-updater/internal/master/model"
|
||||
"github.com/shcizo/package-updater/internal/pin"
|
||||
"github.com/shcizo/package-updater/internal/protocol"
|
||||
)
|
||||
|
||||
func TestIdentity_CreatedOnceThenReused(t *testing.T) {
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
k1, err := s.LoadOrCreateIdentity(ctx, pin.GenerateKey)
|
||||
require.NoError(t, err)
|
||||
k2, err := s.LoadOrCreateIdentity(ctx, pin.GenerateKey)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, k1.Equal(k2), "a restarted master must keep its key or every agent breaks")
|
||||
}
|
||||
|
||||
func TestMigrationsAreIdempotent(t *testing.T) {
|
||||
s := openTestStore(t)
|
||||
require.NoError(t, migrate(context.Background(), s.pool))
|
||||
}
|
||||
|
||||
func TestAgents_CreateLookupRevoke(t *testing.T) {
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
created, err := s.CreateAgent(ctx, "node-2", "hash-2")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "node-2", created.Name)
|
||||
assert.Nil(t, created.RevokedAt)
|
||||
|
||||
found, err := s.AgentByTokenHash(ctx, "hash-2")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, created.ID, found.ID)
|
||||
|
||||
_, err = s.AgentByTokenHash(ctx, "nope")
|
||||
assert.ErrorIs(t, err, model.ErrNotFound)
|
||||
|
||||
require.NoError(t, s.RevokeAgent(ctx, created.ID, time.Now()))
|
||||
found, err = s.AgentByTokenHash(ctx, "hash-2")
|
||||
require.NoError(t, err)
|
||||
assert.NotNil(t, found.RevokedAt)
|
||||
}
|
||||
|
||||
func TestAgents_DuplicateNameIsConflict(t *testing.T) {
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
_, err := s.CreateAgent(ctx, "node-2", "h1")
|
||||
require.NoError(t, err)
|
||||
_, err = s.CreateAgent(ctx, "node-2", "h2")
|
||||
assert.ErrorIs(t, err, model.ErrConflict)
|
||||
}
|
||||
|
||||
func TestAgents_RecordHelloAndMarkSeen(t *testing.T) {
|
||||
s := openTestStore(t)
|
||||
ctx := context.Background()
|
||||
a, err := s.CreateAgent(ctx, "swarm-mgr", "h")
|
||||
require.NoError(t, err)
|
||||
|
||||
at := time.Date(2026, 10, 12, 12, 0, 0, 0, time.UTC)
|
||||
require.NoError(t, s.RecordHello(ctx, a.ID, protocol.Hello{
|
||||
Hostname: "swarm-mgr", AgentVersion: "v1", SwarmManager: true, SwarmClusterID: "c1",
|
||||
}, at))
|
||||
require.NoError(t, s.MarkSeen(ctx, a.ID, at.Add(time.Minute)))
|
||||
|
||||
agents, err := s.ListAgents(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, agents, 1)
|
||||
assert.Equal(t, "swarm-mgr", agents[0].Hostname)
|
||||
assert.Equal(t, "v1", agents[0].AgentVersion)
|
||||
assert.True(t, agents[0].SwarmManager)
|
||||
require.NotNil(t, agents[0].LastSeenAt)
|
||||
assert.True(t, agents[0].LastSeenAt.Equal(at.Add(time.Minute)))
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
//go:build integration
|
||||
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// openTestStore returns a Store migrated into a fresh schema that is dropped
|
||||
// when the test ends, so tests never see each other's rows.
|
||||
func openTestStore(t *testing.T) *Store {
|
||||
t.Helper()
|
||||
url := os.Getenv("DATABASE_URL")
|
||||
if url == "" {
|
||||
t.Skip("DATABASE_URL not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
admin, err := pgx.Connect(ctx, url)
|
||||
require.NoError(t, err)
|
||||
schema := fmt.Sprintf("test_%d", time.Now().UnixNano())
|
||||
_, err = admin.Exec(ctx, "CREATE SCHEMA "+schema)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
_, _ = admin.Exec(ctx, "DROP SCHEMA "+schema+" CASCADE")
|
||||
_ = admin.Close(ctx)
|
||||
})
|
||||
|
||||
cfg, err := pgxpool.ParseConfig(url)
|
||||
require.NoError(t, err)
|
||||
cfg.ConnConfig.RuntimeParams["search_path"] = schema
|
||||
s, err := openWithConfig(ctx, cfg)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(s.Close)
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"crypto/x509"
|
||||
"errors"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// LoadOrCreateIdentity returns the master's private key, creating and
|
||||
// storing one with generate on first start. Keeping it in Postgres means it
|
||||
// survives container recreation and is covered by database backups; losing
|
||||
// it would force every agent to be reconfigured with a new MASTER_KEY.
|
||||
func (s *Store) LoadOrCreateIdentity(ctx context.Context, generate func() (ed25519.PrivateKey, error)) (ed25519.PrivateKey, error) {
|
||||
fresh, err := generate()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
der, err := x509.MarshalPKCS8PrivateKey(fresh)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal master key: %w", err)
|
||||
}
|
||||
if _, err := s.pool.Exec(ctx,
|
||||
`INSERT INTO master_identity (id, private_key) VALUES (1, $1) ON CONFLICT (id) DO NOTHING`, der,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("store master key: %w", err)
|
||||
}
|
||||
var stored []byte
|
||||
if err := s.pool.QueryRow(ctx, `SELECT private_key FROM master_identity WHERE id = 1`).Scan(&stored); err != nil {
|
||||
return nil, fmt.Errorf("load master key: %w", err)
|
||||
}
|
||||
parsed, err := x509.ParsePKCS8PrivateKey(stored)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse master key: %w", err)
|
||||
}
|
||||
key, ok := parsed.(ed25519.PrivateKey)
|
||||
if !ok {
|
||||
return nil, errors.New("stored master key is not ed25519")
|
||||
}
|
||||
return key, nil
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"embed"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
//go:embed migrations/*.sql
|
||||
var migrationsFS embed.FS
|
||||
|
||||
// migrate applies every embedded migration not yet recorded in
|
||||
// schema_migrations, in file-name order, each in its own transaction.
|
||||
func migrate(ctx context.Context, pool *pgxpool.Pool) error {
|
||||
if _, err := pool.Exec(ctx, `CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
version text PRIMARY KEY,
|
||||
applied_at timestamptz NOT NULL DEFAULT now()
|
||||
)`); err != nil {
|
||||
return fmt.Errorf("create schema_migrations: %w", err)
|
||||
}
|
||||
entries, err := fs.ReadDir(migrationsFS, "migrations")
|
||||
if err != nil {
|
||||
return fmt.Errorf("read migrations: %w", err)
|
||||
}
|
||||
for _, e := range entries {
|
||||
name := e.Name()
|
||||
var applied bool
|
||||
if err := pool.QueryRow(ctx,
|
||||
`SELECT EXISTS (SELECT 1 FROM schema_migrations WHERE version = $1)`, name,
|
||||
).Scan(&applied); err != nil {
|
||||
return fmt.Errorf("check migration %s: %w", name, err)
|
||||
}
|
||||
if applied {
|
||||
continue
|
||||
}
|
||||
sql, err := migrationsFS.ReadFile("migrations/" + name)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read migration %s: %w", name, err)
|
||||
}
|
||||
if err := applyMigration(ctx, pool, name, string(sql)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func applyMigration(ctx context.Context, pool *pgxpool.Pool, name, sql string) error {
|
||||
tx, err := pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin migration %s: %w", name, err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
// Exec without arguments uses the simple protocol, which allows the
|
||||
// several statements a migration file contains.
|
||||
if _, err := tx.Exec(ctx, sql); err != nil {
|
||||
return fmt.Errorf("migration %s: %w", name, err)
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `INSERT INTO schema_migrations (version) VALUES ($1)`, name); err != nil {
|
||||
return fmt.Errorf("record migration %s: %w", name, err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return fmt.Errorf("commit migration %s: %w", name, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
CREATE TABLE master_identity (
|
||||
id int PRIMARY KEY CHECK (id = 1),
|
||||
private_key bytea NOT NULL,
|
||||
created_at timestamptz NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE TABLE agents (
|
||||
id bigserial PRIMARY KEY,
|
||||
name text NOT NULL UNIQUE,
|
||||
token_hash text NOT NULL UNIQUE,
|
||||
hostname text NOT NULL DEFAULT '',
|
||||
agent_version text NOT NULL DEFAULT '',
|
||||
swarm_manager boolean NOT NULL DEFAULT false,
|
||||
last_seen_at timestamptz,
|
||||
revoked_at timestamptz,
|
||||
created_at timestamptz NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE TABLE services (
|
||||
id bigserial PRIMARY KEY,
|
||||
kind text NOT NULL CHECK (kind IN ('compose', 'swarm', 'container')),
|
||||
location_key text NOT NULL UNIQUE,
|
||||
agent_id bigint REFERENCES agents (id),
|
||||
compose_project text NOT NULL DEFAULT '',
|
||||
compose_service text NOT NULL DEFAULT '',
|
||||
working_dir text NOT NULL DEFAULT '',
|
||||
swarm_cluster_id text NOT NULL DEFAULT '',
|
||||
swarm_service text NOT NULL DEFAULT '',
|
||||
container_name text NOT NULL DEFAULT '',
|
||||
image_repo text NOT NULL,
|
||||
actual_ref text NOT NULL,
|
||||
actual_digest text NOT NULL DEFAULT '',
|
||||
status text NOT NULL,
|
||||
self boolean NOT NULL DEFAULT false,
|
||||
update_mode text NOT NULL DEFAULT 'manual' CHECK (update_mode IN ('manual', 'auto')),
|
||||
first_seen_at timestamptz NOT NULL,
|
||||
updated_at timestamptz NOT NULL,
|
||||
removed_at timestamptz
|
||||
);
|
||||
|
||||
CREATE TABLE sessions (
|
||||
id text PRIMARY KEY,
|
||||
subject text NOT NULL,
|
||||
name text NOT NULL DEFAULT '',
|
||||
email text NOT NULL DEFAULT '',
|
||||
csrf_token text NOT NULL,
|
||||
expires_at timestamptz NOT NULL
|
||||
);
|
||||
@@ -0,0 +1,41 @@
|
||||
// Package store persists the master's state in Postgres.
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
// Store is the Postgres-backed storage for the master.
|
||||
type Store struct {
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
// Open connects to Postgres and applies pending migrations.
|
||||
func Open(ctx context.Context, url string) (*Store, error) {
|
||||
cfg, err := pgxpool.ParseConfig(url)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse DATABASE_URL: %w", err)
|
||||
}
|
||||
return openWithConfig(ctx, cfg)
|
||||
}
|
||||
|
||||
func openWithConfig(ctx context.Context, cfg *pgxpool.Config) (*Store, error) {
|
||||
pool, err := pgxpool.NewWithConfig(ctx, cfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("connect postgres: %w", err)
|
||||
}
|
||||
if err := migrate(ctx, pool); err != nil {
|
||||
pool.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &Store{pool: pool}, nil
|
||||
}
|
||||
|
||||
// Close releases the connection pool.
|
||||
func (s *Store) Close() { s.pool.Close() }
|
||||
|
||||
// Ping checks that Postgres answers. Used by /healthz.
|
||||
func (s *Store) Ping(ctx context.Context) error { return s.pool.Ping(ctx) }
|
||||
Reference in New Issue
Block a user