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

package pocketbase

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

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

// DelayedEvents is the PocketBase implementation of
// storage.DelayedEventStore.
type DelayedEvents struct{ *base }

// delayedEventColumns reads state_key as NULL for a message event, as the
// SQL drivers store it.
const delayedEventColumns = `
	delay_id, user_id, device_id, room_id, event_type, CASE WHEN has_state_key THEN state_key END, content,
	delay_ms, running_since_ms, send_at_ms,
	finalised_at_ms, outcome, reason, event_id, error`

// InsertDelayedEvent implements storage.DelayedEventStore.
func (s *DelayedEvents) InsertDelayedEvent(ctx context.Context, e storage.DelayedEvent) error {
	stateKey, hasStateKey := "", false
	if e.StateKey != nil {
		stateKey, hasStateKey = *e.StateKey, true
	}
	err := s.write(ctx, func(t *tx) error {
		t.put("delayed_events", rec{
			"delay_id": e.DelayID, "user_id": e.UserID, "device_id": e.DeviceID, "room_id": e.RoomID,
			"event_type": e.Type, "state_key": stateKey, "has_state_key": hasStateKey,
			"content": string(e.Content), "delay_ms": e.DelayMS, "running_since_ms": e.RunningSinceMS,
			"send_at_ms": e.SendAtMS, "finalised_at_ms": int64(0),
			"outcome": "", "reason": "", "event_id": "", "error": "",
		})
		return nil
	})
	if err != nil {
		return fmt.Errorf("insert delayed event: %w", err)
	}
	return nil
}

// DelayedEvent implements storage.DelayedEventStore.
func (s *DelayedEvents) DelayedEvent(ctx context.Context, userID, delayID string) (storage.DelayedEvent, error) {
	return s.one(ctx, `SELECT`+delayedEventColumns+`
		FROM delayed_events WHERE user_id = ? AND delay_id = ?`, userID, delayID)
}

// ScheduledDelayedEvent implements storage.DelayedEventStore.
func (s *DelayedEvents) ScheduledDelayedEvent(ctx context.Context, delayID string) (storage.DelayedEvent, error) {
	return s.one(ctx, `SELECT`+delayedEventColumns+`
		FROM delayed_events WHERE delay_id = ? AND finalised_at_ms = 0`, delayID)
}

func (s *DelayedEvents) one(ctx context.Context, query string, args ...any) (storage.DelayedEvent, error) {
	events, err := s.list(ctx, query, args...)
	if err != nil {
		return storage.DelayedEvent{}, err
	}
	if len(events) == 0 {
		return storage.DelayedEvent{}, storage.ErrNoSuchDelayedEvent
	}
	return events[0], nil
}

// DelayedEvents implements storage.DelayedEventStore.
func (s *DelayedEvents) DelayedEvents(ctx context.Context, userID string) ([]storage.DelayedEvent, error) {
	return s.list(ctx, `SELECT`+delayedEventColumns+`
		FROM delayed_events WHERE user_id = ?
		ORDER BY finalised_at_ms = 0 DESC, send_at_ms ASC, finalised_at_ms DESC`, userID)
}

// CountScheduledDelayedEvents implements storage.DelayedEventStore.
func (s *DelayedEvents) CountScheduledDelayedEvents(ctx context.Context, userID string) (int, error) {
	var n int
	if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM delayed_events
		WHERE user_id = ? AND finalised_at_ms = 0`, userID).Scan(&n); err != nil {
		return 0, fmt.Errorf("count delayed events: %w", err)
	}
	return n, nil
}

// RestartDelayedEvent implements storage.DelayedEventStore.
func (s *DelayedEvents) RestartDelayedEvent(ctx context.Context, delayID string, nowMS int64) error {
	return s.write(ctx, func(t *tx) error {
		var delay int64
		err := t.queryRow(`SELECT delay_ms FROM delayed_events WHERE delay_id = ? AND finalised_at_ms = 0`, delayID).
			Scan(&delay)
		switch {
		case errors.Is(err, sql.ErrNoRows):
			return storage.ErrNoSuchDelayedEvent
		case err != nil:
			return fmt.Errorf("restart delayed event: %w", err)
		}
		t.put("delayed_events", rec{"delay_id": delayID, "running_since_ms": nowMS, "send_at_ms": nowMS + delay})
		return nil
	})
}

// FinaliseDelayedEvent implements storage.DelayedEventStore.
func (s *DelayedEvents) FinaliseDelayedEvent(ctx context.Context, e storage.DelayedEvent) error {
	return s.write(ctx, func(t *tx) error {
		n, err := t.setWhere("delayed_events", rec{
			"finalised_at_ms": e.FinalisedAtMS, "outcome": e.Outcome, "reason": e.Reason,
			"event_id": e.EventID, "error": string(e.Error),
		}, `delay_id = ? AND finalised_at_ms = 0`, e.DelayID)
		if err != nil {
			return fmt.Errorf("finalise delayed event: %w", err)
		}
		if n == 0 {
			return storage.ErrNoSuchDelayedEvent
		}
		return nil
	})
}

// DueDelayedEvents implements storage.DelayedEventStore.
func (s *DelayedEvents) DueDelayedEvents(ctx context.Context, nowMS int64, limit int) ([]storage.DelayedEvent, error) {
	return s.list(ctx, `SELECT`+delayedEventColumns+`
		FROM delayed_events WHERE finalised_at_ms = 0 AND send_at_ms <= ?
		ORDER BY send_at_ms LIMIT ?`, nowMS, limit)
}

// NextDelayedEventDue implements storage.DelayedEventStore.
func (s *DelayedEvents) NextDelayedEventDue(ctx context.Context) (int64, bool, error) {
	var next sql.NullInt64
	if err := s.db.QueryRowContext(ctx, `SELECT MIN(send_at_ms) FROM delayed_events
		WHERE finalised_at_ms = 0`).Scan(&next); err != nil {
		return 0, false, fmt.Errorf("next delayed event: %w", err)
	}
	return next.Int64, next.Valid, nil
}

// ScheduledDelayedStateEvents implements storage.DelayedEventStore.
func (s *DelayedEvents) ScheduledDelayedStateEvents(ctx context.Context, roomID, eventType, stateKey string) ([]storage.DelayedEvent, error) {
	return s.list(ctx, `SELECT`+delayedEventColumns+`
		FROM delayed_events
		WHERE room_id = ? AND event_type = ? AND has_state_key AND state_key = ? AND finalised_at_ms = 0`,
		roomID, eventType, stateKey)
}

// PurgeFinalisedDelayedEvents implements storage.DelayedEventStore.
func (s *DelayedEvents) PurgeFinalisedDelayedEvents(ctx context.Context, beforeMS int64) error {
	err := s.write(ctx, func(t *tx) error {
		_, err := t.delWhere("delayed_events", `finalised_at_ms <> 0 AND finalised_at_ms < ?`, beforeMS)
		return err
	})
	if err != nil {
		return fmt.Errorf("purge delayed events: %w", err)
	}
	return nil
}

func (s *DelayedEvents) list(ctx context.Context, query string, args ...any) ([]storage.DelayedEvent, error) {
	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("read delayed events: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.DelayedEvent
	for rows.Next() {
		var (
			e                     storage.DelayedEvent
			stateKey              sql.NullString
			content, errorContent string
		)
		if err := rows.Scan(&e.DelayID, &e.UserID, &e.DeviceID, &e.RoomID, &e.Type, &stateKey, &content,
			&e.DelayMS, &e.RunningSinceMS, &e.SendAtMS,
			&e.FinalisedAtMS, &e.Outcome, &e.Reason, &e.EventID, &errorContent,
		); err != nil {
			return nil, fmt.Errorf("read delayed event: %w", err)
		}
		if stateKey.Valid {
			e.StateKey = &stateKey.String
		}
		e.Content = []byte(content)
		if errorContent != "" {
			e.Error = []byte(errorContent)
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

// OpenIDTokens is the PocketBase implementation of storage.OpenIDTokenStore.
type OpenIDTokens struct{ *base }

// InsertOpenIDToken implements storage.OpenIDTokenStore.
func (s *OpenIDTokens) InsertOpenIDToken(ctx context.Context, tokenHash, userID string, expiresAtMS int64) error {
	err := s.write(ctx, func(t *tx) error {
		t.put("openid_tokens", rec{"token_hash": tokenHash, "user_id": userID, "expires_at_ms": expiresAtMS})
		return nil
	})
	if err != nil {
		return fmt.Errorf("insert OpenID token: %w", err)
	}
	return nil
}

// OpenIDTokenUser implements storage.OpenIDTokenStore.
func (s *OpenIDTokens) OpenIDTokenUser(ctx context.Context, tokenHash string, nowMS int64) (string, error) {
	var userID string
	err := s.db.QueryRowContext(ctx, `SELECT user_id FROM openid_tokens
		WHERE token_hash = ? AND expires_at_ms > ?`, tokenHash, nowMS).Scan(&userID)
	if errors.Is(err, sql.ErrNoRows) {
		return "", storage.ErrNoSuchOpenIDToken
	}
	if err != nil {
		return "", fmt.Errorf("read OpenID token: %w", err)
	}
	return userID, nil
}

// PurgeOpenIDTokens implements storage.OpenIDTokenStore.
func (s *OpenIDTokens) PurgeOpenIDTokens(ctx context.Context, beforeMS int64) error {
	err := s.write(ctx, func(t *tx) error {
		_, err := t.delWhere("openid_tokens", `expires_at_ms < ?`, beforeMS)
		return err
	})
	if err != nil {
		return fmt.Errorf("purge OpenID tokens: %w", err)
	}
	return nil
}

var (
	_ storage.DelayedEventStore = (*DelayedEvents)(nil)
	_ storage.OpenIDTokenStore  = (*OpenIDTokens)(nil)
)