refactor: switch to using service paginator

Author: Hweinstock

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<C, CommandInput extends object, CommandOutput extends object> = (
+    config: Omit<PaginationConfiguration, 'client'> & { client: C },
+    input: CommandInput,
+    ...rest: any[]
+) => Paginator<CommandOutput>
 export abstract class ClientWrapper<C extends AwsClient> implements vscode.Disposable {
     protected client?: C
 
@@ -31,24 +36,21 @@ export abstract class ClientWrapper<C extends AwsClient> 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<C, CommandInput, CommandOutput>,
+        input: CommandInput,
+        extractPage: (page: CommandOutput) => Output[] | undefined
+    ): Promise<Output[]> {
+        const p = paginator({ client: await this.getClient() }, input)
+        const results = []
+        for await (const page of p) {
+            results.push(extractPage(page))
+        }
+        return results.flat().filter((result) => result !== undefined)
     }
 
     public dispose() {

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<SSMClient> {
     }
 
     public async describeInstance(target: string): Promise<InstanceInformation> {
-        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<string> {
@@ -115,6 +110,6 @@ export class SsmClient extends ClientWrapper<SSMClient> {
     }
 
     public async describeSessions(state: SessionState) {
-        return await this.makePaginatedRequest(DescribeSessionsCommand, { State: state }, 'Sessions').promise()
+        return await this.makePaginatedRequest(paginateDescribeSessions, { State: state }, (page) => page.Sessions)
     }
 }