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

package pocketbase

import (
	"context"
	"crypto/rand"
	"database/sql/driver"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"
)

// The read side of the driver is a database/sql driver over PocketBase's SQL
// endpoint, so that the stores can read the way the SQL drivers do: with
// `?` placeholders, joins and aggregates, and rows.Scan. It cannot write —
// every write goes through the batch endpoint, as record operations — and
// it cannot begin a transaction.
//
// The endpoint has three limits the driver hides. It takes no parameters, so
// every argument is inlined as a SQL literal. It takes at most 5000
// characters, so a query that would be longer has its longest arguments
// moved into temporary records of warp_params and read back by subquery. And
// it returns at most 1000 rows, so a result that fills a page is read on in
// further pages.

// errReadOnly is returned for anything but a query.
var errReadOnly = errors.New("pocketbase: the SQL connection is read-only; writes go through the batch API")

const (
	// maxQueryLength is the SQL endpoint's limit, with room for the paging
	// wrapper.
	maxQueryLength = 5000 - len(pageWrapperPrefix) - len(pageWrapperSuffix) - 32

	// pageSize is how many rows the SQL endpoint returns at most.
	pageSize = 1000

	pageWrapperPrefix = `SELECT * FROM (`
	pageWrapperSuffix = `) LIMIT 1000 OFFSET `

	// paramsCollection holds the arguments too long to inline.
	paramsCollection = "warp_params"
)

// sqlConnector opens connections to one PocketBase.
type sqlConnector struct {
	c *client

	// lock is held for reading around every query that runs outside a
	// write, so that a query read in several pages is not split by one.
	lock *sync.RWMutex
}

func (k *sqlConnector) Connect(context.Context) (driver.Conn, error) { return &sqlConn{k: k}, nil }
func (k *sqlConnector) Driver() driver.Driver                        { return sqlDriver{k: k} }

// sqlDriver exists because database/sql wants one; connections come from the
// connector.
type sqlDriver struct{ k *sqlConnector }

func (d sqlDriver) Open(string) (driver.Conn, error) { return &sqlConn{k: d.k}, nil }

type sqlConn struct{ k *sqlConnector }

func (c *sqlConn) Prepare(query string) (driver.Stmt, error) {
	return &sqlStmt{c: c, query: query}, nil
}
func (c *sqlConn) Close() error              { return nil }
func (c *sqlConn) Begin() (driver.Tx, error) { return nil, errReadOnly }

// Ping implements driver.Pinger.
func (c *sqlConn) Ping(ctx context.Context) error { return c.k.c.health(ctx) }

// QueryContext implements driver.QueryerContext.
func (c *sqlConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	if !holdsWriteLock(ctx) {
		c.k.lock.RLock()
		defer c.k.lock.RUnlock()
	}
	return c.k.run(ctx, query, args)
}

type sqlStmt struct {
	c     *sqlConn
	query string
}

func (s *sqlStmt) Close() error                               { return nil }
func (s *sqlStmt) NumInput() int                              { return -1 }
func (s *sqlStmt) Exec([]driver.Value) (driver.Result, error) { return nil, errReadOnly }
func (s *sqlStmt) Query(args []driver.Value) (driver.Rows, error) {
	named := make([]driver.NamedValue, len(args))
	for i, v := range args {
		named[i] = driver.NamedValue{Ordinal: i + 1, Value: v}
	}
	return s.c.QueryContext(context.Background(), s.query, named)
}

// writeLockKey marks a context whose caller already holds the store's lock,
// for writing or for a multi-query read, so that its queries do not take
// the read lock again behind it.
type writeLockKey struct{}

func withWriteLock(ctx context.Context) context.Context {
	return context.WithValue(ctx, writeLockKey{}, true)
}

func holdsWriteLock(ctx context.Context) bool {
	held, _ := ctx.Value(writeLockKey{}).(bool)
	return held
}

// run inlines the arguments, moves what does not fit into warp_params, and
// reads every page of the result.
func (k *sqlConnector) run(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	parts, literals, err := splitPlaceholders(query, args)
	if err != nil {
		return nil, err
	}

	params, err := k.fitLiterals(ctx, parts, literals)
	if err != nil {
		return nil, err
	}
	if len(params) > 0 {
		defer k.dropParams(params)
	}
	full := joinLiterals(parts, literals)

	res, err := k.c.query(ctx, full)
	if err != nil {
		return nil, fmt.Errorf("%w\n\tin query: %s", err, abbreviate(full))
	}
	rows := &sqlRows{columns: make([]string, len(res.Columns)), data: res.Rows}
	for i, col := range res.Columns {
		rows.columns[i] = col.Name
	}
	for page := 1; len(res.Rows) == pageSize; page++ {
		res, err = k.c.query(ctx, pageWrapperPrefix+full+pageWrapperSuffix+strconv.Itoa(page*pageSize))
		if err != nil {
			return nil, fmt.Errorf("%w\n\tin page %d of query: %s", err, page, abbreviate(full))
		}
		rows.data = append(rows.data, res.Rows...)
	}
	return rows, nil
}

// splitPlaceholders cuts a query at its `?` placeholders, skipping those
// inside quoted strings and identifiers, and renders each argument as a SQL
// literal.
func splitPlaceholders(query string, args []driver.NamedValue) ([]string, []string, error) {
	var (
		parts    []string
		literals []string
		start    int
		quote    byte
	)
	for i := 0; i < len(query); i++ {
		ch := query[i]
		switch {
		case quote != 0:
			if ch == quote {
				quote = 0
			}
		case ch == '\'' || ch == '"':
			quote = ch
		case ch == '?':
			n := len(literals)
			if n >= len(args) {
				return nil, nil, fmt.Errorf("pocketbase: query has more placeholders than its %d arguments", len(args))
			}
			lit, err := literal(args[n].Value)
			if err != nil {
				return nil, nil, err
			}
			parts = append(parts, query[start:i])
			literals = append(literals, lit)
			start = i + 1
		}
	}
	if len(literals) != len(args) {
		return nil, nil, fmt.Errorf("pocketbase: query has %d placeholders for %d arguments", len(literals), len(args))
	}
	return append(parts, query[start:]), literals, nil
}

// literal renders one value as SQLite would read it back.
func literal(v any) (string, error) {
	switch v := v.(type) {
	case nil:
		return "NULL", nil
	case bool:
		if v {
			return "1", nil
		}
		return "0", nil
	case int64:
		return strconv.FormatInt(v, 10), nil
	case float64:
		return strconv.FormatFloat(v, 'g', -1, 64), nil
	case string:
		return textLiteral(v), nil
	case []byte:
		return textLiteral(string(v)), nil
	case time.Time:
		return textLiteral(v.UTC().Format(time.RFC3339Nano)), nil
	default:
		return "", fmt.Errorf("pocketbase: cannot inline a %T into SQL", v)
	}
}

// textLiteral quotes a string. Text that is not valid UTF-8, or holds a NUL,
// goes in as hex and is cast back, so that nothing in it can end the literal
// early.
func textLiteral(s string) string {
	if utf8.ValidString(s) && !strings.ContainsRune(s, 0) {
		return "'" + strings.ReplaceAll(s, "'", "''") + "'"
	}
	return "CAST(X'" + hex.EncodeToString([]byte(s)) + "' AS TEXT)"
}

func joinLiterals(parts, literals []string) string {
	var b strings.Builder
	for i, p := range parts {
		b.WriteString(p)
		if i < len(literals) {
			b.WriteString(literals[i])
		}
	}
	return b.String()
}

// queryLength is what the SQL endpoint counts: characters, not bytes.
func queryLength(parts, literals []string) int {
	n := 0
	for _, p := range parts {
		n += utf8.RuneCountInString(p)
	}
	for _, l := range literals {
		n += utf8.RuneCountInString(l)
	}
	return n
}

// fitLiterals moves the longest text literals into warp_params until the
// query fits, rewriting each into a subquery reading it back. It returns the
// IDs of the records it created.
func (k *sqlConnector) fitLiterals(ctx context.Context, parts, literals []string) ([]string, error) {
	if queryLength(parts, literals) <= maxQueryLength {
		return nil, nil
	}
	order := make([]int, len(literals))
	for i := range order {
		order[i] = i
	}
	slices.SortFunc(order, func(a, b int) int { return len(literals[b]) - len(literals[a]) })

	var (
		ids  []string
		reqs []batchRequest
	)
	for _, i := range order {
		if queryLength(parts, literals) <= maxQueryLength {
			break
		}
		text, ok := unquote(literals[i])
		if !ok {
			continue
		}
		id, err := randomID()
		if err != nil {
			return nil, err
		}
		reqs = append(reqs, batchRequest{
			Method: "POST",
			URL:    "/api/collections/" + paramsCollection + "/records",
			Body:   map[string]any{"id": id, "value": text},
		})
		ids = append(ids, id)
		literals[i] = `(SELECT value FROM ` + paramsCollection + ` WHERE id = '` + id + `')`
	}
	if queryLength(parts, literals) > maxQueryLength {
		return nil, fmt.Errorf("pocketbase: query is too long for the SQL endpoint even with its arguments moved out")
	}
	if err := k.c.batch(ctx, reqs); err != nil {
		return nil, fmt.Errorf("pocketbase: store query arguments: %w", err)
	}
	return ids, nil
}

// dropParams deletes the records a query's arguments were moved into. It
// runs after the query whatever happened to it, on a context of its own, and
// a failure only leaves rows behind for the next start to clear.
func (k *sqlConnector) dropParams(ids []string) {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	reqs := make([]batchRequest, len(ids))
	for i, id := range ids {
		reqs[i] = batchRequest{Method: "DELETE", URL: "/api/collections/" + paramsCollection + "/records/" + id}
	}
	_ = k.c.batch(ctx, reqs)
}

// unquote recovers the text of a literal textLiteral made.
func unquote(lit string) (string, bool) {
	if strings.HasPrefix(lit, "'") {
		return strings.ReplaceAll(lit[1:len(lit)-1], "''", "'"), true
	}
	if hexText, ok := strings.CutPrefix(lit, "CAST(X'"); ok {
		raw, err := hex.DecodeString(strings.TrimSuffix(hexText, "' AS TEXT)"))
		return string(raw), err == nil
	}
	return "", false
}

// randomID makes a record ID in PocketBase's default alphabet.
func randomID() (string, error) {
	var raw [8]byte
	if _, err := rand.Read(raw[:]); err != nil {
		return "", err
	}
	return "p" + hex.EncodeToString(raw[:])[:14], nil
}

func abbreviate(query string) string {
	query = strings.Join(strings.Fields(query), " ")
	if len(query) > 500 {
		return query[:500] + "…"
	}
	return query
}

// sqlRows is a whole result, already read.
type sqlRows struct {
	columns []string
	data    [][]*string
	next    int
}

func (r *sqlRows) Columns() []string { return r.columns }
func (r *sqlRows) Close() error      { return nil }

func (r *sqlRows) Next(dest []driver.Value) error {
	if r.next >= len(r.data) {
		return io.EOF
	}
	row := r.data[r.next]
	r.next++
	for i := range dest {
		if i < len(row) && row[i] != nil {
			dest[i] = *row[i]
		} else {
			dest[i] = nil
		}
	}
	return nil
}