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() }