[Perf] Use bundled related entities in notifications, drop N+1 fetches

packages/common/src/api/tan-query/notifications/useNotifications.ts

@@ -1,7 +1,11 @@
 import { useEffect, useMemo } from 'react'
 
-import { Id, type NotificationsResponse } from '@audius/sdk'
-import { InfiniteData, useInfiniteQuery } from '@tanstack/react-query'
+import { Id } from '@audius/sdk'
+import {
+  InfiniteData,
+  useInfiniteQuery,
+  useQueryClient
+} from '@tanstack/react-query'
 import { usePrevious } from 'react-use'
 
 import { notificationFromSDK, transformAndCleanList } from '~/adapters'
@@ -11,156 +15,24 @@ import { ChallengeRewardID } from '~/models'
 import { ID } from '~/models/Identifiers'
 import { StringKeys } from '~/services'
 import {
-  Entity,
   NotificationType,
   Notification
 } from '~/store/notifications/types'
 
-import { useCollections } from '../collection/useCollections'
 import { QUERY_KEYS } from '../queryKeys'
-import { useTracks } from '../tracks/useTracks'
 import { QueryKey, QueryOptions } from '../types'
 import { useCurrentUserId } from '../users/account/useCurrentUserId'
-import { useUsers } from '../users/useUsers'
+import { primeRelatedData } from '../utils/primeRelatedData'
 
 import { useNotificationUnreadCount } from './useNotificationUnreadCount'
 
 const DEFAULT_LIMIT = 20
-const USER_INITIAL_LOAD_COUNT = 9
 
 type PageParam = {
   timestamp: number
   groupId: string | undefined
 } | null
 
-type EntityIds = {
-  userIds: ID[]
-  trackIds: ID[]
-  collectionIds: ID[]
-}
-
-const collectEntityIds = (notifications: Notification[]): EntityIds => {
-  const trackIds = new Set<ID>()
-  const collectionIds = new Set<ID>()
-  const userIds = new Set<ID>()
-
-  notifications.forEach((notification) => {
-    const { type } = notification
-    if (type === NotificationType.UserSubscription) {
-      if (notification.entityType === Entity.Track) {
-        if (notification.entityIds.length === 1) {
-          trackIds.add(notification.entityIds[0])
-        }
-      } else if (
-        notification.entityType === Entity.Playlist ||
-        notification.entityType === Entity.Album
-      ) {
-        if (notification.entityIds.length === 1) {
-          collectionIds.add(notification.entityIds[0])
-        }
-      }
-      userIds.add(notification.userId)
-    }
-    if (
-      type === NotificationType.Repost ||
-      type === NotificationType.RepostOfRepost ||
-      type === NotificationType.Favorite ||
-      type === NotificationType.FavoriteOfRepost ||
-      (type === NotificationType.Milestone && 'entityType' in notification)
-    ) {
-      if (notification.entityType === Entity.Track) {
-        trackIds.add(notification.entityId)
-      } else if (
-        notification.entityType === Entity.Playlist ||
-        notification.entityType === Entity.Album
-      ) {
-        collectionIds.add(notification.entityId)
-      } else if (notification.entityType === Entity.User) {
-        userIds.add(notification.entityId)
-      }
-    }
-    if (
-      type === NotificationType.Follow ||
-      type === NotificationType.Repost ||
-      type === NotificationType.RepostOfRepost ||
-      type === NotificationType.Favorite ||
-      type === NotificationType.FavoriteOfRepost
-    ) {
-      notification.userIds
-        .slice(0, USER_INITIAL_LOAD_COUNT)
-        .forEach((id) => userIds.add(id))
-    }
-    if (type === NotificationType.RemixCreate) {
-      trackIds.add(notification.parentTrackId).add(notification.childTrackId)
-    }
-    if (type === NotificationType.RemixCosign) {
-      notification.entityIds.forEach((id) => trackIds.add(id))
-      userIds.add(notification.parentTrackUserId)
-    }
-    if (
-      type === NotificationType.TrendingTrack ||
-      type === NotificationType.TrendingUnderground
-    ) {
-      trackIds.add(notification.entityId)
-    }
-    if (
-      type === NotificationType.AddTrackToPlaylist ||
-      type === NotificationType.TrackAddedToPurchasedAlbum
-    ) {
-      trackIds.add(notification.trackId)
-      userIds.add(notification.playlistOwnerId)
-      collectionIds.add(notification.playlistId)
-    }
-    if (type === NotificationType.Tastemaker) {
-      userIds.add(notification.userId)
-      trackIds.add(notification.entityId)
-    }
-    if (
-      type === NotificationType.USDCPurchaseBuyer ||
-      type === NotificationType.USDCPurchaseSeller
-    ) {
-      notification.userIds.forEach((id) => userIds.add(id))
-      if (notification.entityType === Entity.Track) {
-        trackIds.add(notification.entityId)
-      } else if (notification.entityType === Entity.Album) {
-        collectionIds.add(notification.entityId)
-      }
-    }
-    if (
-      type === NotificationType.RequestManager ||
-      type === NotificationType.ApproveManagerRequest
-    ) {
-      userIds.add(notification.userId)
-    }
-    if (
-      type === NotificationType.Comment ||
-      type === NotificationType.CommentThread ||
-      type === NotificationType.CommentMention ||
-      type === NotificationType.CommentReaction
-    ) {
-      if (notification.entityType === Entity.Track) {
-        trackIds.add(notification.entityId)
-      }
-      notification.userIds
-        .slice(0, USER_INITIAL_LOAD_COUNT)
-        .forEach((id) => userIds.add(id))
-    }
-    if (type === NotificationType.RemixCreate) {
-      trackIds.add(notification.parentTrackId)
-      trackIds.add(notification.childTrackId)
-    }
-    if (type === NotificationType.FanClubTextPost) {
-      userIds.add(notification.entityUserId)
-    }
-  })
-
-  return {
-    userIds: Array.from(userIds),
-    trackIds: Array.from(trackIds),
-    collectionIds: Array.from(collectionIds)
-  }
-}
-
 export const getNotificationsQueryKey = ({
   currentUserId,
   pageSize
@@ -181,6 +53,7 @@ export const getNotificationsQueryKey = ({
  */
 export const useNotifications = (options?: QueryOptions) => {
   const { audiusSdk } = useQueryContext()
+  const queryClient = useQueryClient()
   const { data: currentUserId } = useCurrentUserId()
   const pageSize = DEFAULT_LIMIT
   const { data: unreadCount } = useNotificationUnreadCount()
@@ -205,25 +78,17 @@ export const useNotifications = (options?: QueryOptions) => {
     initialPageParam: null as PageParam,
     queryFn: async ({ pageParam = null }) => {
       const sdk = await audiusSdk()
-      const response = await (
-        sdk.notifications as {
-          getNotifications: (params: {
-            userId: string
-            limit?: number
-            timestamp?: number
-            groupId?: string
-          }) => Promise<NotificationsResponse>
-        }
-      ).getNotifications({
+      const response = await sdk.notifications.getNotifications({
         userId: Id.parse(currentUserId),
         limit: DEFAULT_LIMIT,
         timestamp: pageParam?.timestamp,
         groupId: pageParam?.groupId
       })
-      const data = response
+
+      primeRelatedData({ related: response.related, queryClient })
 
       const notifications = transformAndCleanList(
-        data?.data?.notifications,
+        response?.data?.notifications,
         notificationFromSDK
       ) as Notification[]
 
@@ -263,38 +128,12 @@ export const useNotifications = (options?: QueryOptions) => {
     }
   }, [unreadCount, prevUnreadCount, query])
 
-  const lastPage = query.data?.pages[query.data.pages.length - 1]
-  const { userIds, trackIds, collectionIds } = lastPage
-    ? collectEntityIds(lastPage)
-    : { userIds: undefined, trackIds: undefined, collectionIds: undefined }
-
-  // Pre-fetch related entities
-  const { isPending: isUsersPending } = useUsers(userIds)
-  const { isPending: isTracksPending } = useTracks(trackIds)
-  const { isPending: isCollectionsPending } = useCollections(collectionIds)
-
-  // Return all pages except the last one if it's still loading entity data
-  const notifications = query.data?.pages.slice(0, -1).flat() ?? []
-  if (
-    !query.isPending &&
-    !isUsersPending &&
-    !isTracksPending &&
-    !isCollectionsPending &&
-    lastPage
-  ) {
-    notifications.push(...lastPage)
-  }
+  const notifications = query.data?.pages.flat() ?? []
 
   const queryResults = query as typeof query & {
     notifications: Notification[]
-    isAllPending: boolean
   }
   queryResults.notifications = notifications
-  queryResults.isAllPending =
-    queryResults.isPending ||
-    isUsersPending ||
-    isTracksPending ||
-    isCollectionsPending
 
   return queryResults
 }

packages/common/src/api/tan-query/utils/primeRelatedData.ts

@@ -3,16 +3,19 @@ import { QueryClient } from '@tanstack/react-query'
 
 import {
   transformAndCleanList,
+  userCollectionMetadataFromSDK,
   userTrackMetadataFromSDK,
   userMetadataFromSDK
 } from '~/adapters'
 
+import { primeCollectionData } from './primeCollectionData'
 import { primeTrackData } from './primeTrackData'
 import { primeUserData } from './primeUserData'
 
 /**
  * Utility function to prime related data from API responses
- * This handles both users and tracks that may be included in the related field
+ * This handles users, tracks, and playlists that may be included in the
+ * related field.
  */
 export const primeRelatedData = ({
   related,
@@ -27,9 +30,8 @@ export const primeRelatedData = ({
 }) => {
   if (!related) return
 
-  const { users, tracks } = related
+  const { users, tracks, playlists } = related
 
-  // Prime user data if available
   if (users && users.length > 0) {
     primeUserData({
       users: transformAndCleanList(users, userMetadataFromSDK),
@@ -39,7 +41,6 @@ export const primeRelatedData = ({
     })
   }
 
-  // Prime track data if available
   if (tracks && tracks.length > 0) {
     primeTrackData({
       tracks: transformAndCleanList(tracks, userTrackMetadataFromSDK),
@@ -48,4 +49,16 @@ export const primeRelatedData = ({
       skipQueryData
     })
   }
+
+  if (playlists && playlists.length > 0) {
+    primeCollectionData({
+      collections: transformAndCleanList(
+        playlists,
+        userCollectionMetadataFromSDK
+      ),
+      queryClient,
+      forceReplace,
+      skipQueryData
+    })
+  }
 }

packages/common/src/models/Analytics.ts

@@ -578,6 +578,9 @@ export enum Name {
   REMIX_CONTEST_DELETE = 'Remix Contest: Delete',
   REMIX_CONTEST_PICK_WINNERS_OPEN = 'Remix Contest: Pick Winners Open',
   REMIX_CONTEST_PICK_WINNERS_FINALIZE = 'Remix Contest: Finalize Winners',
+  REMIX_CONTEST_VIEW = 'Remix Contest: View',
+  REMIX_CONTEST_ENTER = 'Remix Contest: Enter',
+  REMIX_CONTEST_VIEW_SUBMISSIONS = 'Remix Contest: View Submissions',
 
   // Android App Lifecycle
   ANDROID_APP_RESTART_HEARTBEAT = 'Android App: Restart Due to Heartbeat',
@@ -2868,6 +2871,24 @@ export type RemixContestPickWinnersFinalize = {
   trackId: ID
 }
 
+export type RemixContestView = {
+  eventName: Name.REMIX_CONTEST_VIEW
+  remixContestId: ID
+  trackId: ID
+}
+
+export type RemixContestEnter = {
+  eventName: Name.REMIX_CONTEST_ENTER
+  remixContestId: ID
+  trackId: ID
+}
+
+export type RemixContestViewSubmissions = {
+  eventName: Name.REMIX_CONTEST_VIEW_SUBMISSIONS
+  remixContestId: ID
+  trackId: ID
+}
+
 export type AndroidAppRestartHeartbeat = {
   eventName: Name.ANDROID_APP_RESTART_HEARTBEAT
   timeSinceLastHeartbeat: number
@@ -3504,6 +3525,9 @@ export type AllTrackingEvents =
   | RemixContestDelete
   | RemixContestPickWinnersOpen
   | RemixContestPickWinnersFinalize
+  | RemixContestView
+  | RemixContestEnter
+  | RemixContestViewSubmissions
   | AndroidAppRestartHeartbeat
   | AndroidAppRestartStale
   | AndroidAppRestartForceQuit