Warp/server/internal/storage/pocketbase/federation.go

package pocketbase

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// Federation is the PocketBase implementation of storage.FederationStore.
type Federation struct{ *base }

// StoreServerKeys implements storage.FederationStore.
func (s *Federation) StoreServerKeys(ctx context.Context, keys []storage.ServerKey) error {
	if len(keys) == 0 {
		return nil
	}
	return s.write(ctx, func(t *tx) error {
		for _, k := range keys {
			t.put("federation_server_keys", rec{
				"server_name": k.ServerName, "key_id": k.KeyID, "public_key": k.PublicKey,
				"valid_until_ts": k.ValidUntilTS, "expired_ts": k.ExpiredTS, "fetched_at_ms": k.FetchedAtMS,
			})
		}
		return nil
	})
}

// ServerKeys implements storage.FederationStore.
func (s *Federation) ServerKeys(ctx context.Context, serverName string, keyIDs []string) ([]storage.ServerKey, error) {
	query := `
		SELECT server_name, key_id, public_key, valid_until_ts, expired_ts, fetched_at_ms
		FROM federation_server_keys WHERE server_name = ?`
	args := []any{serverName}
	if len(keyIDs) > 0 {
		query += ` AND key_id IN ` + inList
		args = append(args, jsonList(keyIDs))
	}
	query += ` ORDER BY key_id`

	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("read the keys of %q: %w", serverName, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.ServerKey
	for rows.Next() {
		var k storage.ServerKey
		if err := rows.Scan(&k.ServerName, &k.KeyID, &k.PublicKey,
			&k.ValidUntilTS, &k.ExpiredTS, &k.FetchedAtMS); err != nil {
			return nil, fmt.Errorf("read the keys of %q: %w", serverName, err)
		}
		out = append(out, k)
	}
	return out, rows.Err()
}

// Enqueue implements storage.FederationStore.
func (s *Federation) Enqueue(ctx context.Context, items []storage.QueuedItem) ([]storage.QueuedItem, error) {
	if len(items) == 0 {
		return nil, nil
	}
	out := make([]storage.QueuedItem, len(items))
	copy(out, items)
	err := s.write(ctx, func(t *tx) error {
		top, err := t.advance(federationSequences, seqQueue, int64(len(out)))
		if err != nil {
			return err
		}
		first := top - int64(len(out)) + 1
		for i := range out {
			out[i].ID = first + int64(i)
			item := out[i]
			t.put("federation_queue", rec{
				"queue_id": item.ID, "destination": item.Destination, "kind": item.Kind,
				"room_id": item.RoomID, "event_id": item.EventID,
				"payload": string(item.Payload), "queued_at_ms": item.QueuedAtMS,
			})
		}
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("enqueue %d items: %w", len(items), err)
	}
	return out, nil
}

// Queued implements storage.FederationStore.
func (s *Federation) Queued(ctx context.Context, destination, kind string, limit int) ([]storage.QueuedItem, error) {
	if limit <= 0 {
		return nil, nil
	}
	rows, err := s.db.QueryContext(ctx, `
		SELECT queue_id, destination, kind, room_id, event_id, payload, queued_at_ms
		FROM federation_queue
		WHERE destination = ? AND kind = ?
		ORDER BY queue_id
		LIMIT ?`, destination, kind, limit)
	if err != nil {
		return nil, fmt.Errorf("read the queue for %q: %w", destination, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.QueuedItem
	for rows.Next() {
		var (
			item    storage.QueuedItem
			payload string
		)
		if err := rows.Scan(&item.ID, &item.Destination, &item.Kind, &item.RoomID,
			&item.EventID, &payload, &item.QueuedAtMS); err != nil {
			return nil, fmt.Errorf("read the queue for %q: %w", destination, err)
		}
		item.Payload = []byte(payload)
		out = append(out, item)
	}
	return out, rows.Err()
}

// Dequeue implements storage.FederationStore.
func (s *Federation) Dequeue(ctx context.Context, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}
	err := s.write(ctx, func(t *tx) error {
		_, err := t.delWhere("federation_queue", `queue_id IN `+inList, jsonInts(ids))
		return err
	})
	if err != nil {
		return fmt.Errorf("dequeue %d items: %w", len(ids), err)
	}
	return nil
}

// QueuedDestinations implements storage.FederationStore.
func (s *Federation) QueuedDestinations(ctx context.Context) ([]string, error) {
	out, err := s.queryStrings(ctx, `SELECT DISTINCT destination FROM federation_queue ORDER BY destination`)
	if err != nil {
		return nil, fmt.Errorf("read the queued destinations: %w", err)
	}
	return out, nil
}

const destinationColumns = `
	server_name, failures, retry_at_ms, last_success_ms, last_failure_ms, last_error`

// Destination implements storage.FederationStore.
func (s *Federation) Destination(ctx context.Context, serverName string) (storage.Destination, error) {
	var d storage.Destination
	err := s.db.QueryRowContext(ctx, `SELECT `+destinationColumns+`
		FROM federation_destinations WHERE server_name = ?`, serverName).
		Scan(&d.ServerName, &d.Failures, &d.RetryAtMS, &d.LastSuccessMS, &d.LastFailureMS, &d.LastError)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.Destination{ServerName: serverName}, nil
	case err != nil:
		return storage.Destination{}, fmt.Errorf("read destination %q: %w", serverName, err)
	}
	return d, nil
}

// SetDestination implements storage.FederationStore.
func (s *Federation) SetDestination(ctx context.Context, d storage.Destination) error {
	err := s.write(ctx, func(t *tx) error {
		t.put("federation_destinations", rec{
			"server_name": d.ServerName, "failures": d.Failures, "retry_at_ms": d.RetryAtMS,
			"last_success_ms": d.LastSuccessMS, "last_failure_ms": d.LastFailureMS, "last_error": d.LastError,
		})
		return nil
	})
	if err != nil {
		return fmt.Errorf("record destination %q: %w", d.ServerName, err)
	}
	return nil
}

// Destinations implements storage.FederationStore: every server with a
// health row or something queued.
func (s *Federation) Destinations(ctx context.Context) ([]storage.DestinationSummary, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT n.server_name,
			COALESCE(d.failures, 0), COALESCE(d.retry_at_ms, 0),
			COALESCE(d.last_success_ms, 0), COALESCE(d.last_failure_ms, 0),
			COALESCE(d.last_error, ''),
			(SELECT COUNT(*) FROM federation_queue q WHERE q.destination = n.server_name AND q.kind = 'pdu'),
			(SELECT COUNT(*) FROM federation_queue q WHERE q.destination = n.server_name AND q.kind = 'edu')
		FROM (
			SELECT server_name FROM federation_destinations
			UNION
			SELECT DISTINCT destination FROM federation_queue
		) n
		LEFT JOIN federation_destinations d ON d.server_name = n.server_name
		ORDER BY n.server_name`)
	if err != nil {
		return nil, fmt.Errorf("read the destinations: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.DestinationSummary
	for rows.Next() {
		var d storage.DestinationSummary
		if err := rows.Scan(&d.ServerName, &d.Failures, &d.RetryAtMS, &d.LastSuccessMS,
			&d.LastFailureMS, &d.LastError, &d.QueuedPDUs, &d.QueuedEDUs); err != nil {
			return nil, fmt.Errorf("read the destinations: %w", err)
		}
		out = append(out, d)
	}
	return out, rows.Err()
}

// InboundTransaction implements storage.FederationStore.
func (s *Federation) InboundTransaction(ctx context.Context, origin, txnID string) ([]byte, bool, error) {
	var response string
	err := s.db.QueryRowContext(ctx, `SELECT response_json FROM federation_inbound_txns
		WHERE origin = ? AND txn_id = ?`, origin, txnID).Scan(&response)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return nil, false, nil
	case err != nil:
		return nil, false, fmt.Errorf("read transaction %q from %q: %w", txnID, origin, err)
	}
	return []byte(response), true, nil
}

// RecordInboundTransaction implements storage.FederationStore. A
// transaction already recorded keeps its first response.
func (s *Federation) RecordInboundTransaction(ctx context.Context, origin, txnID string, response []byte, atMS int64) error {
	err := s.write(ctx, func(t *tx) error {
		found, err := t.exists("federation_inbound_txns", `origin = ? AND txn_id = ?`, origin, txnID)
		if err != nil || found {
			return err
		}
		t.put("federation_inbound_txns", rec{
			"origin": origin, "txn_id": txnID, "response_json": string(response), "received_at_ms": atMS,
		})
		return nil
	})
	if err != nil {
		return fmt.Errorf("record transaction %q from %q: %w", txnID, origin, err)
	}
	return nil
}

var _ storage.FederationStore = (*Federation)(nil)