Add batching for events by entity ID (#11972)

Author: Isaac Solo

packages/common/src/api/tan-query/batchers/getEventsByEntityIdBatcher.ts

@@ -0,0 +1,29 @@
+import { Id, OptionalId } from '@audius/sdk'
+import { create, windowScheduler } from '@yornaath/batshit'
+import { memoize } from 'lodash'
+
+import { eventMetadataListFromSDK } from '~/adapters/event'
+import { ID, Event } from '~/models'
+
+import { contextCacheResolver } from './contextCacheResolver'
+import { BatchContext } from './types'
+
+export const getEventsByEntityIdBatcher = memoize(
+  (context: BatchContext) =>
+    create({
+      fetcher: async (entityIds: ID[]): Promise<Event[]> => {
+        const { sdk, currentUserId } = context
+        if (!entityIds.length) return []
+        const { data } = await sdk.events.getEntityEvents({
+          entityId: entityIds.map((entityId) => Id.parse(entityId)),
+          userId: OptionalId.parse(currentUserId)
+        })
+
+        return eventMetadataListFromSDK(data)
+      },
+      resolver: (events: Event[], entityId: ID) =>
+        events.filter((event) => event.entityId === entityId), // resolve array of events for entity ID
+      scheduler: windowScheduler(10)
+    }),
+  contextCacheResolver()
+)

packages/common/src/api/tan-query/events/useEventsByEntityId.ts

@@ -1,28 +1,27 @@
 import { useMemo } from 'react'
 
-import { Id, OptionalId } from '@audius/sdk'
 import { useQuery, useQueryClient } from '@tanstack/react-query'
+import { useDispatch } from 'react-redux'
 
-import { eventMetadataFromSDK } from '~/adapters/event'
 import { useAudiusQueryContext } from '~/audius-query'
 import { ID } from '~/models'
-import { removeNullable } from '~/utils'
 
+import { getEventsByEntityIdBatcher } from '../batchers/getEventsByEntityIdBatcher'
 import { SelectableQueryOptions } from '../types'
 import { useCurrentUserId } from '../useCurrentUserId'
 
 import {
   getEventIdsByEntityIdQueryKey,
-  EventIdsByEntityIdOptions,
-  getEventQueryKey
+  EventIdsByEntityIdOptions
 } from './utils'
 
 export const useEventIdsByEntityId = (
   args: EventIdsByEntityIdOptions,
   options?: SelectableQueryOptions<ID[]>
 ) => {
-  const { entityId, ...restArgs } = args ?? {}
+  const { entityId } = args ?? {}
   const { audiusSdk } = useAudiusQueryContext()
+  const dispatch = useDispatch()
   const { data: currentUserId } = useCurrentUserId()
   const queryClient = useQueryClient()
 
@@ -33,21 +32,17 @@ export const useEventIdsByEntityId = (
     queryKey: getEventIdsByEntityIdQueryKey(args),
     queryFn: async () => {
       const sdk = await audiusSdk()
-      const response = await sdk.events.getEntityEvents({
-        entityId: Id.parse(entityId),
-        userId: OptionalId.parse(currentUserId),
-        ...restArgs
-      })
-      const events = response.data ?? []
-      const eventsMetadata = events
-        .map(eventMetadataFromSDK)
-        .filter(removeNullable)
 
-      eventsMetadata.forEach((event) => {
-        queryClient.setQueryData(getEventQueryKey(event.eventId), event)
+      const batchGetEvents = getEventsByEntityIdBatcher({
+        sdk,
+        currentUserId,
+        queryClient,
+        dispatch
       })
 
-      return eventsMetadata.map((event) => event.eventId)
+      const events = await batchGetEvents.fetch(entityId!)
+
+      return events.map((event) => event.eventId)
     },
     ...options,
     enabled: options?.enabled !== false && !!entityId,

packages/discovery-provider/src/api/v1/events.py

@@ -4,7 +4,6 @@
     abort_not_found,
     current_user_parser,
     decode_ids_array,
-    decode_with_abort,
     format_limit,
     format_offset,
     make_response,
@@ -120,6 +119,7 @@ def get(self):
     "entity_id",
     required=True,
     type=str,
+    action="append",
     description="The ID of the entity to get events for",
 )
 entity_events_parser.add_argument(
@@ -139,7 +139,7 @@ def get(self):
 
 
 @ns.route("/entity")
-class EntityEvents(Resource):
+class BulkEntityEvents(Resource):
     @ns.doc(
         id="Get Entity Events",
         description="Get events for a specific entity",
@@ -151,10 +151,13 @@ class EntityEvents(Resource):
     def get(self):
         """Get events for a specific entity"""
         args = entity_events_parser.parse_args()
-        decoded_entity_id = decode_with_abort(args.get("entity_id"), ns)
+        decoded_entity_ids = decode_ids_array(
+            args.get("entity_id") if args.get("entity_id") else []
+        )
+
         events = get_events(
             {
-                "entity_id": decoded_entity_id,
+                "entity_ids": decoded_entity_ids,
                 "entity_type": args.get("entity_type"),
                 "filter_deleted": args.get("filter_deleted", True),
                 "limit": format_limit(args),

packages/discovery-provider/src/queries/get_events.py

@@ -64,8 +64,8 @@ def get_events_by_ids(args) -> List[Event]:
             query = query.filter(Event.event_id.in_(ids))
         if args.get("event_type") is not None:
             query = query.filter(Event.event_type == args.get("event_type"))
-        events = query.all()
 
+        events = helpers.query_result_to_list(query.all())
         return format_events(events)
 
 
@@ -74,9 +74,9 @@ def _get_events(session, args):
     # Create initial query
     base_query = session.query(Event)
 
-    # Filter by entity_id if provided
-    if args.get("entity_id") is not None:
-        base_query = base_query.filter(Event.entity_id == args.get("entity_id"))
+    # Filter by entity_ids if provided
+    if args.get("entity_ids") is not None:
+        base_query = base_query.filter(Event.entity_id.in_(args.get("entity_ids")))
 
     # Filter by entity_type if provided
     if args.get("entity_type") is not None:

packages/sdk/src/sdk/api/generated/default/apis/EventsApi.ts

@@ -41,7 +41,7 @@ export interface GetBulkEventsRequest {
 }
 
 export interface GetEntityEventsRequest {
-    entityId: string;
+    entityId: Array<string>;
     offset?: number;
     limit?: number;
     userId?: string;
@@ -166,7 +166,7 @@ export class EventsApi extends runtime.BaseAPI {
             queryParameters['user_id'] = params.userId;
         }
 
-        if (params.entityId !== undefined) {
+        if (params.entityId) {
             queryParameters['entity_id'] = params.entityId;
         }