Warp/server/internal/storage/storage.go

// Package storage defines the interfaces every Warp storage driver
// implements, plus the registry that turns a config into a Database.
//
// Drivers live in sub-packages (sqlite, postgres, pocketbase) and register
// themselves from an init function, the database/sql way. Nothing in this
// package imports a driver, so components can depend on the interfaces
// without pulling in every database engine.
package storage

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"maps"
	"slices"
	"sync"
	"time"
)

// Driver names. These are the values accepted by database.driver in the
// config file.
const (
	DriverSQLite     = "sqlite"
	DriverPostgres   = "postgres"
	DriverPocketBase = "pocketbase"
)

// ErrNotImplemented is returned by a driver that is compiled in but not
// finished — currently PocketBase.
var ErrNotImplemented = errors.New("not implemented by this storage driver")

// ErrNoSuchSetting is returned by SettingsStore.Get for an unknown key.
var ErrNoSuchSetting = errors.New("no such setting")

// Config is the subset of the file configuration a driver needs.
type Config struct {
	Driver           string
	ConnectionString string
	MaxOpenConns     int
	MaxIdleConns     int
	ConnMaxLifetime  time.Duration
}

// Database is the handle every component gets. Component-specific stores
// (userapi, roomserver, syncapi...) hang off it as they are implemented.
type Database interface {
	// Driver reports which driver this is, for logs and admin output.
	Driver() string

	// Ping checks the connection is alive; used by /healthz.
	Ping(ctx context.Context) error

	// Migrate applies pending migrations.
	Migrate(ctx context.Context) error

	// PendingMigrations lists migrations that Migrate would apply.
	PendingMigrations(ctx context.Context) ([]string, error)

	// Settings is the server_settings key/value store backing the runtime
	// settings service (federation and registration switches).
	Settings() SettingsStore

	// Users is the userapi store: accounts, devices, tokens, account data.
	Users() UserStore

	// Rooms is the roomserver store: rooms, events, room state, aliases.
	Rooms() RoomStore

	// Sync is the syncapi store: receipts, to-device messages, presence and
	// uploaded filters.
	Sync() SyncStore

	// Federation is the federation store: other servers' keys, the
	// outbound queue and its backoff, and the inbound transaction log.
	Federation() FederationStore

	// Close releases the connection pool.
	Close() error
}

// SettingsStore persists the handful of settings that are mutable at
// runtime through the admin API.
type SettingsStore interface {
	// Get returns the raw value for key. The boolean reports presence.
	Get(ctx context.Context, key string) (string, bool, error)

	// GetAll returns every setting.
	GetAll(ctx context.Context) (map[string]string, error)

	// Set writes key, overwriting any previous value.
	Set(ctx context.Context, key, value string) error

	// Delete removes key, reverting it to its configured default. Deleting
	// a key that is not there is not an error.
	Delete(ctx context.Context, key string) error
}

// OpenFunc opens a database for a driver.
type OpenFunc func(ctx context.Context, cfg Config, log *slog.Logger) (Database, error)

var (
	registryMu sync.RWMutex
	registry   = map[string]OpenFunc{}
)

// Register makes a driver available to Open. It panics on a duplicate
// registration, which can only be a programming error.
func Register(name string, open OpenFunc) {
	registryMu.Lock()
	defer registryMu.Unlock()
	if _, dup := registry[name]; dup {
		panic("storage: driver " + name + " registered twice")
	}
	registry[name] = open
}

// Drivers lists the registered driver names, sorted.
func Drivers() []string {
	registryMu.RLock()
	defer registryMu.RUnlock()
	return slices.Sorted(maps.Keys(registry))
}

// Open connects using the driver named in cfg.
func Open(ctx context.Context, cfg Config, log *slog.Logger) (Database, error) {
	registryMu.RLock()
	open, ok := registry[cfg.Driver]
	registryMu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("storage: unknown driver %q (compiled in: %v)", cfg.Driver, Drivers())
	}
	if log == nil {
		log = slog.Default()
	}
	db, err := open(ctx, cfg, log.With("driver", cfg.Driver))
	if err != nil {
		return nil, fmt.Errorf("storage: open %s: %w", cfg.Driver, err)
	}
	return db, nil
}