rollcall/cmd/rollcall/store.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

const schema = `
CREATE TABLE IF NOT EXISTS machines (
  id              text PRIMARY KEY,
  uuid            text,
  serial          text,
  hostname        text,
  manufacturer    text,
  product         text,
  state           text NOT NULL DEFAULT 'new',
  first_seen      timestamptz NOT NULL DEFAULT now(),
  last_seen       timestamptz,
  commissioned_at timestamptz,
  macs            jsonb NOT NULL DEFAULT '[]',
  bmc             jsonb,
  bmc_secret      bytea,
  inventory       jsonb
);
CREATE INDEX IF NOT EXISTS machines_uuid ON machines(uuid);
CREATE INDEX IF NOT EXISTS machines_macs ON machines USING gin(macs);
CREATE TABLE IF NOT EXISTS events (
  id         bigserial PRIMARY KEY,
  ts         timestamptz NOT NULL DEFAULT now(),
  machine_id text NOT NULL,
  event      text NOT NULL,
  detail     text
);
CREATE INDEX IF NOT EXISTS events_machine ON events(machine_id, ts);
CREATE TABLE IF NOT EXISTS profiles (
  name        text PRIMARY KEY,
  image       text NOT NULL,
  append      text,
  overlays    jsonb NOT NULL DEFAULT '[]',
  name_format text,
  vars        jsonb NOT NULL DEFAULT '{}',
  created     timestamptz NOT NULL DEFAULT now()
);
ALTER TABLE machines ADD COLUMN IF NOT EXISTS profile text REFERENCES profiles(name);
ALTER TABLE machines ADD COLUMN IF NOT EXISTS node_name text UNIQUE;
ALTER TABLE machines ADD COLUMN IF NOT EXISTS mgmt_ip text;
ALTER TABLE machines ADD COLUMN IF NOT EXISTS ib_ip text;
ALTER TABLE machines ADD COLUMN IF NOT EXISTS node_token_hash text;
ALTER TABLE machines ADD COLUMN IF NOT EXISTS node_token bytea;
ALTER TABLE machines ADD COLUMN IF NOT EXISTS host_key bytea;
ALTER TABLE machines ADD COLUMN IF NOT EXISTS booted_at timestamptz;
CREATE INDEX IF NOT EXISTS machines_node_token ON machines(node_token_hash);
ALTER TABLE machines ADD COLUMN IF NOT EXISTS os text;
ALTER TABLE machines ADD COLUMN IF NOT EXISTS os_version text;
ALTER TABLE machines ADD COLUMN IF NOT EXISTS deployed_at timestamptz;
CREATE TABLE IF NOT EXISTS jobs (
  id         text PRIMARY KEY,
  machine_id text NOT NULL,
  kind       text NOT NULL,
  state      text NOT NULL DEFAULT 'pending',
  spec       jsonb NOT NULL,
  progress   jsonb,
  token_hash text,
  token      bytea,
  created    timestamptz NOT NULL DEFAULT now(),
  started    timestamptz,
  finished   timestamptz,
  error      text
);
CREATE INDEX IF NOT EXISTS jobs_machine ON jobs(machine_id, created);
CREATE INDEX IF NOT EXISTS jobs_token ON jobs(token_hash);
`

var errNotFound = errors.New("not found")

type Store struct{ pool *pgxpool.Pool }

type Machine struct {
	ID             string          `json:"id"`
	UUID           *string         `json:"uuid"`
	Serial         *string         `json:"serial"`
	Hostname       *string         `json:"hostname"`
	Manufacturer   *string         `json:"manufacturer"`
	Product        *string         `json:"product"`
	State          string          `json:"state"`
	FirstSeen      time.Time       `json:"first_seen"`
	LastSeen       *time.Time      `json:"last_seen"`
	CommissionedAt *time.Time      `json:"commissioned_at"`
	MACs           []string        `json:"macs"`
	BMC            json.RawMessage `json:"bmc,omitempty"`
	Inventory      json.RawMessage `json:"inventory,omitempty"`
	Profile        *string         `json:"profile"`
	NodeName       *string         `json:"node_name"`
	MgmtIP         *string         `json:"mgmt_ip"`
	IBIP           *string         `json:"ib_ip"`
	BootedAt       *time.Time      `json:"booted_at"`
	OS             *string         `json:"os"`
	OSVersion      *string         `json:"os_version"`
	DeployedAt     *time.Time      `json:"deployed_at"`
	bmcSecret      []byte
	nodeToken      []byte
	hostKey        []byte
}

type Event struct {
	TS     time.Time `json:"ts"`
	Event  string    `json:"event"`
	Detail string    `json:"detail"`
}

func openStore(ctx context.Context, url string) (*Store, error) {
	pool, err := pgxpool.New(ctx, url)
	if err != nil {
		return nil, err
	}
	if _, err := pool.Exec(ctx, schema); err != nil {
		return nil, fmt.Errorf("schema: %w", err)
	}
	return &Store{pool: pool}, nil
}

const machineCols = `id, uuid, serial, hostname, manufacturer, product, state, first_seen, last_seen,
	commissioned_at, macs, bmc, bmc_secret, inventory, profile, node_name, mgmt_ip, ib_ip, node_token, host_key, booted_at,
	os, os_version, deployed_at`

func scanMachine(row pgx.Row, full bool) (*Machine, error) {
	m := &Machine{}
	var macs []byte
	err := row.Scan(&m.ID, &m.UUID, &m.Serial, &m.Hostname, &m.Manufacturer, &m.Product, &m.State,
		&m.FirstSeen, &m.LastSeen, &m.CommissionedAt, &macs, &m.BMC, &m.bmcSecret, &m.Inventory,
		&m.Profile, &m.NodeName, &m.MgmtIP, &m.IBIP, &m.nodeToken, &m.hostKey, &m.BootedAt,
		&m.OS, &m.OSVersion, &m.DeployedAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, errNotFound
	}
	if err != nil {
		return nil, err
	}
	_ = json.Unmarshal(macs, &m.MACs)
	if !full {
		m.Inventory = nil
	}
	return m, nil
}

func (s *Store) List(ctx context.Context) ([]*Machine, error) {
	rows, err := s.pool.Query(ctx, `SELECT `+machineCols+` FROM machines ORDER BY first_seen`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []*Machine
	for rows.Next() {
		m, err := scanMachine(rows, true)
		if err != nil {
			return nil, err
		}
		out = append(out, m)
	}
	if out == nil {
		out = []*Machine{}
	}
	return out, rows.Err()
}

func (s *Store) Get(ctx context.Context, id string) (*Machine, error) {
	return scanMachine(s.pool.QueryRow(ctx, `SELECT `+machineCols+` FROM machines WHERE id=$1`, id), true)
}

func (s *Store) Events(ctx context.Context, id string) ([]Event, error) {
	rows, err := s.pool.Query(ctx, `SELECT ts, event, coalesce(detail,'') FROM events WHERE machine_id=$1 ORDER BY ts`, id)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	evs := []Event{}
	for rows.Next() {
		var e Event
		if err := rows.Scan(&e.TS, &e.Event, &e.Detail); err != nil {
			return nil, err
		}
		evs = append(evs, e)
	}
	return evs, rows.Err()
}

func (s *Store) AddEvent(ctx context.Context, id, event, detail string) error {
	_, err := s.pool.Exec(ctx, `INSERT INTO events(machine_id, event, detail) VALUES ($1,$2,$3)`, id, event, detail)
	return err
}

// UpsertCommissioned stores a fresh inventory; returns true when the machine was new.
func (s *Store) UpsertCommissioned(ctx context.Context, m *Machine, secret []byte) (bool, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return false, err
	}
	defer tx.Rollback(ctx)
	macs, _ := json.Marshal(m.MACs)
	tag, err := tx.Exec(ctx, `INSERT INTO machines(id, state) VALUES ($1,'new') ON CONFLICT (id) DO NOTHING`, m.ID)
	if err != nil {
		return false, err
	}
	isNew := tag.RowsAffected() == 1
	if isNew {
		if _, err := tx.Exec(ctx, `INSERT INTO events(machine_id, event) VALUES ($1,'enlisted')`, m.ID); err != nil {
			return false, err
		}
	}
	// a node that already has a profile goes back to 'stateless', everything else to 'ready'
	_, err = tx.Exec(ctx, `UPDATE machines SET uuid=$2, serial=$3, hostname=$4, manufacturer=$5, product=$6,
		state=CASE WHEN profile IS NULL THEN 'ready' ELSE 'stateless' END, booted_at=NULL,
		last_seen=now(), commissioned_at=now(), macs=$7, bmc=$8, bmc_secret=$9, inventory=$10 WHERE id=$1`,
		m.ID, m.UUID, m.Serial, m.Hostname, m.Manufacturer, m.Product, macs, m.BMC, secret, m.Inventory)
	if err != nil {
		return false, err
	}
	return isNew, tx.Commit(ctx)
}

func (s *Store) SetState(ctx context.Context, id, state string) error {
	tag, err := s.pool.Exec(ctx, `UPDATE machines SET state=$2 WHERE id=$1`, id, state)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return errNotFound
	}
	return nil
}

func (s *Store) Delete(ctx context.Context, id string) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)
	if _, err := tx.Exec(ctx, `DELETE FROM events WHERE machine_id=$1`, id); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `DELETE FROM jobs WHERE machine_id=$1`, id); err != nil {
		return err
	}
	tag, err := tx.Exec(ctx, `DELETE FROM machines WHERE id=$1`, id)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return errNotFound
	}
	return tx.Commit(ctx)
}