Warp/server/internal/storage/sqlstore/sqlstore.go

// Package sqlstore holds the parts of a storage driver that are identical
// for every database/sql backend: the migration plumbing and the
// server_settings table. The sqlite and postgres drivers differ only in
// their DSN handling, placeholder style and upsert syntax, all of which
// they pass in here.
package sqlstore

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io/fs"
	"log/slog"

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

// Options describe one database/sql backend.
type Options struct {
	// DriverName is the Warp driver name, e.g. "sqlite".
	DriverName string

	// DB is the already-opened pool.
	DB *sql.DB

	// Dialect supplies the migration bookkeeping statements.
	Dialect migrations.Dialect

	// MigrationsFS and MigrationsDir locate the embedded .sql files.
	MigrationsFS  fs.FS
	MigrationsDir string

	// Settings holds the statements for the server_settings table.
	Settings SettingsQueries

	// Rebind translates the `?` placeholders that the shared component
	// stores are written with into this driver's style. SQLite needs no
	// translation and leaves it nil; PostgreSQL turns them into $1, $2...
	Rebind func(query string) string

	// IsUniqueViolation recognises this driver's unique-constraint error, so
	// that a taken user name can be told apart from a database failure
	// without a check-then-insert race.
	IsUniqueViolation func(error) bool

	// Logger is used for migration progress.
	Logger *slog.Logger
}

// Store is the shared implementation of storage.Database. Drivers embed it
// and add their own component stores as those are built.
type Store struct {
	opts     Options
	db       *sql.DB
	log      *slog.Logger
	settings *Settings
	users    *Users
	rooms    *Rooms
	sync     *Sync
	fed      *Federation
	media    *Media
	keys     *Keys
	push     *Push
	admin    *Admin
	apps     *AppServices
}

// New builds a Store. It does not touch the database.
func New(opts Options) (*Store, error) {
	if opts.DB == nil {
		return nil, errors.New("sqlstore: DB must not be nil")
	}
	if opts.Logger == nil {
		opts.Logger = slog.Default()
	}
	return &Store{
		opts:     opts,
		db:       opts.DB,
		log:      opts.Logger,
		settings: NewSettings(opts.DB, opts.Settings),
		users:    NewUsers(opts.DB, opts.Rebind, opts.IsUniqueViolation),
		rooms:    NewRooms(opts.DB, opts.Rebind, opts.IsUniqueViolation),
		sync:     NewSync(opts.DB, opts.Rebind, opts.IsUniqueViolation),
		fed:      NewFederation(opts.DB, opts.Rebind),
		media:    NewMedia(opts.DB, opts.Rebind, opts.IsUniqueViolation),
		keys:     NewKeys(opts.DB, opts.Rebind),
		push:     NewPush(opts.DB, opts.Rebind),
		admin:    NewAdmin(opts.DB, opts.Rebind, opts.IsUniqueViolation),
		apps:     NewAppServices(opts.DB, opts.Rebind),
	}, nil
}

// DB exposes the pool for driver-specific code and tests.
func (s *Store) DB() *sql.DB { return s.db }

// Driver implements storage.Database.
func (s *Store) Driver() string { return s.opts.DriverName }

// Ping implements storage.Database.
func (s *Store) Ping(ctx context.Context) error { return s.db.PingContext(ctx) }

// Close implements storage.Database.
func (s *Store) Close() error { return s.db.Close() }

// Settings implements storage.Database.
func (s *Store) Settings() storage.SettingsStore { return s.settings }

// Users implements storage.Database.
func (s *Store) Users() storage.UserStore { return s.users }

// Rooms implements storage.Database.
func (s *Store) Rooms() storage.RoomStore { return s.rooms }

// Sync implements storage.Database.
func (s *Store) Sync() storage.SyncStore { return s.sync }

// Federation implements storage.Database.
func (s *Store) Federation() storage.FederationStore { return s.fed }

// Media implements storage.Database.
func (s *Store) Media() storage.MediaStore { return s.media }

// Keys implements storage.Database.
func (s *Store) Keys() storage.KeyStore { return s.keys }

// Push implements storage.Database.
func (s *Store) Push() storage.PushStore { return s.push }

// Admin implements storage.Database.
func (s *Store) Admin() storage.AdminStore { return s.admin }

// AppServices implements storage.Database.
func (s *Store) AppServices() storage.AppServiceStore { return s.apps }

// Migrate implements storage.Database.
func (s *Store) Migrate(ctx context.Context) error {
	ms, err := s.migrations()
	if err != nil {
		return err
	}
	applied, err := migrations.Status(ctx, s.db, s.opts.Dialect)
	if err != nil {
		return err
	}
	if err := migrations.CheckNotAhead(applied, ms); err != nil {
		return err
	}
	n, err := migrations.Run(ctx, s.db, s.opts.Dialect, ms, s.log)
	if err != nil {
		return err
	}
	if n == 0 {
		s.log.Debug("database schema is up to date", "migrations", len(ms))
	}
	return nil
}

// PendingMigrations implements storage.Database.
func (s *Store) PendingMigrations(ctx context.Context) ([]string, error) {
	ms, err := s.migrations()
	if err != nil {
		return nil, err
	}
	pending, err := migrations.Pending(ctx, s.db, s.opts.Dialect, ms)
	if err != nil {
		return nil, err
	}
	out := make([]string, 0, len(pending))
	for _, m := range pending {
		out = append(out, fmt.Sprintf("%04d_%s", m.Version, m.Name))
	}
	return out, nil
}

func (s *Store) migrations() ([]migrations.Migration, error) {
	return migrations.Load(s.opts.MigrationsFS, s.opts.MigrationsDir)
}