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, thread_id, 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 <= ?` // markThreadReadSQL is markReadSQL for one thread, or for the main // timeline as thread ''. markThreadReadSQL = markReadSQL + ` AND thread_id = ?` 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, thread_id, 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.ThreadID, 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. The room's counts stay the room's total across every // thread whichever thread the receipt read. func (s *Push) MarkNotificationsRead(ctx context.Context, userID, roomID, thread string, upTo int64) (int64, bool, error) { var ( pos int64 moved bool ) err := s.tx(ctx, func(tx *sql.Tx) error { query, args := markReadSQL, []any{true, userID, roomID, false, upTo} if thread != storage.AllThreads { query, args = markThreadReadSQL, append(args, thread) } res, err := tx.ExecContext(ctx, s.q(query), args...) 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(¬ifications, &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 } // ThreadNotificationCounts implements storage.PushStore. // // These are counted from the unread rows rather than kept like the room // totals: a room has few unread notifications at a time, they are a range // of an index, and a table of per-thread counts would be one more thing // every write and every receipt has to keep in step. 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 } query := `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 (` + placeholders(len(roomIDs)) + `) GROUP BY room_id, thread_id` args := make([]any, 0, len(roomIDs)+2) args = append(args, userID, false) 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 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, ¬ifications, &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) { 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.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) { 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 }