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

package sqlstore

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

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

// DelayedEvents is the shared implementation of storage.DelayedEventStore.
type DelayedEvents struct {
	db     *sql.DB
	rebind func(string) string
}

// NewDelayedEvents builds a delayed event store.
func NewDelayedEvents(db *sql.DB, rebind func(string) string) *DelayedEvents {
	if rebind == nil {
		rebind = func(query string) string { return query }
	}
	return &DelayedEvents{db: db, rebind: rebind}
}

func (s *DelayedEvents) q(query string) string { return s.rebind(query) }

const delayedEventColumns = `
	delay_id, user_id, device_id, room_id, event_type, state_key, content,
	delay_ms, running_since_ms, send_at_ms,
	finalised_at_ms, outcome, reason, event_id, error`

const (
	insertDelayedEventSQL = `
		INSERT INTO delayed_events (
			delay_id, user_id, device_id, room_id, event_type, state_key, content,
			delay_ms, running_since_ms, send_at_ms)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`

	selectDelayedEventSQL = `SELECT` + delayedEventColumns + `
		FROM delayed_events WHERE user_id = ? AND delay_id = ?`

	selectScheduledDelayedEventSQL = `SELECT` + delayedEventColumns + `
		FROM delayed_events WHERE delay_id = ? AND finalised_at_ms = 0`

	selectUserDelayedEventsSQL = `SELECT` + delayedEventColumns + `
		FROM delayed_events WHERE user_id = ?
		ORDER BY finalised_at_ms = 0 DESC, send_at_ms ASC, finalised_at_ms DESC`

	countScheduledDelayedEventsSQL = `
		SELECT COUNT(*) FROM delayed_events WHERE user_id = ? AND finalised_at_ms = 0`

	restartDelayedEventSQL = `
		UPDATE delayed_events SET running_since_ms = ?, send_at_ms = ? + delay_ms
		WHERE delay_id = ? AND finalised_at_ms = 0`

	finaliseDelayedEventSQL = `
		UPDATE delayed_events
		SET finalised_at_ms = ?, outcome = ?, reason = ?, event_id = ?, error = ?
		WHERE delay_id = ? AND finalised_at_ms = 0`

	selectDueDelayedEventsSQL = `SELECT` + delayedEventColumns + `
		FROM delayed_events WHERE finalised_at_ms = 0 AND send_at_ms <= ?
		ORDER BY send_at_ms LIMIT ?`

	selectNextDelayedEventDueSQL = `
		SELECT MIN(send_at_ms) FROM delayed_events WHERE finalised_at_ms = 0`

	selectScheduledStateEventsSQL = `SELECT` + delayedEventColumns + `
		FROM delayed_events
		WHERE room_id = ? AND event_type = ? AND state_key = ? AND finalised_at_ms = 0`

	purgeFinalisedDelayedEventsSQL = `
		DELETE FROM delayed_events WHERE finalised_at_ms <> 0 AND finalised_at_ms < ?`
)

// InsertDelayedEvent implements storage.DelayedEventStore.
func (s *DelayedEvents) InsertDelayedEvent(ctx context.Context, e storage.DelayedEvent) error {
	var stateKey sql.NullString
	if e.StateKey != nil {
		stateKey = sql.NullString{String: *e.StateKey, Valid: true}
	}
	if _, err := s.db.ExecContext(ctx, s.q(insertDelayedEventSQL),
		e.DelayID, e.UserID, e.DeviceID, e.RoomID, e.Type, stateKey, string(e.Content),
		e.DelayMS, e.RunningSinceMS, e.SendAtMS,
	); 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, selectDelayedEventSQL, userID, delayID)
}

// ScheduledDelayedEvent implements storage.DelayedEventStore.
func (s *DelayedEvents) ScheduledDelayedEvent(ctx context.Context, delayID string) (storage.DelayedEvent, error) {
	return s.one(ctx, selectScheduledDelayedEventSQL, 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, selectUserDelayedEventsSQL, 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, s.q(countScheduledDelayedEventsSQL), 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 {
	res, err := s.db.ExecContext(ctx, s.q(restartDelayedEventSQL), nowMS, nowMS, delayID)
	if err != nil {
		return fmt.Errorf("restart delayed event: %w", err)
	}
	return oneRow(res)
}

// FinaliseDelayedEvent implements storage.DelayedEventStore.
func (s *DelayedEvents) FinaliseDelayedEvent(ctx context.Context, e storage.DelayedEvent) error {
	res, err := s.db.ExecContext(ctx, s.q(finaliseDelayedEventSQL),
		e.FinalisedAtMS, e.Outcome, e.Reason, e.EventID, string(e.Error), e.DelayID)
	if err != nil {
		return fmt.Errorf("finalise delayed event: %w", err)
	}
	return oneRow(res)
}

func oneRow(res sql.Result) error {
	n, err := res.RowsAffected()
	if err != nil {
		return fmt.Errorf("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, selectDueDelayedEventsSQL, 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, s.q(selectNextDelayedEventDueSQL)).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, selectScheduledStateEventsSQL, roomID, eventType, stateKey)
}

// PurgeFinalisedDelayedEvents implements storage.DelayedEventStore.
func (s *DelayedEvents) PurgeFinalisedDelayedEvents(ctx context.Context, beforeMS int64) error {
	if _, err := s.db.ExecContext(ctx, s.q(purgeFinalisedDelayedEventsSQL), beforeMS); 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, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read delayed events: %w", err)
	}
	defer 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 shared implementation of storage.OpenIDTokenStore.
type OpenIDTokens struct {
	db     *sql.DB
	rebind func(string) string
}

// NewOpenIDTokens builds an OpenID token store.
func NewOpenIDTokens(db *sql.DB, rebind func(string) string) *OpenIDTokens {
	if rebind == nil {
		rebind = func(query string) string { return query }
	}
	return &OpenIDTokens{db: db, rebind: rebind}
}

const (
	insertOpenIDTokenSQL = `
		INSERT INTO openid_tokens (token_hash, user_id, expires_at_ms) VALUES (?, ?, ?)`

	selectOpenIDTokenSQL = `
		SELECT user_id FROM openid_tokens WHERE token_hash = ? AND expires_at_ms > ?`

	purgeOpenIDTokensSQL = `DELETE FROM openid_tokens WHERE expires_at_ms < ?`
)

// InsertOpenIDToken implements storage.OpenIDTokenStore.
func (s *OpenIDTokens) InsertOpenIDToken(ctx context.Context, tokenHash, userID string, expiresAtMS int64) error {
	if _, err := s.db.ExecContext(ctx, s.rebind(insertOpenIDTokenSQL), tokenHash, userID, expiresAtMS); 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, s.rebind(selectOpenIDTokenSQL), 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 {
	if _, err := s.db.ExecContext(ctx, s.rebind(purgeOpenIDTokensSQL), beforeMS); err != nil {
		return fmt.Errorf("purge OpenID tokens: %w", err)
	}
	return nil
}

// Compile-time proof that the shared implementations satisfy the interfaces.
var (
	_ storage.DelayedEventStore = (*DelayedEvents)(nil)
	_ storage.OpenIDTokenStore  = (*OpenIDTokens)(nil)
)