Use BatchGetItem action to get multiple installations

packages/bolt-dynamodb/src/DynamoDbInstallationStore.ts

@@ -1,6 +1,7 @@
 import {Logger} from '@slack/logger';
 import {
   AttributeValue,
+  BatchGetItemCommandInput,
   BatchWriteItemCommandInput,
   DynamoDB,
   GetItemCommandInput,
@@ -12,7 +13,6 @@ import {
   InstallationStoreBase,
   KeyGenerator,
   KeyGeneratorArgs,
-  Storage,
   StorageBase,
 } from './InstallationStoreBase';
 
@@ -30,14 +30,20 @@ type DeletionOption = 'DELETE_ITEM' | 'DELETE_ATTRIBUTE';
 
 export interface DynamoDbKeyGenerator
   extends KeyGenerator<DynamoDbKey, DynamoDbDeletionKey> {
+  readonly keyAttributeNames: string[];
   extractKeyFrom(item: Record<string, AttributeValue>): DynamoDbKey;
+  equals(key1: DynamoDbKey, key2: DynamoDbKey): boolean;
 }
 
 export class SimpleKeyGenerator implements DynamoDbKeyGenerator {
+  readonly keyAttributeNames: string[];
+
   private constructor(
     private readonly partitionKeyName: string,
     private readonly sortKeyName: string
-  ) {}
+  ) {
+    this.keyAttributeNames = [partitionKeyName, sortKeyName];
+  }
 
   static create(
     partitionKeyName = 'PK',
@@ -78,6 +84,18 @@ export class SimpleKeyGenerator implements DynamoDbKeyGenerator {
     ]);
   }
 
+  equals(key1: DynamoDbKey, key2: DynamoDbKey): boolean {
+    const pk1 = key1[this.partitionKeyName]?.S;
+    const pk2 = key2[this.partitionKeyName]?.S;
+    if (pk1 === undefined || pk2 === undefined || pk1 !== pk2) {
+      return false;
+    }
+
+    const sk1 = key1[this.sortKeyName]?.S;
+    const sk2 = key2[this.sortKeyName]?.S;
+    return !(sk1 === undefined || sk2 === undefined || sk1 !== sk2);
+  }
+
   private generatePartitionKey(args: KeyGeneratorArgs): string {
     return [
       `Client#${args.clientId}`,
@@ -195,12 +213,90 @@ class DynamoDbStorage extends StorageBase<DynamoDbKey, DynamoDbDeletionKey> {
       return undefined;
     }
 
-    const b = response.Item[this.attributeName].B;
+    return this.extractInstallation(response.Item);
+  }
+
+  private extractInstallation(
+    item: Record<string, AttributeValue>
+  ): Buffer | undefined {
+    const b = item[this.attributeName]?.B;
     return b ? Buffer.from(b) : undefined;
   }
 
   // ---
 
+  async fetchMultiple(
+    keys: DynamoDbKey[],
+    logger: Logger | undefined
+  ): Promise<(Buffer | undefined)[]> {
+    if (keys.length === 1) {
+      return [await this.fetch(keys[0], logger)];
+    }
+
+    const entries = this.keyGenerator.keyAttributeNames.map(
+      (attrName, index) => {
+        return [`#key${index}`, attrName];
+      }
+    );
+
+    const input: BatchGetItemCommandInput = {
+      RequestItems: Object.fromEntries([
+        [
+          this.tableName,
+          {
+            Keys: keys,
+            ProjectionExpression: `#inst, ${entries
+              .map(([expAttrName]) => expAttrName)
+              .join(', ')}`,
+            ExpressionAttributeNames: {
+              '#inst': this.attributeName,
+              ...Object.fromEntries(entries),
+            },
+          },
+        ],
+      ]),
+      ReturnConsumedCapacity: 'TOTAL',
+    };
+
+    const response = await this.client.batchGetItem(input);
+    logger?.debug(
+      '[fetchMultiple] BatchGetItem consumed capacity',
+      response.ConsumedCapacity
+    );
+
+    if (
+      response.Responses === undefined ||
+      response.Responses[this.tableName] === undefined
+    ) {
+      return [];
+    }
+    const items = response.Responses[this.tableName];
+
+    const result: (Buffer | undefined)[] = [];
+
+    for (const key of keys) {
+      let found: Record<string, AttributeValue> | undefined;
+      for (const item of items) {
+        const keyFromItem = this.keyGenerator.extractKeyFrom(item);
+        if (this.keyGenerator.equals(keyFromItem, key)) {
+          found = item;
+          break;
+        }
+      }
+
+      if (found) {
+        result.push(this.extractInstallation(found));
+      } else {
+        logger?.debug('Item not found', key);
+        result.push(undefined);
+      }
+    }
+
+    return result;
+  }
+
+  // ---
+
   async delete(
     key: DynamoDbDeletionKey,
     logger: Logger | undefined

packages/bolt-dynamodb/test/DynamoDbInstallationStore.test.ts

@@ -2,11 +2,16 @@ import {ConsoleLogger, LogLevel} from '@slack/logger';
 import {Installation, InstallationQuery} from '@slack/oauth';
 import {
   AttributeValue,
+  BatchGetItemCommandInput,
+  BatchGetItemCommandOutput,
   DynamoDB,
   PutItemCommandInput,
   ScanCommandOutput,
 } from '@aws-sdk/client-dynamodb';
-import {DynamoDbInstallationStore} from '../src/DynamoDbInstallationStore';
+import {
+  DynamoDbInstallationStore,
+  SimpleKeyGenerator,
+} from '../src/DynamoDbInstallationStore';
 import {generateTestData, TestInstallation} from './test-data';
 
 const logger = new ConsoleLogger();
@@ -610,4 +615,154 @@ describe('DynamoDbInstallationStore', () => {
       });
     });
   });
+
+  describe('Behavior of the fetchMultiple', () => {
+    const testData = generateTestData();
+    const userInstallation = testData.installation.teamA.userA1;
+    const botLatestInstallation = testData.installation.teamA.userA2;
+
+    const clientId = 'bolt-dynamodb-test';
+    const tableName = 'FetchMultipleTestTable';
+    const teamId = userInstallation.team.id;
+    const userId = userInstallation.user.id;
+
+    const itemOfUser: Record<string, AttributeValue> = {
+      PK: {S: `Client#${clientId}$Enterprise#none$Team#${teamId}`},
+      SK: {S: `Type#Token$User#${userId}$Version#latest`},
+      Installation: {B: Buffer.from(JSON.stringify(userInstallation))},
+    };
+    const itemOfBot: Record<string, AttributeValue> = {
+      PK: {S: `Client#${clientId}$Enterprise#none$Team#${teamId}`},
+      SK: {S: 'Type#Token$User#___bot___$Version#latest'},
+      Installation: {B: Buffer.from(JSON.stringify(botLatestInstallation))},
+    };
+
+    const fetchQuery: InstallationQuery<false> = {
+      isEnterpriseInstall: false,
+      enterpriseId: undefined,
+      teamId,
+      userId,
+    };
+
+    function setUp(
+      ...items: Record<string, AttributeValue>[]
+    ): DynamoDbInstallationStore {
+      const mockedBatchGetItem: (
+        args: BatchGetItemCommandInput
+      ) => Promise<BatchGetItemCommandOutput> = jest.fn(() => {
+        return new Promise(resolve =>
+          resolve({
+            Responses: Object.fromEntries([[tableName, items]]),
+          } as BatchGetItemCommandOutput)
+        );
+      });
+
+      const mockedDynamoDbClient = {
+        batchGetItem: mockedBatchGetItem,
+      } as DynamoDB;
+
+      return DynamoDbInstallationStore.create({
+        clientId: 'bolt-dynamodb-test',
+        dynamoDb: mockedDynamoDbClient,
+        tableName: 'FetchMultipleTestTable',
+        partitionKeyName: 'PK',
+        sortKeyName: 'SK',
+        attributeName: 'Installation',
+      });
+    }
+
+    test.each([[[itemOfUser, itemOfBot]], [[itemOfBot, itemOfUser]]])(
+      'fetchInstallation() handles BatchGetItem response regardless of item order',
+      async items => {
+        // arrange
+        const sut = setUp(...items);
+
+        // act
+        const result = await sut.fetchInstallation(fetchQuery);
+
+        // assert
+        expect(result.user.id).toEqual(userId);
+      }
+    );
+  });
+});
+
+describe('SimpleKeyGenerator', () => {
+  const sut = SimpleKeyGenerator.create('PK', 'SK');
+
+  describe('equals()', () => {
+    it('returns true when two objects are equal', () => {
+      const result = sut.equals(
+        {
+          PK: {S: 'PartitionKey-0'},
+          SK: {S: 'SortKey-0'},
+        },
+        {
+          PK: {S: 'PartitionKey-0'},
+          SK: {S: 'SortKey-0'},
+        }
+      );
+
+      expect(result).toEqual(true);
+    });
+
+    it('returns false when partition keys of two objects are not equal', () => {
+      const result = sut.equals(
+        {
+          PK: {S: 'PartitionKey-0'},
+          SK: {S: 'SortKey-0'},
+        },
+        {
+          PK: {S: 'PartitionKey-99999'},
+          SK: {S: 'SortKey-0'},
+        }
+      );
+      expect(result).toEqual(false);
+    });
+
+    it('return false when sort keys of two objects are not equal', () => {
+      const result = sut.equals(
+        {
+          PK: {S: 'PartitionKey-0'},
+          SK: {S: 'SortKey-0'},
+        },
+        {
+          PK: {S: 'PartitionKey-0'},
+          SK: {S: 'SortKey-99999'},
+        }
+      );
+
+      expect(result).toEqual(false);
+    });
+
+    it('returns false when two objects are equal but the data type of the partition key is not string', () => {
+      const result = sut.equals(
+        {
+          PK: {N: '0'},
+          SK: {S: 'SortKey-0'},
+        },
+        {
+          PK: {N: '0'},
+          SK: {S: 'SortKey-0'},
+        }
+      );
+
+      expect(result).toEqual(false);
+    });
+
+    it('returns false when two objects are equal but the data type of the sort key is not string', () => {
+      const result = sut.equals(
+        {
+          PK: {S: 'PartitionKey-0'},
+          SK: {N: '0'},
+        },
+        {
+          PK: {S: 'PartitionKey-0'},
+          SK: {N: '0'},
+        }
+      );
+
+      expect(result).toEqual(false);
+    });
+  });
 });