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

package sqlstore

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

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

// The push store's sequence table and its one counter.
const (
	pushSequences   = "push_sequences"
	seqNotification = "notification"
)

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

// NewPush builds a push store.
func NewPush(db *sql.DB, rebind func(string) string) *Push {
	if rebind == nil {
		rebind = func(query string) string { return query }
	}
	return &Push{db: db, rebind: rebind}
}

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

// tx runs fn in a transaction, rolling back on any error.
func (s *Push) tx(ctx context.Context, fn func(*sql.Tx) error) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin transaction: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	if err := fn(tx); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit transaction: %w", err)
	}
	return nil
}

const (
	pusherColumns = `localpart, app_id, pushkey, device_id, kind, app_display_name,
		device_display_name, profile_tag, lang, data_json, pushkey_ts_ms`

	selectPushersSQL = `SELECT ` + pusherColumns + ` FROM push_pushers
		WHERE localpart = ? ORDER BY app_id, pushkey`

	upsertPusherSQL = `INSERT INTO push_pushers (` + pusherColumns + `)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT (localpart, app_id, pushkey) DO UPDATE SET
			device_id = excluded.device_id,
			kind = excluded.kind,
			app_display_name = excluded.app_display_name,
			device_display_name = excluded.device_display_name,
			profile_tag = excluded.profile_tag,
			lang = excluded.lang,
			data_json = excluded.data_json,
			pushkey_ts_ms = excluded.pushkey_ts_ms`

	deleteOtherUsersPushersSQL = `DELETE FROM push_pushers
		WHERE app_id = ? AND pushkey = ? AND localpart <> ?`

	deletePusherSQL = `DELETE FROM push_pushers
		WHERE localpart = ? AND app_id = ? AND pushkey = ?`

	deleteAllPushersSQL = `DELETE FROM push_pushers WHERE localpart = ?`
)

// Pushers implements storage.PushStore.
func (s *Push) Pushers(ctx context.Context, localpart string) ([]storage.Pusher, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectPushersSQL), localpart)
	if err != nil {
		return nil, fmt.Errorf("read the pushers of %q: %w", localpart, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.Pusher
	for rows.Next() {
		var p storage.Pusher
		var data string
		if err := rows.Scan(&p.Localpart, &p.AppID, &p.PushKey, &p.DeviceID, &p.Kind, &p.AppDisplayName,
			&p.DeviceDisplayName, &p.ProfileTag, &p.Lang, &data, &p.PushKeyTSMS); err != nil {
			return nil, fmt.Errorf("read the pushers of %q: %w", localpart, err)
		}
		p.Data = []byte(data)
		out = append(out, p)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read the pushers of %q: %w", localpart, err)
	}
	return out, nil
}

// SetPusher implements storage.PushStore.
func (s *Push) SetPusher(ctx context.Context, p storage.Pusher, appendOthers bool) error {
	return s.tx(ctx, func(tx *sql.Tx) error {
		if !appendOthers {
			if _, err := tx.ExecContext(ctx, s.q(deleteOtherUsersPushersSQL), p.AppID, p.PushKey, p.Localpart); err != nil {
				return fmt.Errorf("remove other users' pushers for the key: %w", err)
			}
		}
		if _, err := tx.ExecContext(ctx, s.q(upsertPusherSQL), p.Localpart, p.AppID, p.PushKey, p.DeviceID, p.Kind,
			p.AppDisplayName, p.DeviceDisplayName, p.ProfileTag, p.Lang, string(p.Data), p.PushKeyTSMS); err != nil {
			return fmt.Errorf("write the pusher of %q: %w", p.Localpart, err)
		}
		return nil
	})
}

// DeletePusher implements storage.PushStore.
func (s *Push) DeletePusher(ctx context.Context, localpart, appID, pushKey string) error {
	if _, err := s.db.ExecContext(ctx, s.q(deletePusherSQL), localpart, appID, pushKey); err != nil {
		return fmt.Errorf("delete a pusher of %q: %w", localpart, err)
	}
	return nil
}

// DeletePushersForDevices implements storage.PushStore.
func (s *Push) DeletePushersForDevices(ctx context.Context, localpart string, deviceIDs []string) error {
	query, args := deleteAllPushersSQL, []any{localpart}
	if len(deviceIDs) > 0 {
		query += ` AND device_id IN (` + placeholders(len(deviceIDs)) + `)`
		for _, id := range deviceIDs {
			args = append(args, id)
		}
	}
	if _, err := s.db.ExecContext(ctx, s.q(query), args...); err != nil {
		return fmt.Errorf("delete the pushers of %q's devices: %w", localpart, err)
	}
	return nil
}

// The sums are cast back to BIGINT because PostgreSQL widens SUM(BIGINT)
// to NUMERIC, which does not scan into an int64.
const (
	insertNotificationSQL = `INSERT INTO push_notifications
		(id, user_id, room_id, event_id, event_pos, ts_ms, highlight, actions, is_read)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`

	bumpRoomCountsSQL = `INSERT INTO push_room_counts
		(user_id, room_id, notification_count, highlight_count, stream_pos)
		VALUES (?, ?, 1, ?, ?)
		ON CONFLICT (user_id, room_id) DO UPDATE SET
			notification_count = push_room_counts.notification_count + 1,
			highlight_count = push_room_counts.highlight_count + excluded.highlight_count,
			stream_pos = excluded.stream_pos`

	markReadSQL = `UPDATE push_notifications SET is_read = ?
		WHERE user_id = ? AND room_id = ? AND is_read = ? AND event_pos <= ?`

	countUnreadSQL = `SELECT COUNT(*), CAST(COALESCE(SUM(CASE WHEN highlight THEN 1 ELSE 0 END), 0) AS BIGINT)
		FROM push_notifications WHERE user_id = ? AND room_id = ? AND is_read = ?`

	setRoomCountsSQL = `INSERT INTO push_room_counts
		(user_id, room_id, notification_count, highlight_count, stream_pos)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT (user_id, room_id) DO UPDATE SET
			notification_count = excluded.notification_count,
			highlight_count = excluded.highlight_count,
			stream_pos = excluded.stream_pos`

	selectCountChangesSQL = `SELECT room_id FROM push_room_counts
		WHERE user_id = ? AND stream_pos > ? AND stream_pos <= ? ORDER BY room_id`

	selectUnreadTotalSQL = `SELECT CAST(COALESCE(SUM(notification_count), 0) AS BIGINT)
		FROM push_room_counts WHERE user_id = ?`

	notificationColumns = `id, user_id, room_id, event_id, event_pos, ts_ms, highlight, actions, is_read`

	pruneNotificationsSQL = `DELETE FROM push_notifications WHERE is_read = ? AND ts_ms < ?`
)

// AddNotifications implements storage.PushStore.
//
// Every notification takes a position of its own, which is both its ID
// and the stamp on its room's counts: a sync that has seen a position has
// seen every count change at or below it.
func (s *Push) AddNotifications(ctx context.Context, ns []storage.NewNotification) (int64, error) {
	if len(ns) == 0 {
		return currentSequence(ctx, s.db, s.q, pushSequences, seqNotification)
	}
	var last int64
	err := s.tx(ctx, func(tx *sql.Tx) error {
		for _, n := range ns {
			pos, err := nextSequence(ctx, tx, s.q, pushSequences, seqNotification)
			if err != nil {
				return err
			}
			if _, err := tx.ExecContext(ctx, s.q(insertNotificationSQL), pos, n.UserID, n.RoomID, n.EventID,
				n.EventPos, n.TimestampMS, n.Highlight, string(n.Actions), false); err != nil {
				return fmt.Errorf("write a notification for %q: %w", n.UserID, err)
			}
			highlight := 0
			if n.Highlight {
				highlight = 1
			}
			if _, err := tx.ExecContext(ctx, s.q(bumpRoomCountsSQL), n.UserID, n.RoomID, highlight, pos); err != nil {
				return fmt.Errorf("count a notification for %q: %w", n.UserID, err)
			}
			last = pos
		}
		return nil
	})
	return last, err
}

// MarkNotificationsRead implements storage.PushStore.
//
// The counts are recounted from the rows that are still unread rather than
// decremented, so that whatever got them out of step — a crash between
// two writes of an older version, a row pruned by hand — is put right by
// the next receipt.
func (s *Push) MarkNotificationsRead(ctx context.Context, userID, roomID string, upTo int64) (int64, bool, error) {
	var (
		pos   int64
		moved bool
	)
	err := s.tx(ctx, func(tx *sql.Tx) error {
		res, err := tx.ExecContext(ctx, s.q(markReadSQL), true, userID, roomID, false, upTo)
		if err != nil {
			return fmt.Errorf("mark %q's notifications read: %w", userID, err)
		}
		n, err := res.RowsAffected()
		if err != nil {
			return fmt.Errorf("mark %q's notifications read: %w", userID, err)
		}
		if n == 0 {
			pos, err = currentSequence(ctx, tx, s.q, pushSequences, seqNotification)
			return err
		}

		var notifications, highlights int64
		if err := tx.QueryRowContext(ctx, s.q(countUnreadSQL), userID, roomID, false).
			Scan(&notifications, &highlights); err != nil {
			return fmt.Errorf("count %q's unread notifications: %w", userID, err)
		}
		if pos, err = nextSequence(ctx, tx, s.q, pushSequences, seqNotification); err != nil {
			return err
		}
		if _, err := tx.ExecContext(ctx, s.q(setRoomCountsSQL), userID, roomID, notifications, highlights, pos); err != nil {
			return fmt.Errorf("write %q's notification counts: %w", userID, err)
		}
		moved = true
		return nil
	})
	return pos, moved, err
}

// NotificationCounts implements storage.PushStore.
func (s *Push) NotificationCounts(ctx context.Context, userID string, roomIDs []string) (map[string]storage.NotificationCounts, error) {
	out := make(map[string]storage.NotificationCounts, len(roomIDs))
	if len(roomIDs) == 0 {
		return out, nil
	}
	for _, id := range roomIDs {
		out[id] = storage.NotificationCounts{}
	}

	query := `SELECT room_id, notification_count, highlight_count FROM push_room_counts
		WHERE user_id = ? AND room_id IN (` + placeholders(len(roomIDs)) + `)`
	args := make([]any, 0, len(roomIDs)+1)
	args = append(args, userID)
	for _, id := range roomIDs {
		args = append(args, id)
	}
	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read %q's notification counts: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()
	for rows.Next() {
		var roomID string
		var c storage.NotificationCounts
		if err := rows.Scan(&roomID, &c.Notifications, &c.Highlights); err != nil {
			return nil, fmt.Errorf("read %q's notification counts: %w", userID, err)
		}
		out[roomID] = c
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read %q's notification counts: %w", userID, err)
	}
	return out, nil
}

// NotificationCountChanges implements storage.PushStore.
func (s *Push) NotificationCountChanges(ctx context.Context, userID string, from, to int64) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectCountChangesSQL), userID, from, to)
	if err != nil {
		return nil, fmt.Errorf("read %q's notification count changes: %w", userID, 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("read %q's notification count changes: %w", userID, err)
		}
		out = append(out, roomID)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read %q's notification count changes: %w", userID, err)
	}
	return out, nil
}

// UnreadNotifications implements storage.PushStore.
func (s *Push) UnreadNotifications(ctx context.Context, userID string) (int, error) {
	var n int64
	if err := s.db.QueryRowContext(ctx, s.q(selectUnreadTotalSQL), userID).Scan(&n); err != nil {
		return 0, fmt.Errorf("count %q's unread notifications: %w", userID, err)
	}
	return int(n), nil
}

// Notifications implements storage.PushStore.
func (s *Push) Notifications(ctx context.Context, userID string, before int64, limit int, onlyHighlight bool) ([]storage.Notification, error) {
	query := `SELECT ` + notificationColumns + ` FROM push_notifications WHERE user_id = ?`
	args := []any{userID}
	if before > 0 {
		query += ` AND id < ?`
		args = append(args, before)
	}
	if onlyHighlight {
		query += ` AND highlight = ?`
		args = append(args, true)
	}
	query += ` ORDER BY id DESC LIMIT ?`
	args = append(args, limit)

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read %q's notifications: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()
	var out []storage.Notification
	for rows.Next() {
		var n storage.Notification
		var actions string
		if err := rows.Scan(&n.ID, &n.UserID, &n.RoomID, &n.EventID, &n.EventPos, &n.TimestampMS,
			&n.Highlight, &actions, &n.Read); err != nil {
			return nil, fmt.Errorf("read %q's notifications: %w", userID, err)
		}
		n.Actions = []byte(actions)
		out = append(out, n)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read %q's notifications: %w", userID, err)
	}
	return out, nil
}

// PruneNotifications implements storage.PushStore.
func (s *Push) PruneNotifications(ctx context.Context, beforeMS int64) (int64, error) {
	res, err := s.db.ExecContext(ctx, s.q(pruneNotificationsSQL), true, beforeMS)
	if err != nil {
		return 0, fmt.Errorf("prune read notifications: %w", err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("prune read notifications: %w", err)
	}
	return n, nil
}