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

package sqlstore

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

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

// The relation and search half of the room store: the reads behind
// /relations, /threads, the bundled summaries in unsigned.m.relations, and
// /search.
//
// Relations need nothing but the two columns on roomserver_events and their
// index, and are written once for both engines like everything else here.
// Search is the one place the engines really differ — SQLite has FTS5,
// PostgreSQL has tsvector — so the statements that touch the index come
// from the driver as a SearchDialect, and only the shape of the query is
// shared.

// SearchDialect is one engine's full-text index.
//
// Every statement is written against the table roomserver_event_search,
// referred to as Name where it is read, with `?` placeholders that are
// rebound like any other statement.
type SearchDialect struct {
	// Insert adds one event, taking (stream_ordering, room_id, event_id,
	// search_key, sender, event_type, body) in that order.
	Insert string

	// DeleteEvent removes an event's row, taking (room_id, event_id).
	DeleteEvent string

	// DeleteRoom removes a room's rows, taking (room_id).
	DeleteRoom string

	// From is the FROM clause the search reads, and Name what its columns
	// are qualified with. SQLite's FTS5 cannot be aliased — MATCH and
	// bm25() take the table's own name — so the two are separate.
	From string
	Name string

	// Position is the column holding a row's stream ordering.
	Position string

	// Match is the condition selecting the rows that match. It takes the
	// engine query once.
	Match string

	// Rank is the relevance of a matching row, higher being better. It
	// takes the engine query RankArgs times.
	Rank     string
	RankArgs int

	// Query turns what a user typed into the engine's query syntax. It
	// returns "" for a term with nothing searchable in it, which matches
	// nothing.
	Query func(term string) string
}

// WithSearch sets the full-text index the store writes and reads.
func (s *Rooms) WithSearch(d *SearchDialect) *Rooms {
	s.search = d
	return s
}

// indexEvent adds an event that has just joined a timeline to the search
// index, if it has anything to find it by.
func (s *Rooms) indexEvent(ctx context.Context, tx *sql.Tx, e storage.StoredEvent, ordering int64) error {
	if s.search == nil || e.SearchKey == "" || strings.TrimSpace(e.SearchText) == "" {
		return nil
	}
	_, err := tx.ExecContext(ctx, s.q(s.search.Insert),
		ordering, e.RoomID, e.EventID, e.SearchKey, e.Sender, e.Type, e.SearchText)
	if err != nil {
		return fmt.Errorf("index event %s for search: %w", e.EventID, err)
	}
	return nil
}

// unindexEvent removes an event from the search index.
func (s *Rooms) unindexEvent(ctx context.Context, tx *sql.Tx, roomID, eventID string) error {
	if s.search == nil {
		return nil
	}
	if _, err := tx.ExecContext(ctx, s.q(s.search.DeleteEvent), roomID, eventID); err != nil {
		return fmt.Errorf("remove event %s from the search index: %w", eventID, err)
	}
	return nil
}

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

	conditions := []string{
		`e.room_id = ?`,
		`e.relates_to_id IN (` + placeholders(len(req.ParentIDs)) + `)`,
		`NOT e.outlier`,
	}
	args := []any{req.RoomID}
	for _, id := range req.ParentIDs {
		args = append(args, id)
	}
	if req.RelType != "" {
		conditions = append(conditions, `e.rel_type = ?`)
		args = append(args, req.RelType)
	} else {
		// A parent's own row has relates_to_id '' and can never match an
		// IN list of event IDs, but '' itself could be passed; it names no
		// event and has no children.
		conditions = append(conditions, `e.rel_type <> ''`)
	}
	if req.EventType != "" {
		conditions = append(conditions, `e.event_type = ?`)
		args = append(args, req.EventType)
	}
	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`
	}
	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 relations of %v in %q: %w", req.ParentIDs, req.RoomID, err)
	}
	events, err := scanEvents(rows)
	if err != nil {
		return nil, fmt.Errorf("read the relations of %v in %q: %w", req.ParentIDs, req.RoomID, err)
	}
	return events, nil
}

// ChildEventIDs implements storage.RoomStore.
func (s *Rooms) ChildEventIDs(ctx context.Context, roomID string, parentIDs []string, relType, eventType string) ([]string, error) {
	if len(parentIDs) == 0 {
		return nil, nil
	}
	query := `SELECT event_id FROM roomserver_events
		WHERE room_id = ? AND relates_to_id IN (` + placeholders(len(parentIDs)) + `)
		AND rel_type <> '' AND NOT outlier`
	args := []any{roomID}
	for _, id := range parentIDs {
		args = append(args, id)
	}
	if relType != "" {
		query += ` AND rel_type = ?`
		args = append(args, relType)
	}
	if eventType != "" {
		query += ` AND event_type = ?`
		args = append(args, eventType)
	}
	query += ` ORDER BY stream_ordering`

	ids, err := s.queryStrings(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("read the children of %v in %q: %w", parentIDs, roomID, err)
	}
	return ids, nil
}

// RelationAggregates implements storage.RoomStore.
//
// It groups by parent, relation type and sender, and joins each group's
// newest position back onto the events table for its event ID: a portable
// spelling of "the row with the largest stream ordering in each group",
// which neither engine can say in one aggregate.
func (s *Rooms) RelationAggregates(ctx context.Context, roomID string, parentIDs []string, high int64) ([]storage.RelationAggregate, error) {
	if len(parentIDs) == 0 {
		return nil, nil
	}
	inner := `SELECT relates_to_id, rel_type, sender, COUNT(*) AS n, MAX(stream_ordering) AS latest
		FROM roomserver_events
		WHERE room_id = ? AND relates_to_id IN (` + placeholders(len(parentIDs)) + `)
		AND rel_type IN (?, ?, ?) AND NOT outlier`
	args := []any{roomID}
	for _, id := range parentIDs {
		args = append(args, id)
	}
	args = append(args, storage.RelThread, storage.RelReplace, storage.RelReference)
	if high != 0 {
		inner += ` AND stream_ordering <= ?`
		args = append(args, high)
	}
	inner += ` GROUP BY relates_to_id, rel_type, sender`

	query := `SELECT g.relates_to_id, g.rel_type, g.sender, g.n, g.latest, e.event_id
		FROM (` + inner + `) g
		JOIN roomserver_events e ON e.stream_ordering = g.latest`
	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("summarise the relations of %d events in %q: %w", len(parentIDs), roomID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.RelationAggregate
	for rows.Next() {
		var a storage.RelationAggregate
		if err := rows.Scan(&a.ParentID, &a.RelType, &a.Sender, &a.Count, &a.LatestOrdering, &a.LatestEventID); err != nil {
			return nil, fmt.Errorf("summarise the relations in %q: %w", roomID, err)
		}
		out = append(out, a)
	}
	return out, rows.Err()
}

// RelationEventIDs implements storage.RoomStore.
//
// The limit is per parent in spirit and applied over the whole set: the
// summary it feeds is a sample, and a parent with more children than the
// limit is rare enough that one query beats one per parent.
func (s *Rooms) RelationEventIDs(ctx context.Context, roomID string, parentIDs []string, relType string, high int64, limit int) (map[string][]string, error) {
	if len(parentIDs) == 0 || limit <= 0 {
		return nil, nil
	}
	query := `SELECT relates_to_id, event_id FROM roomserver_events
		WHERE room_id = ? AND relates_to_id IN (` + placeholders(len(parentIDs)) + `)
		AND rel_type = ? AND NOT outlier`
	args := []any{roomID}
	for _, id := range parentIDs {
		args = append(args, id)
	}
	args = append(args, relType)
	if high != 0 {
		query += ` AND stream_ordering <= ?`
		args = append(args, high)
	}
	query += ` ORDER BY stream_ordering LIMIT ?`
	args = append(args, limit*len(parentIDs))

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

	out := map[string][]string{}
	for rows.Next() {
		var parent, id string
		if err := rows.Scan(&parent, &id); err != nil {
			return nil, fmt.Errorf("read the %s relations in %q: %w", relType, roomID, err)
		}
		if len(out[parent]) < limit {
			out[parent] = append(out[parent], id)
		}
	}
	return out, rows.Err()
}

// Threads implements storage.RoomStore.
//
// A thread is a root with at least one m.thread child, and is ordered by its
// newest child. The root has to be an event of the room this server holds
// in its timeline: a thread whose root is missing — not yet backfilled, or
// never sent here — cannot be shown as a thread at all.
func (s *Rooms) Threads(ctx context.Context, req storage.ThreadsRequest) ([]storage.Thread, error) {
	if req.Limit <= 0 {
		return nil, nil
	}
	inner := `SELECT relates_to_id AS root, MAX(stream_ordering) AS latest
		FROM roomserver_events
		WHERE room_id = ? AND rel_type = ? AND NOT outlier`
	args := []any{req.RoomID, storage.RelThread}
	if req.High != 0 {
		inner += ` AND stream_ordering <= ?`
		args = append(args, req.High)
	}
	inner += ` GROUP BY relates_to_id`

	conditions := []string{`r.room_id = ?`, `NOT r.outlier`}
	args = append(args, req.RoomID)
	if req.Before != 0 {
		conditions = append(conditions, `g.latest < ?`)
		args = append(args, req.Before)
	}
	if req.ParticipatedBy != "" {
		conditions = append(conditions, `(r.sender = ? OR EXISTS (
			SELECT 1 FROM roomserver_events p
			WHERE p.room_id = r.room_id AND p.relates_to_id = g.root AND p.rel_type = ? AND p.sender = ?))`)
		args = append(args, req.ParticipatedBy, storage.RelThread, req.ParticipatedBy)
	}
	args = append(args, req.Limit)

	query := `SELECT g.root, g.latest
		FROM (` + inner + `) g
		JOIN roomserver_events r ON r.event_id = g.root
		WHERE ` + strings.Join(conditions, " AND ") + `
		ORDER BY g.latest DESC
		LIMIT ?`
	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read the threads of %q: %w", req.RoomID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.Thread
	for rows.Next() {
		var t storage.Thread
		if err := rows.Scan(&t.RootID, &t.LatestOrdering); err != nil {
			return nil, fmt.Errorf("read the threads of %q: %w", req.RoomID, err)
		}
		out = append(out, t)
	}
	return out, rows.Err()
}

// searchKeys are the fields the index holds.
var searchKeys = []string{"content.body", "content.name", "content.topic"}

// Search implements storage.RoomStore.
func (s *Rooms) Search(ctx context.Context, req storage.SearchRequest) ([]storage.SearchResult, int, error) {
	if s.search == nil || req.Limit <= 0 || len(req.Rooms) == 0 {
		return nil, 0, nil
	}
	term := s.search.Query(req.Term)
	if term == "" {
		return nil, 0, nil
	}

	pos := s.search.Position
	col := func(name string) string { return s.search.Name + "." + name }
	conditions := []string{s.search.Match}
	args := []any{term}

	// One bounded clause per room: what a user may search of a room is a
	// window of its stream, and the windows differ from room to room.
	rooms := make([]string, 0, len(req.Rooms))
	for _, r := range req.Rooms {
		clause := col(`room_id`) + ` = ?`
		args = append(args, r.RoomID)
		if r.Low != 0 {
			clause += ` AND ` + pos + ` > ?`
			args = append(args, r.Low)
		}
		if r.High != 0 {
			clause += ` AND ` + pos + ` <= ?`
			args = append(args, r.High)
		}
		rooms = append(rooms, `(`+clause+`)`)
	}
	conditions = append(conditions, `(`+strings.Join(rooms, ` OR `)+`)`)

	keys := req.Keys
	if len(keys) == 0 {
		keys = searchKeys
	}
	conditions = append(conditions, col(`search_key`)+` IN (`+placeholders(len(keys))+`)`)
	for _, k := range keys {
		args = append(args, k)
	}

	if clause, filterArgs := filterConditionsOn(col(`event_type`), col(`sender`), req.Filter); clause != "" {
		conditions = append(conditions, clause)
		args = append(args, filterArgs...)
	}
	where := strings.Join(conditions, " AND ")

	var total int
	countQuery := `SELECT COUNT(*) FROM ` + s.search.From + ` WHERE ` + where
	if err := s.db.QueryRowContext(ctx, s.q(countQuery), args...).Scan(&total); err != nil {
		return nil, 0, fmt.Errorf("count search results: %w", err)
	}

	pageWhere := where
	pageArgs := append([]any{}, args...)
	if req.Order == storage.SearchByRecent && req.Before != 0 {
		pageWhere += ` AND ` + pos + ` < ?`
		pageArgs = append(pageArgs, req.Before)
	}

	rankArgs := make([]any, 0, s.search.RankArgs)
	for range s.search.RankArgs {
		rankArgs = append(rankArgs, term)
	}
	order := `rank DESC, ` + pos + ` DESC`
	if req.Order == storage.SearchByRecent {
		order = pos + ` DESC`
	}
	query := `SELECT ` + pos + `, ` + s.search.Rank + ` AS rank
		FROM ` + s.search.From + `
		WHERE ` + pageWhere + `
		ORDER BY ` + order + `
		LIMIT ? OFFSET ?`
	offset := req.Offset
	if req.Order == storage.SearchByRecent {
		offset = 0
	}
	queryArgs := append(rankArgs, pageArgs...)
	queryArgs = append(queryArgs, req.Limit, offset)

	hits, err := s.searchHits(ctx, query, queryArgs)
	if err != nil {
		return nil, 0, err
	}
	if len(hits) == 0 {
		return nil, total, nil
	}

	// The events themselves, read by position in one query and put back in
	// the order the search ranked them.
	eventArgs := make([]any, 0, len(hits))
	for _, h := range hits {
		eventArgs = append(eventArgs, h.pos)
	}
	eventRows, err := s.db.QueryContext(ctx, s.q(`SELECT `+eventColumns+`
		FROM roomserver_events e WHERE e.stream_ordering IN (`+placeholders(len(hits))+`)`), eventArgs...)
	if err != nil {
		return nil, 0, fmt.Errorf("read search results: %w", err)
	}
	events, err := scanEvents(eventRows)
	if err != nil {
		return nil, 0, fmt.Errorf("read search results: %w", err)
	}
	byPos := make(map[int64]storage.StoredEvent, len(events))
	for _, e := range events {
		byPos[e.StreamOrdering] = e
	}

	out := make([]storage.SearchResult, 0, len(hits))
	for _, h := range hits {
		// An index row whose event has gone — purged under a search in
		// flight — is skipped rather than reported.
		if e, ok := byPos[h.pos]; ok {
			out = append(out, storage.SearchResult{Event: e, Rank: h.rank})
		}
	}
	return out, total, nil
}

// searchHit is one row of a search page: where the event is, and how well
// it matched.
type searchHit struct {
	pos  int64
	rank float64
}

// searchHits runs the page query of a search.
func (s *Rooms) searchHits(ctx context.Context, query string, args []any) ([]searchHit, error) {
	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("search: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var hits []searchHit
	for rows.Next() {
		var h searchHit
		if err := rows.Scan(&h.pos, &h.rank); err != nil {
			return nil, fmt.Errorf("search: %w", err)
		}
		hits = append(hits, h)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("search: %w", err)
	}
	return hits, nil
}

// queryStrings runs a query returning one string column.
func (s *Rooms) queryStrings(ctx context.Context, query string, args ...any) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()

	var out []string
	for rows.Next() {
		var v string
		if err := rows.Scan(&v); err != nil {
			return nil, err
		}
		out = append(out, v)
	}
	return out, rows.Err()
}