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

// Package pocketbase is the PocketBase storage driver: Warp's data in a
// standalone PocketBase server, as collections, reached only through
// PocketBase's REST API — the way the postgres driver reaches PostgreSQL
// over its wire protocol.
//
// Each of the SQL drivers' tables is a collection of the same name and
// columns (see schema.go), so a PocketBase admin sees Warp's rooms, events
// and accounts as records. The stores read them through the superuser SQL
// endpoint with the same queries the SQL drivers use (sqldriver.go), and
// write them as record operations in batch requests (batch.go), which
// PocketBase runs as one transaction each. Media content goes into a file
// field, which is PocketBase's own storage: its pb_data/storage, or an S3
// bucket.
package pocketbase

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

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

func init() {
	storage.Register(storage.DriverPocketBase, Open)
}

// Database is the PocketBase implementation of storage.Database.
type Database struct {
	*base
	log *slog.Logger
	s3  storage.S3Config

	settings *Settings
	users    *Users
	rooms    *Rooms
	sync     *Sync
	fed      *Federation
	media    *Media
	keys     *Keys
	push     *Push
	admin    *Admin
	apps     *AppServices
	delayed  *DelayedEvents
	openID   *OpenIDTokens
	blobs    *blobStore
}

// Open connects to the PocketBase server at cfg.ConnectionString and signs
// in as the configured superuser.
//
// It configures the server for Warp before anything else — batch requests
// on, and the S3 bucket when one is configured — so a server that cannot
// be reached or refuses the account fails here, at startup.
func Open(ctx context.Context, cfg storage.Config, log *slog.Logger) (storage.Database, error) {
	c, err := newClient(cfg.ConnectionString, cfg.SuperuserEmail, cfg.SuperuserPassword)
	if err != nil {
		return nil, err
	}
	lock := &sync.RWMutex{}
	db := sql.OpenDB(&sqlConnector{c: c, lock: lock})
	// The connections hold nothing, so the pool only bounds how many reads
	// are in flight at once.
	maxConns := cfg.MaxOpenConns
	if maxConns <= 0 {
		maxConns = 16
	}
	db.SetMaxOpenConns(maxConns)

	b := &base{c: c, db: db, lock: lock}
	d := &Database{
		base:     b,
		log:      log,
		s3:       cfg.S3,
		settings: &Settings{base: b},
		users:    &Users{base: b},
		rooms:    &Rooms{base: b},
		sync:     &Sync{base: b},
		fed:      &Federation{base: b},
		media:    &Media{base: b},
		keys:     &Keys{base: b},
		push:     &Push{base: b},
		admin:    &Admin{base: b},
		apps:     &AppServices{base: b},
		delayed:  &DelayedEvents{base: b},
		openID:   &OpenIDTokens{base: b},
		blobs:    &blobStore{base: b},
	}
	if err := d.configure(ctx); err != nil {
		_ = db.Close()
		return nil, err
	}
	return d, nil
}

// configure applies the PocketBase settings Warp depends on.
//
// The S3 settings are written from Warp's configuration on every start
// rather than left to the PocketBase admin UI: Warp's config is the
// authority, and a second place for the bucket to be set would be a second
// place for it to disagree. A deployment without S3 in Warp's config leaves
// PocketBase's own setting as it is.
func (d *Database) configure(ctx context.Context) error {
	settings := map[string]any{
		"batch": map[string]any{
			"enabled":     true,
			"maxRequests": maxBatch,
			"timeout":     300,
			"maxBodySize": 0,
		},
	}
	if d.s3.Enabled {
		settings["s3"] = map[string]any{
			"enabled":        true,
			"bucket":         d.s3.Bucket,
			"region":         d.s3.Region,
			"endpoint":       d.s3.Endpoint,
			"accessKey":      d.s3.AccessKey,
			"secret":         d.s3.Secret,
			"forcePathStyle": d.s3.ForcePathStyle,
		}
	}
	if err := d.c.patchSettings(ctx, settings); err != nil {
		return fmt.Errorf("pocketbase: configure the server: %w", err)
	}
	return nil
}

// DB exposes the read-only SQL pool, for tests and diagnostics: queries run
// on the PocketBase server through its SQL endpoint, and anything else is
// refused.
func (d *Database) DB() *sql.DB { return d.db }

// Driver implements storage.Database.
func (d *Database) Driver() string { return storage.DriverPocketBase }

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

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

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

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

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

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

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

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

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

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

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

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

// DelayedEvents implements storage.Database.
func (d *Database) DelayedEvents() storage.DelayedEventStore { return d.delayed }

// OpenIDTokens implements storage.Database.
func (d *Database) OpenIDTokens() storage.OpenIDTokenStore { return d.openID }

// Blobs implements storage.BlobProvider: media content goes into
// PocketBase's storage.
func (d *Database) Blobs() storage.BlobStore { return d.blobs }

// schemaMetaKey is the warp_meta row holding the schema version applied.
const schemaMetaKey = "schema"

// Migrate implements storage.Database.
//
// There are no numbered migrations: the schema is declared in schema.go
// and imported as it is, which PocketBase applies as a diff — a new
// collection is created, a new field added, an unchanged one left alone.
// The version stamped afterwards is a hash of the declaration, so any
// change to it is a pending migration.
func (d *Database) Migrate(ctx context.Context) error {
	pending, err := d.PendingMigrations(ctx)
	if err != nil {
		return err
	}
	if len(pending) == 0 {
		d.log.Debug("database schema is up to date", "version", schemaVersion)
		return nil
	}

	collections := make([]map[string]any, len(schema))
	for i := range schema {
		collections[i] = schema[i].collectionJSON()
	}
	if err := d.c.importCollections(ctx, collections); err != nil {
		return fmt.Errorf("pocketbase: import the schema: %w", err)
	}

	err = d.write(ctx, func(t *tx) error {
		for table, names := range sequenceRows {
			for _, name := range names {
				found, err := t.exists(table, `name = ?`, name)
				if err != nil {
					return err
				}
				if !found {
					t.put(table, rec{"name": name, "value": int64(0)})
				}
			}
		}
		// Arguments a crashed query left behind.
		if _, err := t.delWhere(paramsCollection, `1 = 1`); err != nil {
			return err
		}
		t.put("warp_meta", rec{"meta_key": schemaMetaKey, "meta_value": schemaVersion})
		return nil
	})
	if err != nil {
		return fmt.Errorf("pocketbase: initialise the schema: %w", err)
	}
	d.log.Info("applied the database schema", "version", schemaVersion)
	return nil
}

// PendingMigrations implements storage.Database.
func (d *Database) PendingMigrations(ctx context.Context) ([]string, error) {
	pending := []string{"pocketbase_schema_" + schemaVersion}

	var n int
	if err := d.db.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'warp_meta'`).Scan(&n); err != nil {
		return nil, fmt.Errorf("pocketbase: read the schema version: %w", err)
	}
	if n == 0 {
		return pending, nil
	}
	var version string
	err := d.db.QueryRowContext(ctx,
		`SELECT meta_value FROM warp_meta WHERE meta_key = ?`, schemaMetaKey).Scan(&version)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return pending, nil
	case err != nil:
		return nil, fmt.Errorf("pocketbase: read the schema version: %w", err)
	case version != schemaVersion:
		return pending, nil
	}
	return nil, nil
}

// Settings is the server_settings store.
type Settings struct{ *base }

// Get implements storage.SettingsStore.
func (s *Settings) Get(ctx context.Context, key string) (string, bool, error) {
	var value string
	err := s.db.QueryRowContext(ctx,
		`SELECT setting_value FROM server_settings WHERE setting_key = ?`, key).Scan(&value)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return "", false, nil
	case err != nil:
		return "", false, fmt.Errorf("get setting %q: %w", key, err)
	}
	return value, true, nil
}

// GetAll implements storage.SettingsStore.
func (s *Settings) GetAll(ctx context.Context) (map[string]string, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT setting_key, setting_value FROM server_settings`)
	if err != nil {
		return nil, fmt.Errorf("list settings: %w", err)
	}
	defer func() { _ = rows.Close() }()

	out := map[string]string{}
	for rows.Next() {
		var k, v string
		if err := rows.Scan(&k, &v); err != nil {
			return nil, fmt.Errorf("scan setting: %w", err)
		}
		out[k] = v
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("list settings: %w", err)
	}
	return out, nil
}

// Set implements storage.SettingsStore.
func (s *Settings) Set(ctx context.Context, key, value string) error {
	err := s.write(ctx, func(t *tx) error {
		t.put("server_settings", rec{
			"setting_key":   key,
			"setting_value": value,
			"updated_at":    time.Now().UTC().Format(time.RFC3339Nano),
		})
		return nil
	})
	if err != nil {
		return fmt.Errorf("set setting %q: %w", key, err)
	}
	return nil
}

// Delete implements storage.SettingsStore.
func (s *Settings) Delete(ctx context.Context, key string) error {
	err := s.write(ctx, func(t *tx) error {
		_, err := t.delKey("server_settings", key)
		return err
	})
	if err != nil {
		return fmt.Errorf("delete setting %q: %w", key, err)
	}
	return nil
}

var (
	_ storage.Database      = (*Database)(nil)
	_ storage.BlobProvider  = (*Database)(nil)
	_ storage.SettingsStore = (*Settings)(nil)
)