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

package pocketbase

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

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

// Push is the PocketBase implementation of storage.PushStore.
type Push struct{ *base }

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

// Pushers implements storage.PushStore.
func (s *Push) Pushers(ctx context.Context, localpart string) ([]storage.Pusher, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT `+pusherColumns+` FROM push_pushers
		WHERE localpart = ? ORDER BY app_id, pushkey`, 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.write(ctx, func(t *tx) error {
		if !appendOthers {
			if _, err := t.delWhere("push_pushers", `app_id = ? AND pushkey = ? AND localpart <> ?`,
				p.AppID, p.PushKey, p.Localpart); err != nil {
				return fmt.Errorf("remove other users' pushers for the key: %w", err)
			}
		}
		t.put("push_pushers", rec{
			"localpart": p.Localpart, "app_id": p.AppID, "pushkey": p.PushKey, "device_id": p.DeviceID,
			"kind": p.Kind, "app_display_name": p.AppDisplayName, "device_display_name": p.DeviceDisplayName,
			"profile_tag": p.ProfileTag, "lang": p.Lang, "data_json": string(p.Data), "pushkey_ts_ms": p.PushKeyTSMS,
		})
		return nil
	})
}

// DeletePusher implements storage.PushStore.
func (s *Push) DeletePusher(ctx context.Context, localpart, appID, pushKey string) error {
	err := s.write(ctx, func(t *tx) error {
		_, err := t.delKey("push_pushers", localpart, appID, pushKey)
		return err
	})
	if 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 {
	where, args := `localpart = ?`, []any{localpart}
	if len(deviceIDs) > 0 {
		where += ` AND device_id IN ` + inList
		args = append(args, jsonList(deviceIDs))
	}
	err := s.write(ctx, func(t *tx) error {
		_, err := t.delWhere("push_pushers", where, args...)
		return err
	})
	if err != nil {
		return fmt.Errorf("delete the pushers of %q's devices: %w", localpart, err)
	}
	return nil
}

const notificationColumns = `notification_id, user_id, room_id, event_id, event_pos, thread_id, ts_ms, highlight, actions, is_read`

// roomCounts is one room's counts for one user, as a write leaves them.
type roomCounts struct{ notifications, highlights int64 }

// 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.
func (s *Push) AddNotifications(ctx context.Context, ns []storage.NewNotification) (int64, error) {
	if len(ns) == 0 {
		return s.currentSequence(ctx, pushSequences, seqNotification)
	}
	var last int64
	err := s.write(ctx, func(t *tx) error {
		counts := map[[2]string]roomCounts{}
		for _, n := range ns {
			pos, err := t.next(pushSequences, seqNotification)
			if err != nil {
				return err
			}
			t.put("push_notifications", rec{
				"notification_id": pos, "user_id": n.UserID, "room_id": n.RoomID, "event_id": n.EventID,
				"event_pos": n.EventPos, "thread_id": n.ThreadID, "ts_ms": n.TimestampMS,
				"highlight": n.Highlight, "actions": string(n.Actions), "is_read": false,
			})

			key := [2]string{n.UserID, n.RoomID}
			c, ok := counts[key]
			if !ok {
				err := t.queryRow(`SELECT notification_count, highlight_count FROM push_room_counts
					WHERE user_id = ? AND room_id = ?`, n.UserID, n.RoomID).Scan(&c.notifications, &c.highlights)
				if err != nil && !errors.Is(err, sql.ErrNoRows) {
					return fmt.Errorf("count a notification for %q: %w", n.UserID, err)
				}
			}
			c.notifications++
			if n.Highlight {
				c.highlights++
			}
			counts[key] = c
			t.put("push_room_counts", rec{
				"user_id": n.UserID, "room_id": n.RoomID,
				"notification_count": c.notifications, "highlight_count": c.highlights, "stream_pos": pos,
			})
			last = pos
		}
		return nil
	})
	return last, err
}

// MarkNotificationsRead implements storage.PushStore.
//
// The counts are recounted from the rows that stay unread rather than
// decremented, so whatever got them out of step is put right by the next
// receipt. The room's counts stay the room's total across every thread.
func (s *Push) MarkNotificationsRead(ctx context.Context, userID, roomID, thread string, upTo int64) (int64, bool, error) {
	var (
		pos   int64
		moved bool
	)
	err := s.write(ctx, func(t *tx) error {
		marked := `event_pos <= ?`
		markedArgs := []any{upTo}
		if thread != storage.AllThreads {
			marked += ` AND thread_id = ?`
			markedArgs = append(markedArgs, thread)
		}
		n, err := t.setWhere("push_notifications", rec{"is_read": true},
			`user_id = ? AND room_id = ? AND is_read = ? AND `+marked,
			append([]any{userID, roomID, false}, markedArgs...)...)
		if err != nil {
			return fmt.Errorf("mark %q's notifications read: %w", userID, err)
		}
		if n == 0 {
			pos, err = t.current(pushSequences, seqNotification)
			return err
		}

		var notifications, highlights int64
		if err := t.queryRow(`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 = ? AND NOT (`+marked+`)`,
			append([]any{userID, roomID, false}, markedArgs...)...).
			Scan(&notifications, &highlights); err != nil {
			return fmt.Errorf("count %q's unread notifications: %w", userID, err)
		}
		if pos, err = t.next(pushSequences, seqNotification); err != nil {
			return err
		}
		t.put("push_room_counts", rec{
			"user_id": userID, "room_id": roomID,
			"notification_count": notifications, "highlight_count": highlights, "stream_pos": pos,
		})
		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{}
	}

	rows, err := s.db.QueryContext(ctx, `SELECT room_id, notification_count, highlight_count FROM push_room_counts
		WHERE user_id = ? AND room_id IN `+inList, userID, jsonList(roomIDs))
	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
}

// ThreadNotificationCounts implements storage.PushStore. These are counted
// from the unread rows rather than kept like the room totals.
func (s *Push) ThreadNotificationCounts(ctx context.Context, userID string, roomIDs []string) (map[string]map[string]storage.NotificationCounts, error) {
	out := make(map[string]map[string]storage.NotificationCounts, len(roomIDs))
	if len(roomIDs) == 0 {
		return out, nil
	}

	rows, err := s.db.QueryContext(ctx, `SELECT room_id, thread_id, COUNT(*),
			CAST(COALESCE(SUM(CASE WHEN highlight THEN 1 ELSE 0 END), 0) AS BIGINT)
		FROM push_notifications
		WHERE user_id = ? AND is_read = ? AND room_id IN `+inList+`
		GROUP BY room_id, thread_id`, userID, false, jsonList(roomIDs))
	if err != nil {
		return nil, fmt.Errorf("read %q's thread notification counts: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()
	for rows.Next() {
		var roomID, thread string
		var notifications, highlights int64
		if err := rows.Scan(&roomID, &thread, &notifications, &highlights); err != nil {
			return nil, fmt.Errorf("read %q's thread notification counts: %w", userID, err)
		}
		if out[roomID] == nil {
			out[roomID] = map[string]storage.NotificationCounts{}
		}
		out[roomID][thread] = storage.NotificationCounts{
			Notifications: int(notifications),
			Highlights:    int(highlights),
		}
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read %q's thread 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) {
	out, err := s.queryStrings(ctx, `SELECT room_id FROM push_room_counts
		WHERE user_id = ? AND stream_pos > ? AND stream_pos <= ? ORDER BY room_id`, userID, from, to)
	if 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, `SELECT CAST(COALESCE(SUM(notification_count), 0) AS BIGINT)
		FROM push_room_counts WHERE user_id = ?`, 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 notification_id < ?`
		args = append(args, before)
	}
	if onlyHighlight {
		query += ` AND highlight = ?`
		args = append(args, true)
	}
	query += ` ORDER BY notification_id DESC LIMIT ?`
	args = append(args, limit)

	rows, err := s.db.QueryContext(ctx, 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.ThreadID, &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) {
	var n int
	err := s.write(ctx, func(t *tx) error {
		var err error
		n, err = t.delWhere("push_notifications", `is_read = ? AND ts_ms < ?`, true, beforeMS)
		return err
	})
	if err != nil {
		return 0, fmt.Errorf("prune read notifications: %w", err)
	}
	return int64(n), nil
}

var _ storage.PushStore = (*Push)(nil)