diff --git a/packages/bolt-dynamodb/src/DynamoDbInstallationStore.ts b/packages/bolt-dynamodb/src/DynamoDbInstallationStore.ts index 100e840..ba25e0a 100644 --- a/packages/bolt-dynamodb/src/DynamoDbInstallationStore.ts +++ b/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 { + readonly keyAttributeNames: string[]; extractKeyFrom(item: Record): 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 { return undefined; } - const b = response.Item[this.attributeName].B; + return this.extractInstallation(response.Item); + } + + private extractInstallation( + item: Record + ): 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 | 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 diff --git a/packages/bolt-dynamodb/test/DynamoDbInstallationStore.test.ts b/packages/bolt-dynamodb/test/DynamoDbInstallationStore.test.ts index 5ed117f..163572d 100644 --- a/packages/bolt-dynamodb/test/DynamoDbInstallationStore.test.ts +++ b/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 = { + 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 = { + 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 = { + isEnterpriseInstall: false, + enterpriseId: undefined, + teamId, + userId, + }; + + function setUp( + ...items: Record[] + ): DynamoDbInstallationStore { + const mockedBatchGetItem: ( + args: BatchGetItemCommandInput + ) => Promise = 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); + }); + }); });