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

package pocketbase

import (
	"context"
	"fmt"
	"strings"

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

// The relation and search half of the room store.
//
// Relations are the SQL store's queries. Search has no full-text engine to
// lean on — PocketBase does not expose FTS5 — so roomserver_event_search
// holds each event's searchable text as its words, lower-cased and padded
// with spaces, and a search matches every word of the term as the start of
// a word: what the SQLite driver's prefix query matches, short of folding
// diacritics. Every match ranks the same, so a ranked search comes back
// newest first.

// searchText is how an event's text is stored for searching: its words,
// lower-cased, each with a space before it and the whole with one after.
func searchText(text string) string {
	words := storage.SearchWords(text)
	if len(words) == 0 {
		return ""
	}
	return " " + strings.Join(words, " ") + " "
}

// indexEvent adds an event that has just joined a timeline to the search
// index, if it has anything to find it by.
func (w *roomWrite) indexEvent(e storage.StoredEvent, ordering int64) {
	body := searchText(e.SearchText)
	if e.SearchKey == "" || body == "" {
		return
	}
	w.put("roomserver_event_search", rec{
		"stream_ordering": ordering,
		"room_id":         e.RoomID,
		"event_id":        e.EventID,
		"search_key":      e.SearchKey,
		"sender":          e.Sender,
		"event_type":      e.Type,
		"body":            body,
	})
}

// 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 ` + inList,
		`NOT e.outlier`,
	}
	args := []any{req.RoomID, jsonList(req.ParentIDs)}
	if req.RelType != "" {
		conditions = append(conditions, `e.rel_type = ?`)
		args = append(args, req.RelType)
	} else {
		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)

	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 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 ` + inList + `
		AND rel_type <> '' AND NOT outlier`
	args := []any{roomID, jsonList(parentIDs)}
	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.
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 ` + inList + `
		AND rel_type IN (?, ?, ?) AND NOT outlier`
	args := []any{roomID, jsonList(parentIDs), 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`

	rows, err := s.db.QueryContext(ctx, `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`, 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.
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 ` + inList + `
		AND rel_type = ? AND NOT outlier`
	args := []any{roomID, jsonList(parentIDs), 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, 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.
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)

	rows, err := s.db.QueryContext(ctx, `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 ?`, 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 req.Limit <= 0 || len(req.Rooms) == 0 {
		return nil, 0, nil
	}
	words := storage.SearchWords(req.Term)
	if len(words) == 0 {
		return nil, 0, nil
	}

	var (
		conditions []string
		args       []any
	)
	for _, w := range words {
		conditions = append(conditions, `body LIKE ? ESCAPE '\'`)
		args = append(args, `% `+escapeLike(w)+`%`)
	}

	// 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 := `room_id = ?`
		args = append(args, r.RoomID)
		if r.Low != 0 {
			clause += ` AND stream_ordering > ?`
			args = append(args, r.Low)
		}
		if r.High != 0 {
			clause += ` AND stream_ordering <= ?`
			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, `search_key IN `+inList)
	args = append(args, jsonList(keys))

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

	var total int
	if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM roomserver_event_search WHERE `+where, 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 stream_ordering < ?`
		pageArgs = append(pageArgs, req.Before)
	}
	offset := req.Offset
	if req.Order == storage.SearchByRecent {
		offset = 0
	}
	pageArgs = append(pageArgs, req.Limit, offset)

	positions, err := s.queryInts(ctx, `SELECT stream_ordering FROM roomserver_event_search
		WHERE `+pageWhere+`
		ORDER BY stream_ordering DESC
		LIMIT ? OFFSET ?`, pageArgs...)
	if err != nil {
		return nil, 0, fmt.Errorf("search: %w", err)
	}
	if len(positions) == 0 {
		return nil, total, nil
	}

	eventRows, err := s.db.QueryContext(ctx, `SELECT `+eventColumns+`
		FROM roomserver_events e WHERE e.stream_ordering IN `+inList, jsonInts(positions))
	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(positions))
	for _, pos := range positions {
		if e, ok := byPos[pos]; ok {
			out = append(out, storage.SearchResult{Event: e, Rank: 1})
		}
	}
	return out, total, nil
}

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

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