Improve UseCase API Performance (#1057)

Author: toshikwa

packages/cdk/lambda/useCaseBuilder/useCaseBuilderRepository.ts

@@ -17,6 +17,7 @@ import {
   UpdateCommand,
   BatchWriteCommand,
   TransactWriteCommand,
+  QueryCommandOutput,
 } from '@aws-sdk/lib-dynamodb';
 import * as crypto from 'crypto';
 import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
@@ -35,32 +36,30 @@ const getUserIdFromKey = (key: string): string => {
   return key.split('#').slice(1).join('#');
 };
 
+// Create a query command to get use case by useCaseId
+const createFindUseCaseByUseCaseIdCommand = (useCaseId: string) =>
+  new QueryCommand({
+    TableName: USECASE_TABLE_NAME,
+    IndexName: USECASE_ID_INDEX_NAME,
+    KeyConditionExpression:
+      '#useCaseId = :useCaseId and begins_with(#dataType, :dataTypePrefix)',
+    ExpressionAttributeNames: {
+      '#useCaseId': 'useCaseId',
+      '#dataType': 'dataType',
+    },
+    ExpressionAttributeValues: {
+      ':useCaseId': useCaseId,
+      ':dataTypePrefix': 'useCase',
+    },
+  });
+
 // Get use case by useCaseId
 const innerFindUseCaseByUseCaseId = async (
   useCaseId: string
 ): Promise<UseCaseInTable | null> => {
-  const useCaseInTable = await dynamoDbDocument.send(
-    new QueryCommand({
-      TableName: USECASE_TABLE_NAME,
-      IndexName: USECASE_ID_INDEX_NAME,
-      KeyConditionExpression:
-        '#useCaseId = :useCaseId and begins_with(#dataType, :dataTypePrefix)',
-      ExpressionAttributeNames: {
-        '#useCaseId': 'useCaseId',
-        '#dataType': 'dataType',
-      },
-      ExpressionAttributeValues: {
-        ':useCaseId': useCaseId,
-        ':dataTypePrefix': 'useCase',
-      },
-    })
-  );
-
-  if (useCaseInTable.Items && useCaseInTable.Items.length > 0) {
-    return useCaseInTable.Items[0] as UseCaseInTable;
-  } else {
-    return null;
-  }
+  const command = createFindUseCaseByUseCaseIdCommand(useCaseId);
+  const useCaseInTable = await dynamoDbDocument.send(command);
+  return (useCaseInTable.Items?.[0] as UseCaseInTable) || null;
 };
 
 // Get use case list by userId
@@ -104,17 +103,16 @@ const innerFindUseCasesByUserId = async (
 const innerFindUseCasesByUseCaseIds = async (
   useCaseIds: string[]
 ): Promise<UseCaseInTable[]> => {
-  const useCasesInTable: UseCaseInTable[] = [];
-
-  for (const useCaseId of useCaseIds) {
-    const useCaseInTable = await innerFindUseCaseByUseCaseId(useCaseId);
-
-    if (useCaseInTable) {
-      useCasesInTable.push(useCaseInTable);
-    }
-  }
-
-  return useCasesInTable;
+  // Run multiple queries in parallel
+  const useCasesInTable: QueryCommandOutput[] = await Promise.all(
+    useCaseIds.map((useCaseId) =>
+      dynamoDb.send(createFindUseCaseByUseCaseIdCommand(useCaseId))
+    )
+  );
+  return useCasesInTable.flatMap(
+    (useCaseInTable) =>
+      (useCaseInTable.Items?.slice(0, 1) || []) as UseCaseInTable[]
+  );
 };
 
 // Get list of specific data type (favorite, recently used) by userId (all)
@@ -499,18 +497,19 @@ export const listRecentlyUsedUseCases = async (
       exclusiveStartKey
     );
   const useCaseIds = commons.map((c) => c.useCaseId);
-  const useCasesInTable = await innerFindUseCasesByUseCaseIds(useCaseIds);
 
-  const favorites = await innerFindCommonsByUserIdAndDataType(
-    userId,
-    'favorite'
-  );
-  const favoritesUseCaseIds = favorites.map((f) => f.useCaseId);
+  const [useCasesInTable, favorites] = await Promise.all([
+    // List user's use cases
+    innerFindUseCasesByUseCaseIds(useCaseIds),
+    // List user's favorites
+    innerFindCommonsByUserIdAndDataType(userId, 'favorite'),
+  ]);
+  const favoritesUseCaseIds = new Set(favorites.map((f) => f.useCaseId));
 
   const useCasesAsOutput: UseCaseAsOutput[] = useCasesInTable.map((u) => {
     return {
       ...u,
-      isFavorite: favoritesUseCaseIds.includes(u.useCaseId),
+      isFavorite: favoritesUseCaseIds.has(u.useCaseId),
       isMyUseCase: getUserIdFromKey(u.id) === userId,
     };
   });

packages/cdk/lib/construct/use-case-builder.ts

@@ -64,6 +64,7 @@ export class UseCaseBuilder extends Construct {
     // Add UseCaseBuilder related APIs
     const listUseCasesFunction = new NodejsFunction(this, 'ListUseCases', {
       ...commonProperty,
+      memorySize: 512,
       entry: `${commonPath}/listUseCases.ts`,
     });
     useCaseBuilderTable.grantReadData(listUseCasesFunction);
@@ -73,6 +74,7 @@ export class UseCaseBuilder extends Construct {
       'ListFavoriteUseCases',
       {
         ...commonProperty,
+        memorySize: 512,
         entry: `${commonPath}/listFavoriteUseCases.ts`,
         environment: {
           ...commonProperty.environment,
@@ -84,6 +86,7 @@ export class UseCaseBuilder extends Construct {
 
     const getUseCaseFunction = new NodejsFunction(this, 'GetUseCase', {
       ...commonProperty,
+      memorySize: 512,
       entry: `${commonPath}/getUseCase.ts`,
     });
     useCaseBuilderTable.grantReadData(getUseCaseFunction);
@@ -123,6 +126,7 @@ export class UseCaseBuilder extends Construct {
       'ListRecentlyUsedUseCases',
       {
         ...commonProperty,
+        memorySize: 512,
         entry: `${commonPath}/listRecentlyUsedUseCases.ts`,
       }
     );

packages/cdk/test/__snapshots__/generative-ai-use-cases.test.ts.snap

@@ -17635,6 +17635,7 @@ exports[`GenerativeAiUseCases matches the snapshot 5`] = `
           },
         },
         "Handler": "index.handler",
+        "MemorySize": 512,
         "Role": {
           "Fn::GetAtt": [
             "UseCaseBuilderGetUseCaseServiceRoleF32AAC5F",
@@ -17753,6 +17754,7 @@ exports[`GenerativeAiUseCases matches the snapshot 5`] = `
           },
         },
         "Handler": "index.handler",
+        "MemorySize": 512,
         "Role": {
           "Fn::GetAtt": [
             "UseCaseBuilderListFavoriteUseCasesServiceRole5BCFD9A7",
@@ -17871,6 +17873,7 @@ exports[`GenerativeAiUseCases matches the snapshot 5`] = `
           },
         },
         "Handler": "index.handler",
+        "MemorySize": 512,
         "Role": {
           "Fn::GetAtt": [
             "UseCaseBuilderListRecentlyUsedUseCasesServiceRole977F8D35",
@@ -17989,6 +17992,7 @@ exports[`GenerativeAiUseCases matches the snapshot 5`] = `
           },
         },
         "Handler": "index.handler",
+        "MemorySize": 512,
         "Role": {
           "Fn::GetAtt": [
             "UseCaseBuilderListUseCasesServiceRoleDE660444",

packages/types/src/protocolUseCaseBuilder.d.ts

@@ -13,7 +13,7 @@ export type ListRecentlyUsedUseCasesResponse = Pagination<UseCaseAsOutput>;
 export type GetUseCaseResponse = UseCaseAsOutput;
 
 export type CreateUseCaseRequest = UseCaseContent;
-export type CreateUseCaseRespose = UseCaseAsOutput;
+export type CreateUseCaseResponse = UseCaseAsOutput;
 
 export type UpdateUseCaseRequest = UseCaseContent;
 

packages/web/src/hooks/useCaseBuilder/useUseCaseBuilderApi.ts

@@ -1,6 +1,6 @@
 import {
   CreateUseCaseRequest,
-  CreateUseCaseRespose,
+  CreateUseCaseResponse,
   GetUseCaseResponse,
   ListFavoriteUseCasesResponse,
   ListRecentlyUsedUseCasesResponse,
@@ -62,7 +62,7 @@ const useUseCaseBuilderApi = () => {
     },
     createUseCase: async (params: CreateUseCaseRequest) => {
       return http
-        .post<CreateUseCaseRespose, CreateUseCaseRequest>('/usecases', params)
+        .post<CreateUseCaseResponse, CreateUseCaseRequest>('/usecases', params)
         .then((res) => {
           return res.data;
         });