Skip to content

Commit 15a37f5

Browse files
committed
refactor: migrate ec2 usage to v3
1 parent 5f8ed1e commit 15a37f5

File tree

18 files changed

+337
-377
lines changed

18 files changed

+337
-377
lines changed

packages/core/src/awsService/ec2/commands.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*/
55
import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
66
import { Ec2Node } from './explorer/ec2ParentNode'
7-
import { SafeEc2Instance, Ec2Client } from '../../shared/clients/ec2Client'
7+
import { SafeEc2Instance, Ec2Client } from '../../shared/clients/ec2'
88
import { copyToClipboard } from '../../shared/utilities/messages'
99
import { ec2LogSchema } from './ec2LogDocumentProvider'
1010
import { getAwsConsoleUrl } from '../../shared/awsConsole'

packages/core/src/awsService/ec2/ec2LogDocumentProvider.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*/
55
import * as vscode from 'vscode'
66
import { Ec2Selection } from './prompter'
7-
import { Ec2Client } from '../../shared/clients/ec2Client'
7+
import { Ec2Client } from '../../shared/clients/ec2'
88
import { ec2LogsScheme } from '../../shared/constants'
99
import { UriSchema } from '../../shared/utilities/uriUtils'
1010

packages/core/src/awsService/ec2/explorer/ec2InstanceNode.ts

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,17 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
import * as vscode from 'vscode'
6-
import { Ec2Client, getNameOfInstance } from '../../../shared/clients/ec2Client'
6+
import { Ec2Client, getNameOfInstance } from '../../../shared/clients/ec2'
77
import { AWSResourceNode } from '../../../shared/treeview/nodes/awsResourceNode'
88
import { AWSTreeNodeBase } from '../../../shared/treeview/nodes/awsTreeNodeBase'
9-
import { SafeEc2Instance } from '../../../shared/clients/ec2Client'
9+
import { SafeEc2Instance } from '../../../shared/clients/ec2'
1010
import globals from '../../../shared/extensionGlobals'
1111
import { getIconCode } from '../utils'
1212
import { Ec2Selection } from '../prompter'
1313
import { Ec2Node, Ec2ParentNode } from './ec2ParentNode'
1414
import { EC2 } from 'aws-sdk'
1515
import { getLogger } from '../../../shared/logger/logger'
16+
import { InstanceStateName } from '@aws-sdk/client-ec2'
1617

1718
export const Ec2InstanceRunningContext = 'awsEc2RunningNode'
1819
export const Ec2InstanceStoppedContext = 'awsEc2StoppedNode'
@@ -68,7 +69,7 @@ export class Ec2InstanceNode extends AWSTreeNodeBase implements AWSResourceNode
6869
return Ec2InstancePendingContext
6970
}
7071

71-
public setInstanceStatus(instanceStatus: string) {
72+
public setInstanceStatus(instanceStatus: InstanceStateName) {
7273
this.instance.LastSeenStatus = instanceStatus
7374
}
7475

@@ -84,7 +85,7 @@ export class Ec2InstanceNode extends AWSTreeNodeBase implements AWSResourceNode
8485
}
8586

8687
public get name(): string {
87-
return getNameOfInstance(this.instance) ?? `(no name)`
88+
return this.instance.Name ?? getNameOfInstance(this.instance) ?? `(no name)`
8889
}
8990

9091
public get InstanceId(): string {

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

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

1414
export const parentContextValue = 'awsEc2ParentNode'
@@ -30,7 +30,7 @@ export class Ec2ParentNode extends AWSTreeNodeBase {
3030
}
3131

3232
public override async getChildren(): Promise<AWSTreeNodeBase[]> {
33-
return await makeChildrenNodes({
33+
const result = await makeChildrenNodes({
3434
getChildNodes: async () => {
3535
await this.updateChildren()
3636

@@ -39,6 +39,7 @@ export class Ec2ParentNode extends AWSTreeNodeBase {
3939
getNoChildrenPlaceholderNode: async () => new PlaceholderNode(this, this.placeHolderMessage),
4040
sort: (nodeA, nodeB) => nodeA.name.localeCompare(nodeB.name),
4141
})
42+
return result
4243
}
4344

4445
public trackPendingNode(instanceId: string) {
@@ -49,7 +50,7 @@ export class Ec2ParentNode extends AWSTreeNodeBase {
4950
}
5051

5152
public async updateChildren(): Promise<void> {
52-
const ec2Instances = await (await this.ec2Client.getInstances()).toMap((instance) => instance.InstanceId)
53+
const ec2Instances = toMap(await this.ec2Client.getInstances(), (instance) => instance.InstanceId)
5354
updateInPlace(
5455
this.ec2InstanceNodes,
5556
ec2Instances.keys(),

packages/core/src/awsService/ec2/model.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import { getOrInstallCli } from '../../shared/utilities/cliUtils'
99
import { isCloud9 } from '../../shared/extensionUtilities'
1010
import { ToolkitError } from '../../shared/errors'
1111
import { SsmClient } from '../../shared/clients/ssm'
12-
import { Ec2Client } from '../../shared/clients/ec2Client'
12+
import { Ec2Client } from '../../shared/clients/ec2'
1313
import {
1414
VscodeRemoteConnection,
1515
createBoundProcess,

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,9 @@
55

66
import { RegionSubmenu, RegionSubmenuResponse } from '../../shared/ui/common/regionSubmenu'
77
import { DataQuickPickItem } from '../../shared/ui/pickerPrompter'
8-
import { Ec2Client, SafeEc2Instance } from '../../shared/clients/ec2Client'
8+
import { Ec2Client, SafeEc2Instance } from '../../shared/clients/ec2'
99
import { isValidResponse } from '../../shared/wizards/wizard'
1010
import { CancellationError } from '../../shared/utilities/timeoutUtils'
11-
import { AsyncCollection } from '../../shared/utilities/asyncCollection'
1211
import { getIconCode } from './utils'
1312
import { Ec2Node } from './explorer/ec2ParentNode'
1413
import { Ec2InstanceNode } from './explorer/ec2InstanceNode'
@@ -53,16 +52,15 @@ export class Ec2Prompter {
5352
}
5453
}
5554

56-
protected async getInstancesFromRegion(regionCode: string): Promise<AsyncCollection<SafeEc2Instance>> {
55+
protected async getInstancesFromRegion(regionCode: string): Promise<SafeEc2Instance[]> {
5756
const client = new Ec2Client(regionCode)
5857
return await client.getInstances()
5958
}
6059

6160
protected async getInstancesAsQuickPickItems(region: string): Promise<DataQuickPickItem<string>[]> {
6261
return (await this.getInstancesFromRegion(region))
63-
.filter(this.filter ? (instance) => this.filter!(instance) : (instance) => true)
62+
.filter(this.filter ? (instance) => this.filter!(instance) : (_) => true)
6463
.map((instance) => Ec2Prompter.asQuickPickItem(instance))
65-
.promise()
6664
}
6765

6866
private createEc2ConnectPrompter(): RegionSubmenu<string> {

packages/core/src/awsService/ec2/utils.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
import { SafeEc2Instance } from '../../shared/clients/ec2Client'
6+
import { SafeEc2Instance } from '../../shared/clients/ec2'
77
import { copyToClipboard } from '../../shared/utilities/messages'
88
import { Ec2Selection } from './prompter'
99
import { sshLogFileLocation } from '../../shared/sshConfig'

packages/core/src/awsexplorer/regionNode.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ import { DefaultSchemaClient } from '../shared/clients/schemaClient'
3131
import { getEcsRootNode } from '../awsService/ecs/model'
3232
import { compareTreeItems, TreeShim } from '../shared/treeview/utils'
3333
import { Ec2ParentNode } from '../awsService/ec2/explorer/ec2ParentNode'
34-
import { Ec2Client } from '../shared/clients/ec2Client'
34+
import { Ec2Client } from '../shared/clients/ec2'
3535

3636
interface ServiceNode {
3737
allRegions?: boolean
Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
import {
6+
EC2Client,
7+
Instance,
8+
InstanceStateName,
9+
GetConsoleOutputRequest,
10+
Filter,
11+
paginateDescribeInstances,
12+
DescribeInstancesRequest,
13+
Reservation,
14+
Tag,
15+
paginateDescribeInstanceStatus,
16+
StartInstancesCommandOutput,
17+
StartInstancesCommand,
18+
StopInstancesCommand,
19+
StopInstancesCommandOutput,
20+
RebootInstancesCommand,
21+
IamInstanceProfileAssociation,
22+
paginateDescribeIamInstanceProfileAssociations,
23+
IamInstanceProfile,
24+
GetConsoleOutputCommand,
25+
} from '@aws-sdk/client-ec2'
26+
import { Timeout } from '../utilities/timeoutUtils'
27+
import { showMessageWithCancel } from '../utilities/messages'
28+
import { ToolkitError, isAwsError } from '../errors'
29+
import { decodeBase64 } from '../utilities/textUtilities'
30+
import { ClientWrapper } from './clientWrapper'
31+
32+
/**
33+
* A wrapper around EC2.Instance where we can safely assume InstanceId field exists.
34+
*/
35+
export interface SafeEc2Instance extends Instance {
36+
InstanceId: string
37+
Name?: string
38+
LastSeenStatus: InstanceStateName
39+
}
40+
41+
interface SafeEc2GetConsoleOutputResult extends GetConsoleOutputRequest {
42+
Output: string
43+
InstanceId: string
44+
}
45+
46+
export class Ec2Client extends ClientWrapper<EC2Client> {
47+
public constructor(public override readonly regionCode: string) {
48+
super(regionCode, EC2Client)
49+
}
50+
51+
public async getInstances(filters?: Filter[]): Promise<SafeEc2Instance[]> {
52+
const reservations = await this.makePaginatedRequest(
53+
paginateDescribeInstances,
54+
filters ? { Filters: filters } : ({} satisfies DescribeInstancesRequest),
55+
(page) => page.Reservations
56+
)
57+
58+
return await this.updateInstancesDetail(this.getInstancesFromReservations(reservations))
59+
}
60+
61+
/** Updates status and name in-place for displaying to humans. */
62+
public async updateInstancesDetail(
63+
instances: Instance[],
64+
getStatus: (i: string) => Promise<InstanceStateName> = this.getInstanceStatus.bind(this)
65+
): Promise<SafeEc2Instance[]> {
66+
const instanceWithId = instances.filter(hasId)
67+
const instanceWithStatus = await Promise.all(instanceWithId.map(addStatus))
68+
return instanceWithStatus.map((i) => (instanceHasName(i) ? { ...i, Name: lookupTagKey(i.Tags, 'Name') } : i))
69+
70+
function hasId(i: Instance): i is Instance & { InstanceId: string } {
71+
return i.InstanceId !== undefined
72+
}
73+
74+
async function addStatus(instance: Instance & { InstanceId: string }) {
75+
return { ...instance, LastSeenStatus: await getStatus(instance.InstanceId) }
76+
}
77+
}
78+
79+
public getInstancesFromReservations(reservations: Reservation[]): (Instance & { InstanceId: string })[] {
80+
return reservations
81+
.map((r) => r.Instances)
82+
.flat()
83+
.filter(isNotEmpty)
84+
85+
function isNotEmpty(i: Instance | undefined): i is Instance & { InstanceId: string } {
86+
return i !== undefined && i.InstanceId !== undefined
87+
}
88+
}
89+
90+
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+
97+
return instanceStatuses[0].InstanceState!.Name!
98+
}
99+
100+
public async isInstanceRunning(instanceId: string): Promise<boolean> {
101+
const status = await this.getInstanceStatus(instanceId)
102+
return status === 'running'
103+
}
104+
105+
public getInstancesFilter(instanceIds: string[]): Filter[] {
106+
return [
107+
{
108+
Name: 'instance-id',
109+
Values: instanceIds,
110+
},
111+
]
112+
}
113+
114+
private handleStatusError(instanceId: string, err: unknown) {
115+
if (isAwsError(err)) {
116+
throw new ToolkitError(`EC2: failed to change status of instance ${instanceId}`, {
117+
cause: err as Error,
118+
})
119+
} else {
120+
throw err
121+
}
122+
}
123+
124+
public async assertNotInStatus(
125+
instanceId: string,
126+
targetStatus: string,
127+
getStatus: (i: string) => Promise<InstanceStateName> = this.getInstanceStatus.bind(this)
128+
) {
129+
const isAlreadyInStatus = (await getStatus(instanceId)) === targetStatus
130+
if (isAlreadyInStatus) {
131+
throw new ToolkitError(
132+
`EC2: Instance is currently ${targetStatus}. Unable to update status of ${instanceId}.`
133+
)
134+
}
135+
}
136+
137+
public async startInstance(instanceId: string): Promise<StartInstancesCommandOutput> {
138+
return await this.makeRequest(StartInstancesCommand, { InstanceIds: [instanceId] })
139+
}
140+
141+
public async startInstanceWithCancel(instanceId: string): Promise<void> {
142+
const timeout = new Timeout(5000)
143+
144+
await showMessageWithCancel(`EC2: Starting instance ${instanceId}`, timeout)
145+
146+
try {
147+
await this.assertNotInStatus(instanceId, 'running')
148+
await this.startInstance(instanceId)
149+
} catch (err) {
150+
this.handleStatusError(instanceId, err)
151+
} finally {
152+
timeout.cancel()
153+
}
154+
}
155+
156+
public async stopInstance(instanceId: string): Promise<StopInstancesCommandOutput> {
157+
return await this.makeRequest(StopInstancesCommand, { InstanceIds: [instanceId] })
158+
}
159+
160+
public async stopInstanceWithCancel(instanceId: string): Promise<void> {
161+
const timeout = new Timeout(5000)
162+
163+
await showMessageWithCancel(`EC2: Stopping instance ${instanceId}`, timeout)
164+
165+
try {
166+
await this.assertNotInStatus(instanceId, 'stopped')
167+
await this.stopInstance(instanceId)
168+
} catch (err) {
169+
this.handleStatusError(instanceId, err)
170+
} finally {
171+
timeout.cancel()
172+
}
173+
}
174+
175+
public async rebootInstance(instanceId: string): Promise<void> {
176+
return await this.makeRequest(RebootInstancesCommand, { InstanceIds: [instanceId] })
177+
}
178+
179+
public async rebootInstanceWithCancel(instanceId: string): Promise<void> {
180+
const timeout = new Timeout(5000)
181+
182+
await showMessageWithCancel(`EC2: Rebooting instance ${instanceId}`, timeout)
183+
184+
try {
185+
await this.rebootInstance(instanceId)
186+
} catch (err) {
187+
this.handleStatusError(instanceId, err)
188+
} finally {
189+
timeout.cancel()
190+
}
191+
}
192+
193+
/**
194+
* Retrieve IAM Association for a given EC2 instance.
195+
* @param instanceId target EC2 instance ID
196+
* @returns IAM Association for instance
197+
*/
198+
private async getIamInstanceProfileAssociation(instanceId: string): Promise<IamInstanceProfileAssociation> {
199+
const instanceFilter = this.getInstancesFilter([instanceId])
200+
201+
const associations = await this.makePaginatedRequest(
202+
paginateDescribeIamInstanceProfileAssociations,
203+
{ Filters: instanceFilter },
204+
(page) => page.IamInstanceProfileAssociations
205+
)
206+
207+
return associations[0]!
208+
}
209+
210+
/**
211+
* Gets the IAM Instance Profile (not role) attached to given EC2 instance.
212+
* @param instanceId target EC2 instance ID
213+
* @returns IAM Instance Profile associated with instance or undefined if none exists.
214+
*/
215+
public async getAttachedIamInstanceProfile(instanceId: string): Promise<IamInstanceProfile | undefined> {
216+
const association = await this.getIamInstanceProfileAssociation(instanceId)
217+
return association ? association.IamInstanceProfile : undefined
218+
}
219+
220+
public async getConsoleOutput(instanceId: string, latest: boolean): Promise<SafeEc2GetConsoleOutputResult> {
221+
const response = await this.makeRequest(GetConsoleOutputCommand, { InstanceId: instanceId, Latest: latest })
222+
223+
return {
224+
...response,
225+
InstanceId: instanceId,
226+
Output: response.Output ? decodeBase64(response.Output) : '',
227+
}
228+
}
229+
}
230+
231+
export function getNameOfInstance(instance: Instance): string | undefined {
232+
return instanceHasName(instance) ? lookupTagKey(instance.Tags!, 'Name')! : undefined
233+
}
234+
235+
export function instanceHasName(instance: Instance): instance is Instance & { Tags: Tag[] } {
236+
return instance.Tags !== undefined && instance.Tags.some((tag) => tag.Key === 'Name')
237+
}
238+
239+
function lookupTagKey(tags: Tag[], targetKey: string) {
240+
return tags.filter((tag) => tag.Key === targetKey)[0].Value
241+
}

0 commit comments

Comments
 (0)