llmgw/internal/store/federation.go

package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"time"
)

// FederationPeer is another llmgw cluster whose models this one offers.
type FederationPeer struct {
	Name string
	URL  string
	// KeyFile holds the key this cluster uses at the peer, on the control
	// plane: a path, never the key, for the reason an external model's is.
	KeyFile string
	// SendTags makes requests sent to the peer carry x-llmgw-tags naming
	// this cluster and the sending project, so the peer can attribute them.
	SendTags bool
	// Cluster is the peer's own cluster id, as it last said it: what a
	// request that has passed through the peer carries in its hop header.
	Cluster string
	// Models is what the peer last offered, as JSON.
	Models    json.RawMessage
	SyncedAt  time.Time
	CreatedAt time.Time
	CreatedBy string
}

// PutFederationPeer adds a peer, or replaces one whole: its address, key
// file, tags setting and what it offers. Who added it, and when, stay as
// they were first written.
func (s *DB) PutFederationPeer(ctx context.Context, p FederationPeer) error {
	if p.Name == "" || p.URL == "" {
		return fmt.Errorf("a federation peer needs a name and an address")
	}
	if p.CreatedAt.IsZero() {
		p.CreatedAt = time.Now()
	}
	models := string(p.Models)
	if models == "" {
		models = "[]"
	}
	if !json.Valid([]byte(models)) {
		return fmt.Errorf("peer %q: its models are not valid JSON", p.Name)
	}
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO federation_peers (name, url, key_file, send_tags, cluster, models, synced_at, created_at, created_by)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
		 ON CONFLICT (name) DO UPDATE SET url = excluded.url, key_file = excluded.key_file,
		   send_tags = excluded.send_tags, cluster = excluded.cluster, models = excluded.models,
		   synced_at = excluded.synced_at`,
		p.Name, p.URL, p.KeyFile, boolInt(p.SendTags), p.Cluster, models, syncedMillis(p.SyncedAt),
		millis(p.CreatedAt), p.CreatedBy)
	return err
}

// SetFederationPeerModels records what a peer offers now, and its cluster
// id, without touching how it is reached. ErrNotFound when the peer was
// removed meanwhile: a sync that finishes after a removal must not bring
// the peer back.
func (s *DB) SetFederationPeerModels(ctx context.Context, name, cluster string, models json.RawMessage, at time.Time) error {
	if !json.Valid(models) {
		return fmt.Errorf("peer %q: its models are not valid JSON", name)
	}
	res, err := s.db.ExecContext(ctx,
		`UPDATE federation_peers SET cluster = ?, models = ?, synced_at = ? WHERE name = ?`,
		cluster, string(models), syncedMillis(at), name)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// DeleteFederationPeer removes a peer; removing one that is not there is
// ErrNotFound, so a mistyped name is said rather than silently accepted.
func (s *DB) DeleteFederationPeer(ctx context.Context, name string) error {
	res, err := s.db.ExecContext(ctx, `DELETE FROM federation_peers WHERE name = ?`, name)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

// FederationPeers lists the peers by name.
func (s *DB) FederationPeers(ctx context.Context) ([]FederationPeer, error) {
	return queryAll(ctx, s.db, `SELECT name, url, key_file, send_tags, cluster, models, synced_at, created_at, created_by
		FROM federation_peers ORDER BY name`, nil, func(rows *sql.Rows) (FederationPeer, error) {
		var p FederationPeer
		var tags int
		var models string
		var synced, created int64
		err := rows.Scan(&p.Name, &p.URL, &p.KeyFile, &tags, &p.Cluster, &models, &synced, &created, &p.CreatedBy)
		p.SendTags, p.Models = tags == 1, json.RawMessage(models)
		p.CreatedAt = fromMillis(created)
		if synced > 0 {
			p.SyncedAt = fromMillis(synced)
		}
		return p, err
	})
}

// FederationPeer is one peer by name; ErrNotFound when there is none.
func (s *DB) FederationPeer(ctx context.Context, name string) (FederationPeer, error) {
	peers, err := s.FederationPeers(ctx)
	if err != nil {
		return FederationPeer{}, err
	}
	for _, p := range peers {
		if p.Name == name {
			return p, nil
		}
	}
	return FederationPeer{}, ErrNotFound
}

// syncedMillis is 0 for a peer never synced, which reads back as no time
// rather than as a date in 1754.
func syncedMillis(t time.Time) int64 {
	if t.IsZero() {
		return 0
	}
	return t.UnixMilli()
}