Store Web Push subscriptions in DB

internal/db/bundb/admin.go

@@ -48,9 +48,6 @@ const rsaKeyBits = 2048
 type adminDB struct {
 	db    *bun.DB
 	state *state.State
-
-	// Since the VAPID key pair is very small and never written to concurrently, we can cache it here.
-	vapidKeyPair *gtsmodel.VAPIDKeyPair
 }
 
 func (a *adminDB) IsUsernameAvailable(ctx context.Context, username string) (bool, error) {
@@ -445,39 +442,6 @@ func (a *adminDB) CountUnhandledSignups(ctx context.Context) (int, error) {
 		Count(ctx)
 }
 
-func (a *adminDB) GetVAPIDKeyPair(ctx context.Context) (*gtsmodel.VAPIDKeyPair, error) {
-	// Look for cached keys.
-	if a.vapidKeyPair != nil {
-		return a.vapidKeyPair, nil
-	}
-
-	// Look for previously generated keys in the database.
-	if err := a.db.NewSelect().
-		Model(a.vapidKeyPair).
-		Limit(1).
-		Scan(ctx); // nocollapse
-	err != nil && !errors.Is(err, db.ErrNoEntries) {
-		return nil, gtserror.Newf("DB error getting VAPID key pair: %w", err)
-	}
-
-	return a.vapidKeyPair, nil
-}
-
-func (a *adminDB) PutVAPIDKeyPair(ctx context.Context, vapidKeyPair *gtsmodel.VAPIDKeyPair) error {
-	// Store the keys in the database.
-	if _, err := a.db.NewInsert().
-		Model(a.vapidKeyPair).
-		Exec(ctx); // nocollapse
-	err != nil {
-		return gtserror.Newf("DB error putting VAPID key pair: %w", err)
-	}
-
-	// Cache the keys.
-	a.vapidKeyPair = vapidKeyPair
-
-	return nil
-}
-
 /*
 	ACTION FUNCS
 */

internal/db/bundb/application.go

@@ -174,6 +174,16 @@ func (a *applicationDB) GetAllTokens(ctx context.Context) ([]*gtsmodel.Token, er
 	return tokens, nil
 }
 
+func (a *applicationDB) GetTokenByID(ctx context.Context, code string) (*gtsmodel.Token, error) {
+	return a.getTokenBy(
+		"ID",
+		func(t *gtsmodel.Token) error {
+			return a.db.NewSelect().Model(t).Where("? = ?", bun.Ident("id"), code).Scan(ctx)
+		},
+		code,
+	)
+}
+
 func (a *applicationDB) GetTokenByCode(ctx context.Context, code string) (*gtsmodel.Token, error) {
 	return a.getTokenBy(
 		"Code",

internal/db/bundb/bundb.go

@@ -88,6 +88,7 @@ type DBService struct {
 	db.Timeline
 	db.User
 	db.Tombstone
+	db.WebPush
 	db.WorkerTask
 	db *bun.DB
 }
@@ -301,6 +302,10 @@ func NewBunDBService(ctx context.Context, state *state.State) (db.DB, error) {
 			db:    db,
 			state: state,
 		},
+		WebPush: &webPushDB{
+			db:    db,
+			state: state,
+		},
 		WorkerTask: &workerTaskDB{
 			db: db,
 		},

internal/db/bundb/notification_test.go

@@ -66,7 +66,7 @@ func (suite *NotificationTestSuite) spamNotifs() {
 
 		notif := &gtsmodel.Notification{
 			ID:               notifID,
-			NotificationType: gtsmodel.NotificationFave,
+			NotificationType: gtsmodel.NotificationFavourite,
 			CreatedAt:        time.Now(),
 			TargetAccountID:  targetAccountID,
 			OriginAccountID:  originAccountID,

internal/db/bundb/webpush.go

@@ -0,0 +1,203 @@
+package bundb
+
+import (
+	"context"
+	"errors"
+
+	"github.com/superseriousbusiness/gotosocial/internal/db"
+	"github.com/superseriousbusiness/gotosocial/internal/gtsmodel"
+	"github.com/superseriousbusiness/gotosocial/internal/state"
+	"github.com/superseriousbusiness/gotosocial/internal/util/xslices"
+	"github.com/uptrace/bun"
+)
+
+type webPushDB struct {
+	db    *bun.DB
+	state *state.State
+}
+
+func (w *webPushDB) GetVAPIDKeyPair(ctx context.Context) (*gtsmodel.VAPIDKeyPair, error) {
+	// Look for cached keys.
+	vapidKeyPair := w.state.Caches.DB.VAPIDKeyPair.Load()
+	if vapidKeyPair != nil {
+		return vapidKeyPair, nil
+	}
+
+	// Look for previously generated keys in the database.
+	if err := w.db.NewSelect().
+		Model(vapidKeyPair).
+		Limit(1).
+		Scan(ctx); // nocollapse
+	err != nil && !errors.Is(err, db.ErrNoEntries) {
+		return nil, err
+	}
+
+	// Cache the keys.
+	w.state.Caches.DB.VAPIDKeyPair.Store(vapidKeyPair)
+
+	return vapidKeyPair, nil
+}
+
+func (w *webPushDB) PutVAPIDKeyPair(ctx context.Context, vapidKeyPair *gtsmodel.VAPIDKeyPair) error {
+	// Store the keys in the database.
+	if _, err := w.db.NewInsert().
+		Model(vapidKeyPair).
+		Exec(ctx); // nocollapse
+	err != nil {
+		return err
+	}
+
+	// Cache the keys.
+	w.state.Caches.DB.VAPIDKeyPair.Store(vapidKeyPair)
+
+	return nil
+}
+
+func (w *webPushDB) GetWebPushSubscriptionByTokenID(ctx context.Context, tokenID string) (*gtsmodel.WebPushSubscription, error) {
+	return w.state.Caches.DB.WebPushSubscription.LoadOne(
+		"TokenID",
+		func() (*gtsmodel.WebPushSubscription, error) {
+			var subscription gtsmodel.WebPushSubscription
+			err := w.db.
+				NewSelect().
+				Model(&subscription).
+				Where("? = ?", bun.Ident("token_id"), tokenID).
+				Scan(ctx)
+			return &subscription, err
+		},
+		tokenID,
+	)
+}
+
+func (w *webPushDB) PutWebPushSubscription(ctx context.Context, subscription *gtsmodel.WebPushSubscription) error {
+	return w.state.Caches.DB.WebPushSubscription.Store(subscription, func() error {
+		_, err := w.db.NewInsert().
+			Model(subscription).
+			Exec(ctx)
+		return err
+	})
+}
+
+func (w *webPushDB) UpdateWebPushSubscription(ctx context.Context, subscription *gtsmodel.WebPushSubscription, columns ...string) error {
+	// If we're updating by column, ensure "updated_at" is included.
+	if len(columns) > 0 {
+		columns = append(columns, "updated_at")
+	}
+
+	// Update database.
+	if _, err := w.db.
+		NewUpdate().
+		Model(subscription).
+		Column(columns...).
+		Where("? = ?", bun.Ident("id"), subscription.ID).
+		Exec(ctx); // nocollapse
+	err != nil {
+		return err
+	}
+
+	// Update cache.
+	w.state.Caches.DB.WebPushSubscription.Put(subscription)
+
+	return nil
+}
+
+func (w *webPushDB) DeleteWebPushSubscriptionByTokenID(ctx context.Context, tokenID string) error {
+	// Deleted partial model for cache invalidation.
+	var deleted gtsmodel.WebPushSubscription
+
+	// Delete subscription, returning subset of columns used by invalidation hook.
+	if _, err := w.db.NewDelete().
+		Model(&deleted).
+		Where("? = ?", bun.Ident("token_id"), tokenID).
+		Returning("?", bun.Ident("account_id")).
+		Exec(ctx); // nocollapse
+	err != nil && !errors.Is(err, db.ErrNoEntries) {
+		return err
+	}
+
+	// Invalidate cached subscription by token ID.
+	w.state.Caches.DB.WebPushSubscription.Invalidate("TokenID", tokenID)
+
+	// Call invalidate hook directly.
+	w.state.Caches.OnInvalidateWebPushSubscription(&deleted)
+
+	return nil
+}
+
+func (w *webPushDB) GetWebPushSubscriptionsByAccountID(ctx context.Context, accountID string) ([]*gtsmodel.WebPushSubscription, error) {
+	// Fetch IDs of all subscriptions created by this account.
+	subscriptionIDs, err := loadPagedIDs(&w.state.Caches.DB.WebPushSubscriptionIDs, accountID, nil, func() ([]string, error) {
+		// Subscription IDs not in cache. Perform DB query.
+		var subscriptionIDs []string
+		if _, err := w.db.
+			NewSelect().
+			Model((*gtsmodel.WebPushSubscription)(nil)).
+			Column("id").
+			Where("? = ?", bun.Ident("account_id"), accountID).
+			Order("id DESC").
+			Exec(ctx, &subscriptionIDs); // nocollapse
+		err != nil && !errors.Is(err, db.ErrNoEntries) {
+			return nil, err
+		}
+		return subscriptionIDs, nil
+	})
+	if len(subscriptionIDs) == 0 {
+		return nil, nil
+	}
+
+	// Get each subscription by ID from the cache or DB.
+	subscriptions, err := w.state.Caches.DB.WebPushSubscription.LoadIDs("ID",
+		subscriptionIDs,
+		func(uncached []string) ([]*gtsmodel.WebPushSubscription, error) {
+			subscriptions := make([]*gtsmodel.WebPushSubscription, 0, len(uncached))
+			if err := w.db.
+				NewSelect().
+				Model(&subscriptions).
+				Where("? IN (?)", bun.Ident("id"), bun.In(uncached)).
+				Scan(ctx); // nocollapse
+			err != nil {
+				return nil, err
+			}
+			return subscriptions, nil
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Put the subscription structs in the same order as the filter IDs.
+	xslices.OrderBy(
+		subscriptions,
+		subscriptionIDs,
+		func(subscription *gtsmodel.WebPushSubscription) string {
+			return subscription.ID
+		},
+	)
+
+	return subscriptions, nil
+}
+
+func (w *webPushDB) DeleteWebPushSubscriptionsByAccountID(ctx context.Context, accountID string) error {
+	// Deleted partial models for cache invalidation.
+	var deleted []*gtsmodel.WebPushSubscription
+
+	// Delete subscriptions, returning subset of columns.
+	if _, err := w.db.NewDelete().
+		Model(&deleted).
+		Where("? = ?", bun.Ident("account_id"), accountID).
+		Returning("?", bun.Ident("account_id")).
+		Exec(ctx); // nocollapse
+	err != nil && !errors.Is(err, db.ErrNoEntries) {
+		return err
+	}
+
+	// Invalidate cached subscriptions by account ID.
+	w.state.Caches.DB.WebPushSubscription.Invalidate("AccountID", accountID)
+
+	// Call invalidate hooks directly in case those entries weren't cached.
+	for _, subscription := range deleted {
+		w.state.Caches.OnInvalidateWebPushSubscription(subscription)
+	}
+
+	return nil
+}