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

package sqlstore

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 live apart from the rest of the room store because they are the only
// queries in it that are about *order* rather than about identity. Every
// other read answers "what is the state of this room" or "what is this
// event"; these answer "what happened between here and there", which is a
// different access pattern and the one the (room_id, stream_ordering) index
// exists for.

const selectMaxStreamOrderingSQL = `
	SELECT COALESCE(MAX(stream_ordering), 0) FROM roomserver_events`

// MaxStreamOrdering implements storage.RoomStore.
//
// It reads the events table rather than the counter in
// roomserver_sequences. The counter is advanced inside the writing
// transaction and is therefore already at the new value while that
// transaction is still open; a sync token taken from it would name a
// position whose event no reader can see yet, and the event would be
// skipped for good once the transaction committed.
func (s *Rooms) MaxStreamOrdering(ctx context.Context) (int64, error) {
	var top int64
	if err := s.db.QueryRowContext(ctx, s.q(selectMaxStreamOrderingSQL)).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 (` + placeholders(len(memberships)) + `)`
		for _, m := range memberships {
			args = append(args, m)
		}
	}
	query += ` ORDER BY cs.room_id`

	rows, err := s.db.QueryContext(ctx, s.q(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()
}

var selectUserMembershipEventsSQL = `
	SELECT ` + eventColumns + `
	FROM roomserver_events e
	WHERE e.room_id = ? AND e.event_type = 'm.room.member' AND e.state_key = ? AND NOT e.outlier
	ORDER BY e.stream_ordering`

// UserMembershipEvents implements storage.RoomStore.
func (s *Rooms) UserMembershipEvents(ctx context.Context, roomID, userID string) ([]storage.StoredEvent, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectUserMembershipEventsSQL), 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
	}

	// Outliers have no place in a timeline; see the outlier column.
	conditions := []string{`e.room_id = ?`, `NOT e.outlier`}
	args := []any{req.RoomID}
	// Zero is "unbounded", not a position, at either end: a bound can be
	// negative, inside backfilled history, so "is it set" cannot be
	// spelled "is it positive".
	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)
	}

	// The order decides which end the LIMIT trims: reading backwards keeps
	// the newest of the window, which is what a timeline is.
	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)

	query := `SELECT ` + eventColumns + `
		FROM roomserver_events e
		WHERE ` + strings.Join(conditions, " AND ") + `
		ORDER BY e.stream_ordering ` + order + `
		LIMIT ?`

	rows, err := s.db.QueryContext(ctx, s.q(query), 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.
//
// It reads by stream_ordering alone, which the column's own uniqueness
// index serves. Backfilled history sits below zero and a reader starting at
// a real position never reaches it, which is right for what this is for:
// an application service is told what happens, not what a backfill dug up.
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)

	query := `SELECT ` + eventColumns + `
		FROM roomserver_events e
		WHERE ` + strings.Join(conditions, " AND ") + `
		ORDER BY e.stream_ordering
		LIMIT ?`
	rows, err := s.db.QueryContext(ctx, s.q(query), 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, as a WHERE fragment and its arguments.
//
// The exclusions are ANDed in after the inclusions, which is what makes
// not_types win over types for an event matching both — the behaviour the
// spec requires.
func filterConditions(f storage.TimelineFilter) (string, []any) {
	return filterConditionsOn(`e.event_type`, `e.sender`, f)
}

// filterConditionsOn is filterConditions against named columns, for a
// table other than roomserver_events.
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".
//
// A filter pattern may end in '*', which the spec defines as matching any
// suffix — "m.room.*" is how a client asks for every room event type. Exact
// patterns go into one IN list so that the common case is a single indexed
// lookup; a wildcard becomes a LIKE with an explicit escape character, so
// that a '%' or '_' inside the pattern is matched literally rather than
// being read as a LIKE wildcard of its own.
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 (`+placeholders(len(exact))+`)`)
		for _, p := range exact {
			args = append(args, p)
		}
	}
	for _, prefix := range wildcard {
		parts = append(parts, column+` LIKE ? ESCAPE '\'`)
		args = append(args, escapeLike(prefix)+`%`)
	}
	return strings.Join(parts, " OR "), args
}

// escapeLike makes a literal string safe to embed in a LIKE pattern that
// declares '\' as its escape character.
func escapeLike(s string) string {
	var b strings.Builder
	b.Grow(len(s))
	for _, r := range s {
		if r == '\\' || r == '%' || r == '_' {
			b.WriteRune('\\')
		}
		b.WriteRune(r)
	}
	return b.String()
}

// 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 (` + placeholders(len(roomIDs)) + `) AND NOT outlier`
	args := make([]any, 0, len(roomIDs)+2)
	for _, id := range roomIDs {
		args = append(args, id)
	}
	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`

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("find the rooms with events after %d: %w", from, err)
	}
	defer func() { _ = rows.Close() }()

	var out []string
	for rows.Next() {
		var roomID string
		if err := rows.Scan(&roomID); err != nil {
			return nil, fmt.Errorf("find the rooms with events after %d: %w", from, err)
		}
		out = append(out, roomID)
	}
	return out, rows.Err()
}

// 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 (` + placeholders(len(roomIDs)) + `)
		AND event_type IN (` + placeholders(len(types)) + `) AND NOT outlier`
	args := make([]any, 0, len(roomIDs)+len(types)+1)
	for _, id := range roomIDs {
		args = append(args, id)
	}
	for _, t := range types {
		args = append(args, t)
	}
	if upTo != 0 {
		query += ` AND stream_ordering <= ?`
		args = append(args, upTo)
	}
	query += ` GROUP BY room_id`

	rows, err := s.db.QueryContext(ctx, s.q(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()
}

const selectEventAtStreamSQL = `
	SELECT e.event_id, e.event_type, e.state_key, 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`

const selectSnapshotEntriesSQL = `
	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`

// StateAtStream implements storage.RoomStore.
//
// The state after a position is the snapshot recorded against the last
// event at or before it, plus that event itself when it is a state event:
// the snapshot is the state the event was authorised against, which is the
// state *before* it.
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, s.q(selectEventAtStreamSQL), roomID, pos).
		Scan(&eventID, &eventType, &stateKey, &snapshotID)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		// Nothing had happened in the room by then, which for a position
		// before its create event is the honest answer.
		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, s.q(selectSnapshotEntriesSQL), 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)
}