Skip to content

Commit 1abe4eb

Browse files
HweinstockHweinstock
committed
fix: re-implement proper pagination
1 parent 38ba68b commit 1abe4eb

File tree

6 files changed

+74
-55
lines changed

6 files changed

+74
-55
lines changed

‎packages/core/src/awsService/ec2/explorer/ec2ParentNode.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { makeChildrenNodes } from '../../../shared/treeview/utils'
88
import { PlaceholderNode } from '../../../shared/treeview/nodes/placeholderNode'
99
import { Ec2InstanceNode } from './ec2InstanceNode'
1010
import { Ec2Client } from '../../../shared/clients/ec2'
11-
import { toMap, updateInPlace } from '../../../shared/utilities/collectionUtils'
11+
import { updateInPlace } from '../../../shared/utilities/collectionUtils'
1212
import { PollingSet } from '../../../shared/utilities/pollingSet'
1313

1414
export const parentContextValue = 'awsEc2ParentNode'
@@ -50,7 +50,7 @@ export class Ec2ParentNode extends AWSTreeNodeBase {
5050
}
5151

5252
public async updateChildren(): Promise<void> {
53-
const ec2Instances = toMap(await this.ec2Client.getInstances(), (instance) => instance.InstanceId)
53+
const ec2Instances = await (await this.ec2Client.getInstances()).toMap((instance) => instance.InstanceId)
5454
updateInPlace(
5555
this.ec2InstanceNodes,
5656
ec2Instances.keys(),

‎packages/core/src/awsService/ec2/prompter.ts‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { CancellationError } from '../../shared/utilities/timeoutUtils'
1111
import { getIconCode } from './utils'
1212
import { Ec2Node } from './explorer/ec2ParentNode'
1313
import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
14+
import { AsyncCollection } from '../../shared/utilities/asyncCollection'
1415

1516
export type instanceFilter = (instance: SafeEc2Instance) => boolean
1617
export interface Ec2Selection {
@@ -52,20 +53,24 @@ export class Ec2Prompter {
5253
}
5354
}
5455

55-
protected async getInstancesFromRegion(regionCode: string): Promise<SafeEc2Instance[]> {
56+
protected async getInstancesFromRegion(regionCode: string): Promise<AsyncCollection<SafeEc2Instance>> {
5657
const client = new Ec2Client(regionCode)
5758
return await client.getInstances()
5859
}
5960

61+
// TODO: implement a batched generator to avoid loading all instances into UI.
6062
protected async getInstancesAsQuickPickItems(region: string): Promise<DataQuickPickItem<string>[]> {
61-
return (await this.getInstancesFromRegion(region))
62-
.filter(this.filter ? (instance) => this.filter!(instance) : (_) => true)
63+
return await (
64+
await this.getInstancesFromRegion(region)
65+
)
66+
.filter(this.filter ? (instance) => this.filter!(instance) : (instance) => true)
6367
.map((instance) => Ec2Prompter.asQuickPickItem(instance))
68+
.promise()
6469
}
6570

6671
private createEc2ConnectPrompter(): RegionSubmenu<string> {
6772
return new RegionSubmenu(
68-
async (region) => this.getInstancesAsQuickPickItems(region),
73+
async (region) => await this.getInstancesAsQuickPickItems(region),
6974
{ title: 'Select EC2 Instance', matchOnDetail: true },
7075
{ title: 'Select Region for EC2 Instance' },
7176
'Instances'

‎packages/core/src/shared/clients/clientWrapper.ts‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import * as vscode from 'vscode'
66
import globals from '../extensionGlobals'
77
import { AwsClient, AwsClientConstructor, AwsCommand } from '../awsClientBuilderV3'
88
import { PaginationConfiguration, Paginator } from '@aws-sdk/types'
9+
import { AsyncCollection, toCollection } from '../utilities/asyncCollection'
910

1011
type SDKPaginator<C, CommandInput extends object, CommandOutput extends object> = (
1112
config: Omit<PaginationConfiguration, 'client'> & { client: C },
@@ -43,15 +44,16 @@ export abstract class ClientWrapper<C extends AwsClient> implements vscode.Dispo
4344
>(
4445
paginator: SDKPaginator<C, CommandInput, CommandOutput>,
4546
input: CommandInput,
46-
extractPage: (page: CommandOutput) => Output[] | undefined
47-
): Promise<Output[]> {
47+
extractPage: (page: CommandOutput) => Output[] | undefined,
48+
promisify: boolean = false
49+
): Promise<AsyncCollection<Output>> {
4850
const p = paginator({ client: await this.getClient() }, input)
49-
const results = []
50-
for await (const page of p) {
51-
results.push(extractPage(page))
52-
}
53-
const filteredResult = results.flat().filter((result) => result !== undefined) as Output[]
54-
return filteredResult
51+
const collection = toCollection(() => p)
52+
.map(extractPage)
53+
.flatten()
54+
.filter((i) => i !== undefined)
55+
56+
return collection
5557
}
5658

5759
public dispose() {

‎packages/core/src/shared/clients/ec2.ts‎

Lines changed: 29 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import { showMessageWithCancel } from '../utilities/messages'
2828
import { ToolkitError, isAwsError } from '../errors'
2929
import { decodeBase64 } from '../utilities/textUtilities'
3030
import { ClientWrapper } from './clientWrapper'
31+
import { AsyncCollection } from '../utilities/asyncCollection'
3132

3233
/**
3334
* A wrapper around EC2.Instance where we can safely assume InstanceId field exists.
@@ -38,6 +39,10 @@ export interface SafeEc2Instance extends Instance {
3839
LastSeenStatus: InstanceStateName
3940
}
4041

42+
interface InstanceWithId extends Instance {
43+
InstanceId: string
44+
}
45+
4146
interface SafeEc2GetConsoleOutputResult extends GetConsoleOutputRequest {
4247
Output: string
4348
InstanceId: string
@@ -48,38 +53,38 @@ export class Ec2Client extends ClientWrapper<EC2Client> {
4853
super(regionCode, EC2Client)
4954
}
5055

51-
public async getInstances(filters?: Filter[]): Promise<SafeEc2Instance[]> {
56+
public async getInstances(filters?: Filter[]): Promise<AsyncCollection<SafeEc2Instance>> {
5257
const reservations = await this.makePaginatedRequest(
5358
paginateDescribeInstances,
5459
filters ? { Filters: filters } : ({} satisfies DescribeInstancesRequest),
5560
(page) => page.Reservations
5661
)
5762

58-
return await this.updateInstancesDetail(this.getInstancesFromReservations(reservations))
63+
return this.updateInstancesDetail(this.getInstancesFromReservations(reservations))
5964
}
6065

6166
/** Updates status and name in-place for displaying to humans. */
62-
public async updateInstancesDetail(
63-
instances: Instance[],
67+
public updateInstancesDetail(
68+
instances: AsyncCollection<Instance>,
6469
getStatus: (i: string) => Promise<InstanceStateName> = this.getInstanceStatus.bind(this)
65-
): Promise<SafeEc2Instance[]> {
70+
): AsyncCollection<SafeEc2Instance> {
6671
const instanceWithId = instances.filter(hasId)
67-
const instanceWithStatus = await Promise.all(instanceWithId.map(addStatus))
72+
const instanceWithStatus = instanceWithId.map(addStatus)
6873
return instanceWithStatus.map((i) => (instanceHasName(i) ? { ...i, Name: lookupTagKey(i.Tags, 'Name') } : i))
6974

70-
function hasId(i: Instance): i is Instance & { InstanceId: string } {
75+
function hasId(i: Instance): i is InstanceWithId {
7176
return i.InstanceId !== undefined
7277
}
7378

74-
async function addStatus(instance: Instance & { InstanceId: string }) {
79+
async function addStatus(instance: InstanceWithId) {
7580
return { ...instance, LastSeenStatus: await getStatus(instance.InstanceId) }
7681
}
7782
}
7883

79-
public getInstancesFromReservations(reservations: Reservation[]): (Instance & { InstanceId: string })[] {
84+
public getInstancesFromReservations(reservations: AsyncCollection<Reservation>): AsyncCollection<InstanceWithId> {
8085
return reservations
8186
.map((r) => r.Instances)
82-
.flat()
87+
.flatten()
8388
.filter(isNotEmpty)
8489

8590
function isNotEmpty(i: Instance | undefined): i is Instance & { InstanceId: string } {
@@ -88,11 +93,13 @@ export class Ec2Client extends ClientWrapper<EC2Client> {
8893
}
8994

9095
public async getInstanceStatus(instanceId: string): Promise<InstanceStateName> {
91-
const instanceStatuses = await this.makePaginatedRequest(
92-
paginateDescribeInstanceStatus,
93-
{ InstanceIds: [instanceId], IncludeAllInstances: true },
94-
(page) => page.InstanceStatuses
95-
)
96+
const instanceStatuses = await (
97+
await this.makePaginatedRequest(
98+
paginateDescribeInstanceStatus,
99+
{ InstanceIds: [instanceId], IncludeAllInstances: true },
100+
(page) => page.InstanceStatuses
101+
)
102+
).promise()
96103

97104
return instanceStatuses[0].InstanceState!.Name!
98105
}
@@ -198,11 +205,13 @@ export class Ec2Client extends ClientWrapper<EC2Client> {
198205
private async getIamInstanceProfileAssociation(instanceId: string): Promise<IamInstanceProfileAssociation> {
199206
const instanceFilter = this.getInstancesFilter([instanceId])
200207

201-
const associations = await this.makePaginatedRequest(
202-
paginateDescribeIamInstanceProfileAssociations,
203-
{ Filters: instanceFilter },
204-
(page) => page.IamInstanceProfileAssociations
205-
)
208+
const associations = await (
209+
await this.makePaginatedRequest(
210+
paginateDescribeIamInstanceProfileAssociations,
211+
{ Filters: instanceFilter },
212+
(page) => page.IamInstanceProfileAssociations
213+
)
214+
).promise()
206215

207216
return associations[0]!
208217
}

‎packages/core/src/test/awsService/ec2/explorer/ec2ParentNode.test.ts‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ import { Ec2InstanceNode } from '../../../../awsService/ec2/explorer/ec2Instance
1515
import * as FakeTimers from '@sinonjs/fake-timers'
1616
import { installFakeClock } from '../../../testUtil'
1717
import { Filter } from '@aws-sdk/client-ec2'
18+
import { AsyncCollection } from '../../../../shared/utilities/asyncCollection'
19+
import { intoCollection } from '../../../../shared/utilities/collectionUtils'
1820

1921
export const testInstance = {
2022
InstanceId: 'testId',
@@ -32,7 +34,7 @@ export const testParentNode = new Ec2ParentNode('fake-region', 'testPartition',
3234
describe('ec2ParentNode', function () {
3335
let testNode: Ec2ParentNode
3436
let client: Ec2Client
35-
let getInstanceStub: sinon.SinonStub<[filters?: Filter[] | undefined], Promise<SafeEc2Instance[]>>
37+
let getInstanceStub: sinon.SinonStub<[filters?: Filter[] | undefined], Promise<AsyncCollection<SafeEc2Instance>>>
3638
let clock: FakeTimers.InstalledClock
3739
let refreshStub: sinon.SinonStub<[], Promise<void>>
3840
let statusUpdateStub: sinon.SinonStub<[status: string], Promise<string>>
@@ -64,7 +66,7 @@ describe('ec2ParentNode', function () {
6466
})
6567

6668
it('returns placeholder node if no children are present', async function () {
67-
getInstanceStub.resolves([])
69+
getInstanceStub.resolves(intoCollection([]))
6870

6971
const childNodes = await testNode.getChildren()
7072
assertNodeListOnlyHasPlaceholderNode(childNodes)
@@ -76,7 +78,7 @@ describe('ec2ParentNode', function () {
7678
{ Name: 'firstOne', InstanceId: '0', LastSeenStatus: 'running' },
7779
{ Name: 'secondOne', InstanceId: '1', LastSeenStatus: 'stopped' },
7880
] satisfies SafeEc2Instance[]
79-
getInstanceStub.resolves(instances)
81+
getInstanceStub.resolves(intoCollection(instances))
8082
const childNodes = await testNode.getChildren()
8183

8284
assert.strictEqual(childNodes.length, instances.length, 'Unexpected child count')
@@ -99,7 +101,7 @@ describe('ec2ParentNode', function () {
99101
{ Name: 'cd', InstanceId: '5', LastSeenStatus: 'running' },
100102
] satisfies SafeEc2Instance[]
101103

102-
getInstanceStub.resolves(instances)
104+
getInstanceStub.resolves(intoCollection(instances))
103105

104106
const childNodes = await testNode.getChildren()
105107

@@ -122,7 +124,7 @@ describe('ec2ParentNode', function () {
122124
{ Name: 'firstOne', InstanceId: '2', LastSeenStatus: 'running' },
123125
] satisfies SafeEc2Instance[]
124126

125-
getInstanceStub.resolves(instances)
127+
getInstanceStub.resolves(intoCollection(instances))
126128

127129
const childNodes = await testNode.getChildren()
128130
assert.strictEqual(childNodes.length, instances.length, 'Unexpected child count')
@@ -136,7 +138,7 @@ describe('ec2ParentNode', function () {
136138
{ Name: 'thirdOne', InstanceId: '2', LastSeenStatus: 'running' },
137139
] satisfies SafeEc2Instance[]
138140

139-
getInstanceStub.resolves(instances)
141+
getInstanceStub.resolves(intoCollection(instances))
140142
await testNode.updateChildren()
141143
assert.strictEqual(testNode.pollingSet.size, 1)
142144
getInstanceStub.restore()
@@ -150,7 +152,7 @@ describe('ec2ParentNode', function () {
150152
{ Name: 'thirdOne', InstanceId: '2', LastSeenStatus: 'running' },
151153
] satisfies SafeEc2Instance[]
152154

153-
getInstanceStub.resolves(instances)
155+
getInstanceStub.resolves(intoCollection(instances))
154156

155157
await testNode.updateChildren()
156158
await clock.tickAsync(6000)
@@ -162,7 +164,7 @@ describe('ec2ParentNode', function () {
162164
statusUpdateStub = statusUpdateStub.resolves('running')
163165
const instances = [{ Name: 'firstOne', InstanceId: '0', LastSeenStatus: 'pending' }] satisfies SafeEc2Instance[]
164166

165-
getInstanceStub.resolves(instances)
167+
getInstanceStub.resolves(intoCollection(instances))
166168
await testNode.updateChildren()
167169

168170
sinon.assert.notCalled(refreshStub)
@@ -175,7 +177,7 @@ describe('ec2ParentNode', function () {
175177
{ Name: 'firstOne', InstanceId: 'node1', LastSeenStatus: 'pending' },
176178
] satisfies SafeEc2Instance[]
177179

178-
getInstanceStub.resolves(instances)
180+
getInstanceStub.resolves(intoCollection(instances))
179181
await testNode.updateChildren()
180182
const node = testNode.getInstanceNode('node1')
181183
assert.strictEqual(node.InstanceId, instances[0].InstanceId)
@@ -187,7 +189,7 @@ describe('ec2ParentNode', function () {
187189
{ Name: 'firstOne', InstanceId: 'node1', LastSeenStatus: 'pending' },
188190
] satisfies SafeEc2Instance[]
189191

190-
getInstanceStub.resolves(instances)
192+
getInstanceStub.resolves(intoCollection(instances))
191193
await testNode.updateChildren()
192194
assert.throws(() => testNode.getInstanceNode('node2'))
193195
getInstanceStub.restore()
@@ -198,7 +200,7 @@ describe('ec2ParentNode', function () {
198200
{ Name: 'firstOne', InstanceId: 'node1', LastSeenStatus: 'pending' },
199201
] satisfies SafeEc2Instance[]
200202

201-
getInstanceStub.resolves(instances)
203+
getInstanceStub.resolves(intoCollection(instances))
202204
await testNode.updateChildren()
203205
testNode.trackPendingNode('node1')
204206
assert.strictEqual(testNode.pollingSet.size, 1)
@@ -210,7 +212,7 @@ describe('ec2ParentNode', function () {
210212
{ Name: 'firstOne', InstanceId: 'node1', LastSeenStatus: 'pending' },
211213
] satisfies SafeEc2Instance[]
212214

213-
getInstanceStub.resolves(instances)
215+
getInstanceStub.resolves(intoCollection(instances))
214216
await testNode.updateChildren()
215217
assert.throws(() => testNode.trackPendingNode('node2'))
216218
getInstanceStub.restore()

‎packages/core/src/test/shared/clients/ec2Client.test.ts‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import assert from 'assert'
77
import { Ec2Client, instanceHasName } from '../../../shared/clients/ec2'
88
import { Filter, Instance, InstanceStateName, Reservation } from '@aws-sdk/client-ec2'
9+
import { intoCollection } from '../../../shared/utilities/collectionUtils'
910

1011
const completeReservationsList: Reservation[] = [
1112
{
@@ -78,19 +79,19 @@ const getStatus: (i: string) => Promise<InstanceStateName> = (i) =>
7879
describe('extractInstancesFromReservations', function () {
7980
const client = new Ec2Client('')
8081

81-
it('returns empty when given empty collection', function () {
82-
const actualResult = client.getInstancesFromReservations([])
82+
it('returns empty when given empty collection', async function () {
83+
const actualResult = await client.getInstancesFromReservations(intoCollection([])).promise()
8384

8485
assert.strictEqual(0, actualResult.length)
8586
})
8687

87-
it('flattens the reservationList', function () {
88-
const actualResult = client.getInstancesFromReservations(completeReservationsList)
88+
it('flattens the reservationList', async function () {
89+
const actualResult = client.getInstancesFromReservations(intoCollection(completeReservationsList))
8990
assert.deepStrictEqual(actualResult, completeInstanceList)
9091
})
9192

9293
it('handles undefined and missing pieces in the ReservationList.', async function () {
93-
const actualResult = client.getInstancesFromReservations(incompleteReservationsList)
94+
const actualResult = client.getInstancesFromReservations(intoCollection(incompleteReservationsList))
9495
assert.deepStrictEqual(actualResult, incomepleteInstanceList)
9596
})
9697
})
@@ -101,8 +102,8 @@ describe('updateInstancesDetail', async function () {
101102
client = new Ec2Client('test-region')
102103
})
103104

104-
it('adds appropriate status and name field to the instance', async function () {
105-
const actualResult = await client.updateInstancesDetail(completeInstanceList, getStatus)
105+
it('adds appropriate status and name field to the instance', function () {
106+
const actualResult = client.updateInstancesDetail(intoCollection(completeInstanceList), getStatus)
106107
const expectedResult = [
107108
{
108109
InstanceId: 'running-1',
@@ -133,8 +134,8 @@ describe('updateInstancesDetail', async function () {
133134
assert.deepStrictEqual(actualResult, expectedResult)
134135
})
135136

136-
it('handles incomplete and missing tag fields', async function () {
137-
const actualResult = await client.updateInstancesDetail(incomepleteInstanceList, getStatus)
137+
it('handles incomplete and missing tag fields', function () {
138+
const actualResult = client.updateInstancesDetail(intoCollection(incomepleteInstanceList), getStatus)
138139

139140
const expectedResult = [
140141
{ InstanceId: 'running-1', LastSeenStatus: 'running' },

0 commit comments

Comments
 (0)