diff --git a/packages/core/src/shared/clients/clientWrapper.ts b/packages/core/src/shared/clients/clientWrapper.ts index dbef719f4ee..70ee98a014a 100644 --- a/packages/core/src/shared/clients/clientWrapper.ts +++ b/packages/core/src/shared/clients/clientWrapper.ts @@ -5,8 +5,13 @@ import * as vscode from 'vscode' import globals from '../extensionGlobals' import { AwsClient, AwsClientConstructor, AwsCommand } from '../awsClientBuilderV3' -import { pageableToCollection } from '../utilities/collectionUtils' +import { PaginationConfiguration, Paginator } from '@aws-sdk/types' +type SDKPaginator = ( + config: Omit & { client: C }, + input: CommandInput, + ...rest: any[] +) => Paginator export abstract class ClientWrapper implements vscode.Disposable { protected client?: C @@ -31,24 +36,22 @@ export abstract class ClientWrapper implements vscode.Dispo return await client.send(new command(commandOptions)) } - protected makePaginatedRequest< + protected async makePaginatedRequest< CommandInput extends object, CommandOutput extends object, - Command extends AwsCommand, + Output extends object, >( - command: new (o: CommandInput) => Command, - commandOptions: CommandInput, - collectKey: keyof CommandOutput & string, - nextTokenKey?: keyof CommandOutput & keyof CommandInput & string - ) { - const requester = async (req: CommandInput) => await this.makeRequest(command, req) - const response = pageableToCollection( - requester, - commandOptions, - nextTokenKey ?? ('NextToken' as never), - collectKey - ) - return response + paginator: SDKPaginator, + input: CommandInput, + extractPage: (page: CommandOutput) => Output[] | undefined + ): Promise { + const p = paginator({ client: await this.getClient() }, input) + const results = [] + for await (const page of p) { + results.push(extractPage(page)) + } + const filteredResult = results.flat().filter((result) => result !== undefined) as Output[] + return filteredResult } public dispose() { diff --git a/packages/core/src/shared/clients/ssm.ts b/packages/core/src/shared/clients/ssm.ts index 1ebe8092876..21d172459c3 100644 --- a/packages/core/src/shared/clients/ssm.ts +++ b/packages/core/src/shared/clients/ssm.ts @@ -10,14 +10,14 @@ import { TerminateSessionCommand, TerminateSessionResponse, StartSessionCommandOutput, - DescribeInstanceInformationCommand, DescribeInstanceInformationCommandInput, InstanceInformation, SendCommandCommand, SendCommandCommandOutput, waitUntilCommandExecuted, SessionState, - DescribeSessionsCommand, + paginateDescribeInstanceInformation, + paginateDescribeSessions, } from '@aws-sdk/client-ssm' import { WaiterState } from '@smithy/util-waiter' import { ToolkitError } from '../errors' @@ -52,20 +52,15 @@ export class SsmClient extends ClientWrapper { } public async describeInstance(target: string): Promise { - const response = this.makePaginatedRequest( - DescribeInstanceInformationCommand, - { - InstanceInformationFilterList: [ - { - key: 'InstanceIds', - valueSet: [target], - }, - ], - } as DescribeInstanceInformationCommandInput, - 'InstanceIds' - ) - const resolvedResponse = await response.flatten().flatten().promise() - return resolvedResponse[0]! + return ( + await this.makePaginatedRequest( + paginateDescribeInstanceInformation, + { + InstanceInformationFilterList: [{ key: 'InstanceIds', valueSet: [target] }], + } satisfies DescribeInstanceInformationCommandInput, + (page) => page.InstanceInformationList + ) + )[0]! } public async getTargetPlatformName(target: string): Promise { @@ -115,6 +110,6 @@ export class SsmClient extends ClientWrapper { } public async describeSessions(state: SessionState) { - return await this.makePaginatedRequest(DescribeSessionsCommand, { State: state }, 'Sessions').promise() + return await this.makePaginatedRequest(paginateDescribeSessions, { State: state }, (page) => page.Sessions) } }