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

package pocketbase

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"

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

// The stream reads of the room store: the queries the sync API and
// /messages are built on. They are the SQL store's.

// MaxStreamOrdering implements storage.RoomStore. It reads the events, not
// the counter: a batch writes both at once, but an event skipped as a
// duplicate leaves the counter above every event there is.
func (s *Rooms) MaxStreamOrdering(ctx context.Context) (int64, error) {
	var top int64
	if err := s.db.QueryRowContext(ctx,
		`SELECT COALESCE(MAX(stream_ordering), 0) FROM roomserver_events`).Scan(&top); err != nil {
		return 0, fmt.Errorf("read the top of the event stream: %w", err)
	}
	return top, nil
}

// MembershipsForUser implements storage.RoomStore.
func (s *Rooms) MembershipsForUser(ctx context.Context, userID string, memberships []string) ([]storage.SyncMembership, error) {
	query := `
		SELECT cs.room_id, cs.membership, cs.event_id, e.stream_ordering
		FROM roomserver_current_state cs
		JOIN roomserver_events e ON e.event_id = cs.event_id
		WHERE cs.event_type = 'm.room.member' AND cs.state_key = ?`
	args := []any{userID}
	if len(memberships) > 0 {
		query += ` AND cs.membership IN ` + inList
		args = append(args, jsonList(memberships))
	}
	query += ` ORDER BY cs.room_id`

	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("read the memberships of %q: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.SyncMembership
	for rows.Next() {
		var m storage.SyncMembership
		if err := rows.Scan(&m.RoomID, &m.Membership, &m.EventID, &m.StreamOrdering); err != nil {
			return nil, fmt.Errorf("read the memberships of %q: %w", userID, err)
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

// UserMembershipEvents implements storage.RoomStore.
func (s *Rooms) UserMembershipEvents(ctx context.Context, roomID, userID string) ([]storage.StoredEvent, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT `+eventColumns+`
		FROM roomserver_events e
		WHERE e.room_id = ? AND e.event_type = 'm.room.member' AND e.has_state_key AND e.state_key = ?
			AND NOT e.outlier
		ORDER BY e.stream_ordering`, roomID, userID)
	if err != nil {
		return nil, fmt.Errorf("read the membership events of %q in %q: %w", userID, roomID, err)
	}
	events, err := scanEvents(rows)
	if err != nil {
		return nil, fmt.Errorf("read the membership events of %q in %q: %w", userID, roomID, err)
	}
	return events, nil
}

// Timeline implements storage.RoomStore.
func (s *Rooms) Timeline(ctx context.Context, req storage.TimelineRequest) ([]storage.StoredEvent, error) {
	if req.Limit <= 0 {
		return nil, nil
	}

	conditions := []string{`e.room_id = ?`, `NOT e.outlier`}
	args := []any{req.RoomID}
	if req.Low != 0 {
		conditions = append(conditions, `e.stream_ordering > ?`)
		args = append(args, req.Low)
	}
	if req.High != 0 {
		conditions = append(conditions, `e.stream_ordering <= ?`)
		args = append(args, req.High)
	}

	order := `ASC`
	if req.Direction == storage.DirectionBackward {
		order = `DESC`
	}

	if clause, filterArgs := filterConditions(req.Filter); clause != "" {
		conditions = append(conditions, clause)
		args = append(args, filterArgs...)
	}
	args = append(args, req.Limit)

	rows, err := s.db.QueryContext(ctx, `SELECT `+eventColumns+`
		FROM roomserver_events e
		WHERE `+strings.Join(conditions, " AND ")+`
		ORDER BY e.stream_ordering `+order+`
		LIMIT ?`, args...)
	if err != nil {
		return nil, fmt.Errorf("read the timeline of %q: %w", req.RoomID, err)
	}
	events, err := scanEvents(rows)
	if err != nil {
		return nil, fmt.Errorf("read the timeline of %q: %w", req.RoomID, err)
	}
	return events, nil
}

// StreamEvents implements storage.RoomStore.
func (s *Rooms) StreamEvents(ctx context.Context, after, upTo int64, limit int) ([]storage.StoredEvent, error) {
	if limit <= 0 {
		return nil, nil
	}
	conditions := []string{`e.stream_ordering > ?`, `NOT e.outlier`}
	args := []any{after}
	if upTo != 0 {
		conditions = append(conditions, `e.stream_ordering <= ?`)
		args = append(args, upTo)
	}
	args = append(args, limit)

	rows, err := s.db.QueryContext(ctx, `SELECT `+eventColumns+`
		FROM roomserver_events e
		WHERE `+strings.Join(conditions, " AND ")+`
		ORDER BY e.stream_ordering
		LIMIT ?`, args...)
	if err != nil {
		return nil, fmt.Errorf("read the event stream after %d: %w", after, err)
	}
	events, err := scanEvents(rows)
	if err != nil {
		return nil, fmt.Errorf("read the event stream after %d: %w", after, err)
	}
	return events, nil
}

// filterConditions renders the part of an event filter the store applies in
// SQL. The exclusions are ANDed in after the inclusions, which is what makes
// not_types win over types.
func filterConditions(f storage.TimelineFilter) (string, []any) {
	return filterConditionsOn(`e.event_type`, `e.sender`, f)
}

// filterConditionsOn is filterConditions against named columns.
func filterConditionsOn(typeColumn, senderColumn string, f storage.TimelineFilter) (string, []any) {
	var (
		clauses []string
		args    []any
	)
	add := func(column string, patterns []string, negate bool) {
		if len(patterns) == 0 {
			return
		}
		clause, clauseArgs := patternMatch(column, patterns)
		if negate {
			clause = `NOT (` + clause + `)`
		}
		clauses = append(clauses, clause)
		args = append(args, clauseArgs...)
	}
	add(typeColumn, f.Types, false)
	add(typeColumn, f.NotTypes, true)
	add(senderColumn, f.Senders, false)
	add(senderColumn, f.NotSenders, true)

	if len(clauses) == 0 {
		return "", nil
	}
	return strings.Join(clauses, " AND "), args
}

// patternMatch renders "this column matches any of these patterns", where a
// trailing '*' matches any suffix.
func patternMatch(column string, patterns []string) (string, []any) {
	var (
		exact    []string
		wildcard []string
		args     []any
	)
	for _, p := range patterns {
		if prefix, isWildcard := strings.CutSuffix(p, "*"); isWildcard {
			wildcard = append(wildcard, prefix)
			continue
		}
		exact = append(exact, p)
	}

	var parts []string
	if len(exact) > 0 {
		parts = append(parts, column+` IN `+inList)
		args = append(args, jsonList(exact))
	}
	for _, prefix := range wildcard {
		parts = append(parts, column+` LIKE ? ESCAPE '\'`)
		args = append(args, escapeLike(prefix)+`%`)
	}
	return strings.Join(parts, " OR "), args
}

// RoomsWithEvents implements storage.RoomStore.
func (s *Rooms) RoomsWithEvents(ctx context.Context, roomIDs []string, from, to int64) ([]string, error) {
	if len(roomIDs) == 0 {
		return nil, nil
	}

	query := `
		SELECT DISTINCT room_id FROM roomserver_events
		WHERE room_id IN ` + inList + ` AND NOT outlier`
	args := []any{jsonList(roomIDs)}
	if from != 0 {
		query += ` AND stream_ordering > ?`
		args = append(args, from)
	}
	if to != 0 {
		query += ` AND stream_ordering <= ?`
		args = append(args, to)
	}
	query += ` ORDER BY room_id`

	out, err := s.queryStrings(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("find the rooms with events after %d: %w", from, err)
	}
	return out, nil
}

// LatestEventPositions implements storage.RoomStore.
func (s *Rooms) LatestEventPositions(ctx context.Context, roomIDs, types []string, upTo int64) (map[string]int64, error) {
	if len(roomIDs) == 0 || len(types) == 0 {
		return map[string]int64{}, nil
	}
	query := `
		SELECT room_id, MAX(stream_ordering) FROM roomserver_events
		WHERE room_id IN ` + inList + `
		AND event_type IN ` + inList + ` AND NOT outlier`
	args := []any{jsonList(roomIDs), jsonList(types)}
	if upTo != 0 {
		query += ` AND stream_ordering <= ?`
		args = append(args, upTo)
	}
	query += ` GROUP BY room_id`

	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("find the latest activity of %d rooms: %w", len(roomIDs), err)
	}
	defer func() { _ = rows.Close() }()

	out := make(map[string]int64, len(roomIDs))
	for rows.Next() {
		var (
			roomID string
			pos    int64
		)
		if err := rows.Scan(&roomID, &pos); err != nil {
			return nil, fmt.Errorf("find the latest activity of %d rooms: %w", len(roomIDs), err)
		}
		out[roomID] = pos
	}
	return out, rows.Err()
}

// StateAtStream implements storage.RoomStore: the snapshot recorded against
// the last event at or before the position, plus that event itself when it
// is a state event.
func (s *Rooms) StateAtStream(ctx context.Context, roomID string, pos int64) ([]storage.StateEntry, error) {
	if pos == 0 {
		return nil, nil
	}

	var (
		eventID    string
		eventType  string
		stateKey   sql.NullString
		snapshotID int64
	)
	err := s.db.QueryRowContext(ctx, `
		SELECT e.event_id, e.event_type, CASE WHEN e.has_state_key THEN e.state_key END, e.state_before_id
		FROM roomserver_events e
		WHERE e.room_id = ? AND NOT e.outlier AND e.stream_ordering <= ?
		ORDER BY e.stream_ordering DESC
		LIMIT 1`, roomID, pos).
		Scan(&eventID, &eventType, &stateKey, &snapshotID)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return nil, nil
	case err != nil:
		return nil, fmt.Errorf("read the state of %q at %d: %w", roomID, pos, err)
	}

	var entries []storage.StateEntry
	if snapshotID != 0 {
		rows, qErr := s.db.QueryContext(ctx, `
			SELECT se.event_type, se.state_key, se.event_id, ''
			FROM roomserver_state_entries se
			WHERE se.snapshot_id = ?
			ORDER BY se.event_type, se.state_key`, snapshotID)
		if qErr != nil {
			return nil, fmt.Errorf("read the state of %q at %d: %w", roomID, pos, qErr)
		}
		if entries, err = scanStateEntries(rows); err != nil {
			return nil, fmt.Errorf("read the state of %q at %d: %w", roomID, pos, err)
		}
	}
	if stateKey.Valid {
		entries = upsertStateEntry(entries, storage.StateEntry{
			Type:     eventType,
			StateKey: stateKey.String,
			EventID:  eventID,
		})
	}
	return entries, nil
}

// upsertStateEntry replaces the matching tuple in entries, or appends it.
func upsertStateEntry(entries []storage.StateEntry, entry storage.StateEntry) []storage.StateEntry {
	for i, existing := range entries {
		if existing.Type == entry.Type && existing.StateKey == entry.StateKey {
			entries[i] = entry
			return entries
		}
	}
	return append(entries, entry)
}