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

package pocketbase

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"sync"
)

// base is what every store shares: the client, the read-only SQL pool and
// the write lock.
//
// Writes are serialised by the lock, in this process, the way the sqlite
// driver serialises them onto one connection: PocketBase is SQLite and takes
// one writer at a time anyway, and holding the lock across a store call's
// reads and its batch is what makes "read, decide, write" safe without a
// transaction that spans requests. It does mean one Warp server per
// PocketBase.
type base struct {
	c    *client
	db   *sql.DB
	lock *sync.RWMutex
}

// rec is one row's columns, by name.
type rec map[string]any

// maxBatch is how many operations go into one batch request. PocketBase is
// told to allow this many (see configure); a store call that writes more is
// split, and each part is atomic on its own.
const maxBatch = 1000

// tx collects the writes of one store call.
type tx struct {
	ctx context.Context
	b   *base

	// seqs are the counters this call has read or advanced, by table and
	// name, and seqOrder the order they were first touched in.
	seqs     map[string]int64
	seqOrder []string

	ops     []batchRequest
	deleted map[string]bool
}

// write runs fn with the write lock held and sends what it collected.
//
// The counters it advanced are written first, in the first batch, so that a
// call split over several batches that fails part-way leaves a gap in a
// stream — which every reader already copes with — and never a position
// handed out twice.
func (b *base) write(ctx context.Context, fn func(*tx) error) error {
	b.lock.Lock()
	defer b.lock.Unlock()

	t := &tx{ctx: withWriteLock(ctx), b: b, seqs: map[string]int64{}, deleted: map[string]bool{}}
	if err := fn(t); err != nil {
		return err
	}
	return t.commit()
}

// read runs fn with the read lock held, so that its queries see one state:
// no write lands between them. The SQL drivers get that from a
// transaction; here it is the lock, since every write takes it.
func (b *base) read(ctx context.Context, fn func(ctx context.Context) error) error {
	b.lock.RLock()
	defer b.lock.RUnlock()
	return fn(withWriteLock(ctx))
}

func (t *tx) commit() error {
	ops := make([]batchRequest, 0, len(t.seqOrder)+len(t.ops))
	for _, key := range t.seqOrder {
		table, name, _ := strings.Cut(key, "\x00")
		ops = append(ops, putRequest(table, rec{"name": name, "value": t.seqs[key]}))
	}
	ops = append(ops, t.ops...)
	for len(ops) > 0 {
		n := min(len(ops), maxBatch)
		if err := t.b.c.batch(t.ctx, ops[:n]); err != nil {
			return err
		}
		ops = ops[n:]
	}
	return nil
}

// putRequest is an upsert of one row: created if its key is new, otherwise
// updated in the columns given and left alone in the rest.
func putRequest(collection string, row rec) batchRequest {
	tb, ok := tables[collection]
	if !ok {
		panic("pocketbase: no collection " + collection)
	}
	key := make([]any, len(tb.key))
	for i, col := range tb.key {
		v, ok := row[col]
		if !ok {
			panic(fmt.Sprintf("pocketbase: a %s row without its key column %s", collection, col))
		}
		key[i] = v
	}
	return patchRequest(collection, rowID(collection, key...), row)
}

// patchRequest is an upsert of a row whose id is already known.
func patchRequest(collection, id string, row rec) batchRequest {
	body := make(map[string]any, len(row)+1)
	for col, v := range row {
		body[col] = jsonValue(v)
	}
	body["id"] = id
	return batchRequest{Method: "PUT", URL: "/api/collections/" + collection + "/records", Body: body}
}

// jsonValue turns a column value into what the record API takes.
func jsonValue(v any) any {
	switch v := v.(type) {
	case []byte:
		return string(v)
	case int:
		return int64(v)
	case int32:
		return int64(v)
	default:
		return v
	}
}

// put upserts a row; see putRequest.
func (t *tx) put(collection string, row rec) {
	req := putRequest(collection, row)
	delete(t.deleted, collection+"/"+req.Body["id"].(string))
	t.ops = append(t.ops, req)
}

// patch upserts columns of the row with this id.
func (t *tx) patch(collection, id string, row rec) {
	delete(t.deleted, collection+"/"+id)
	t.ops = append(t.ops, patchRequest(collection, id, row))
}

// del deletes a row that exists. Deleting one that does not fails the whole
// batch, so callers only delete what they have just read.
func (t *tx) del(collection, id string) bool {
	key := collection + "/" + id
	if t.deleted[key] {
		return false
	}
	t.deleted[key] = true
	t.ops = append(t.ops, batchRequest{Method: "DELETE", URL: "/api/collections/" + collection + "/records/" + id})
	return true
}

// delKey deletes the row with this primary key, if it exists.
func (t *tx) delKey(collection string, key ...any) (bool, error) {
	id := rowID(collection, key...)
	found, err := t.exists(collection, `id = ?`, id)
	if err != nil || !found {
		return false, err
	}
	return t.del(collection, id), nil
}

// ids reads the ids of the rows matching a condition.
func (t *tx) ids(collection, where string, args ...any) ([]string, error) {
	rows, err := t.b.db.QueryContext(t.ctx, `SELECT id FROM `+collection+` WHERE `+where, args...)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()
	var out []string
	for rows.Next() {
		var id string
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		out = append(out, id)
	}
	return out, rows.Err()
}

// delWhere deletes the rows matching a condition and reports how many there
// were.
func (t *tx) delWhere(collection, where string, args ...any) (int, error) {
	ids, err := t.ids(collection, where, args...)
	if err != nil {
		return 0, fmt.Errorf("find the %s rows to delete: %w", collection, err)
	}
	n := 0
	for _, id := range ids {
		if t.del(collection, id) {
			n++
		}
	}
	return n, nil
}

// setWhere updates columns of the rows matching a condition and reports how
// many there were.
func (t *tx) setWhere(collection string, row rec, where string, args ...any) (int, error) {
	ids, err := t.ids(collection, where, args...)
	if err != nil {
		return 0, fmt.Errorf("find the %s rows to update: %w", collection, err)
	}
	n := 0
	for _, id := range ids {
		if t.deleted[collection+"/"+id] {
			continue
		}
		t.patch(collection, id, row)
		n++
	}
	return n, nil
}

// exists reports whether a row matches a condition.
func (t *tx) exists(collection, where string, args ...any) (bool, error) {
	var one int
	err := t.b.db.QueryRowContext(t.ctx, `SELECT 1 FROM `+collection+` WHERE `+where+` LIMIT 1`, args...).Scan(&one)
	switch {
	case err == sql.ErrNoRows:
		return false, nil
	case err != nil:
		return false, err
	}
	return true, nil
}

// queryRow reads one row inside the call.
func (t *tx) queryRow(query string, args ...any) *sql.Row {
	return t.b.db.QueryRowContext(t.ctx, query, args...)
}

// query reads rows inside the call.
func (t *tx) query(query string, args ...any) (*sql.Rows, error) {
	return t.b.db.QueryContext(t.ctx, query, args...)
}

// next hands out the next value of a counter.
func (t *tx) next(table, name string) (int64, error) {
	return t.advance(table, name, 1)
}

// advance moves a counter by n and returns the last value of the block, as
// the SQL drivers' advanceSequence does. n may be negative.
func (t *tx) advance(table, name string, n int64) (int64, error) {
	cur, err := t.current(table, name)
	if err != nil {
		return 0, err
	}
	cur += n
	t.seqs[table+"\x00"+name] = cur
	return cur, nil
}

// current reads a counter as this call has left it.
func (t *tx) current(table, name string) (int64, error) {
	key := table + "\x00" + name
	if v, ok := t.seqs[key]; ok {
		return v, nil
	}
	var v int64
	if err := t.queryRow(`SELECT value FROM `+table+` WHERE name = ?`, name).Scan(&v); err != nil {
		return 0, fmt.Errorf("read the %s counter: %w", name, err)
	}
	t.seqs[key] = v
	t.seqOrder = append(t.seqOrder, key)
	return v, nil
}

// currentSequence reads a counter outside any write.
func (b *base) currentSequence(ctx context.Context, table, name string) (int64, error) {
	var v int64
	if err := b.db.QueryRowContext(ctx, `SELECT value FROM `+table+` WHERE name = ?`, name).Scan(&v); err != nil {
		return 0, fmt.Errorf("read the %s counter: %w", name, err)
	}
	return v, nil
}