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

// Package pocketbase is the PocketBase storage driver: Warp on an embedded
// PocketBase app, for "appliance" deployments that want one binary and one
// data directory.
//
// PLAN.md section 4 set out the shape, and open question 4 asked how far
// to take it. The answer here is the cheaper of the two: the event DAG and
// everything else Warp stores are a poor fit for PocketBase's record
// model, so the Warp tables are the sqlite driver's own, migrated by Warp
// and written through the same shared store, on PocketBase's underlying
// SQLite connection rather than through its collection API. What PocketBase
// adds is everything around that: its data directory and its system
// tables, which a PocketBase admin or a later Warp feature can build on, and
// its filesystem abstraction, which puts media content on local disk or in
// an S3 bucket.
//
// The result is at parity with the sqlite driver by construction — it is
// the same SQL on the same engine — and the storage tests run against it
// to keep it that way.
package pocketbase

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"path/filepath"

	"github.com/pocketbase/dbx"
	"github.com/pocketbase/pocketbase/core"
	_ "github.com/pocketbase/pocketbase/migrations" // PocketBase's system migrations; without them its bootstrap has no tables
	"github.com/pocketbase/pocketbase/tools/filesystem"

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

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

// mediaPrefix is where media content lives in PocketBase's filesystem,
// apart from the files of PocketBase's own collections.
const mediaPrefix = "warp_media"

// Database is the PocketBase implementation of storage.Database.
type Database struct {
	*sqlstore.Store
	app   *core.BaseApp
	blobs *blobStore
}

// Open bootstraps a PocketBase app in the directory named by
// cfg.ConnectionString — PocketBase's pb_data — and opens the Warp store on
// its database.
func Open(ctx context.Context, cfg storage.Config, log *slog.Logger) (storage.Database, error) {
	dir := cfg.ConnectionString
	if dir == "" {
		return nil, errors.New("pocketbase: connection_string must name the data directory")
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	app := core.NewBaseApp(core.BaseAppConfig{
		DataDir:          dir,
		DataMaxOpenConns: cfg.MaxOpenConns,
		DataMaxIdleConns: cfg.MaxIdleConns,
	})
	if err := app.Bootstrap(); err != nil {
		_ = app.ResetBootstrapState()
		return nil, fmt.Errorf("pocketbase: bootstrap %s: %w", dir, err)
	}
	if cfg.S3.Enabled {
		// Applied in memory on every start rather than saved: the Warp
		// config is the authority, and PocketBase's own settings would
		// otherwise be a second place for the bucket to be set, and to
		// disagree.
		app.Settings().S3 = core.S3Config{
			Enabled:        true,
			Bucket:         cfg.S3.Bucket,
			Region:         cfg.S3.Region,
			Endpoint:       cfg.S3.Endpoint,
			AccessKey:      cfg.S3.AccessKey,
			Secret:         cfg.S3.Secret,
			ForcePathStyle: cfg.S3.ForcePathStyle,
		}
	}

	// The writer pool: one connection, the way the sqlite driver runs, and
	// the way PocketBase itself writes. Warp's store reads in the same
	// transactions it writes in, so it is given the one pool for both.
	writer, ok := app.NonconcurrentDB().(*dbx.DB)
	if !ok {
		_ = app.ResetBootstrapState()
		return nil, errors.New("pocketbase: the app's database is not a *dbx.DB")
	}

	opts := sqlite.StoreOptions(writer.DB(), log)
	opts.DriverName = storage.DriverPocketBase
	store, err := sqlstore.New(opts)
	if err != nil {
		_ = app.ResetBootstrapState()
		return nil, err
	}
	return &Database{
		Store: store,
		app:   app,
		blobs: &blobStore{app: app, tmpDir: filepath.Join(app.DataDir(), ".warp_tmp")},
	}, nil
}

// App is the PocketBase app, for code that wants PocketBase's own
// features: its collections, its auth, its realtime.
func (d *Database) App() core.App { return d.app }

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

// Close implements storage.Database. The app owns the connection pools.
func (d *Database) Close() error { return d.app.ResetBootstrapState() }

// blobStore is storage.BlobStore over PocketBase's filesystem: the local
// pb_data/storage directory, or the S3 bucket the settings name.
//
// The filesystem can only be written a whole file at a time, and a blob's
// name is its hash, which is not known until the last byte has been read.
// So an upload is spooled to a temporary file first, hashed on the way, and
// handed over once it is complete — which is also what keeps a half-written
// object from ever appearing under a hash.
type blobStore struct {
	app    *core.BaseApp
	tmpDir string
}

func blobKey(sum string) (string, bool) {
	if len(sum) != sha256.Size*2 {
		return "", false
	}
	for _, c := range sum {
		if (c < '0' || c > '9') && (c < 'a' || c > 'f') {
			return "", false
		}
	}
	return mediaPrefix + "/" + sum[0:2] + "/" + sum[2:4] + "/" + sum, true
}

// fs opens the filesystem for one operation. It is cheap for both kinds,
// and an instance per operation is what lets each carry its own context.
func (b *blobStore) fs(ctx context.Context) (*filesystem.System, error) {
	fsys, err := b.app.NewFilesystem()
	if err != nil {
		return nil, fmt.Errorf("open the media filesystem: %w", err)
	}
	fsys.SetContext(ctx)
	return fsys, nil
}

// Put implements storage.BlobStore.
func (b *blobStore) Put(ctx context.Context, r io.Reader, limit int64) (storage.BlobInfo, error) {
	if err := os.MkdirAll(b.tmpDir, 0o750); err != nil {
		return storage.BlobInfo{}, fmt.Errorf("create the media spool directory: %w", err)
	}
	tmp, err := os.CreateTemp(b.tmpDir, "upload-*")
	if err != nil {
		return storage.BlobInfo{}, fmt.Errorf("create a temporary file: %w", err)
	}
	defer func() {
		_ = tmp.Close()
		_ = os.Remove(tmp.Name())
	}()

	hash := sha256.New()
	n, err := io.Copy(io.MultiWriter(tmp, hash), &ctxReader{ctx: ctx, r: io.LimitReader(r, limit+1)})
	if err != nil {
		return storage.BlobInfo{}, fmt.Errorf("store content: %w", err)
	}
	if n > limit {
		return storage.BlobInfo{}, storage.ErrBlobTooLarge
	}
	if err := tmp.Close(); err != nil {
		return storage.BlobInfo{}, fmt.Errorf("store content: %w", err)
	}

	info := storage.BlobInfo{SHA256: hex.EncodeToString(hash.Sum(nil)), Size: n}
	key, _ := blobKey(info.SHA256)

	fsys, err := b.fs(ctx)
	if err != nil {
		return storage.BlobInfo{}, err
	}
	defer func() { _ = fsys.Close() }()

	if exists, err := fsys.Exists(key); err == nil && exists {
		// The same content is already stored.
		return info, nil
	}
	file, err := filesystem.NewFileFromPath(tmp.Name())
	if err != nil {
		return storage.BlobInfo{}, fmt.Errorf("store content: %w", err)
	}
	if err := fsys.UploadFile(file, key); err != nil {
		return storage.BlobInfo{}, fmt.Errorf("store content: %w", err)
	}
	return info, nil
}

// Open implements storage.BlobStore.
func (b *blobStore) Open(ctx context.Context, sum string) (io.ReadSeekCloser, error) {
	key, ok := blobKey(sum)
	if !ok {
		return nil, storage.ErrNoSuchBlob
	}
	fsys, err := b.fs(ctx)
	if err != nil {
		return nil, err
	}
	reader, err := fsys.GetReader(key)
	if err != nil {
		_ = fsys.Close()
		if errors.Is(err, filesystem.ErrNotFound) {
			return nil, storage.ErrNoSuchBlob
		}
		return nil, fmt.Errorf("open media content: %w", err)
	}
	return &blobReader{ReadSeekCloser: reader, fsys: fsys}, nil
}

// Delete implements storage.BlobStore.
func (b *blobStore) Delete(ctx context.Context, sum string) error {
	key, ok := blobKey(sum)
	if !ok {
		return nil
	}
	fsys, err := b.fs(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = fsys.Close() }()
	if err := fsys.Delete(key); err != nil && !errors.Is(err, filesystem.ErrNotFound) {
		return fmt.Errorf("delete media content: %w", err)
	}
	return nil
}

// blobReader closes the filesystem with the reader it opened.
type blobReader struct {
	io.ReadSeekCloser
	fsys *filesystem.System
}

func (r *blobReader) Close() error {
	return errors.Join(r.ReadSeekCloser.Close(), r.fsys.Close())
}

// ctxReader stops a copy when its context ends.
type ctxReader struct {
	ctx context.Context
	r   io.Reader
}

func (c *ctxReader) Read(p []byte) (int, error) {
	if err := c.ctx.Err(); err != nil {
		return 0, err
	}
	return c.r.Read(p)
}

var (
	_ storage.Database     = (*Database)(nil)
	_ storage.BlobProvider = (*Database)(nil)
	_ storage.BlobStore    = (*blobStore)(nil)
)