Skip to content

Commit 424a1af

Browse files
liramon1liramon1
committed
Revert unnecessary region profile changes
1 parent 422d085 commit 424a1af

File tree

5 files changed

+44
-111
lines changed

5 files changed

+44
-111
lines changed

‎packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ describe('RegionProfileManager', async function () {
6262
const mockClient = {
6363
listAvailableProfiles: listProfilesStub,
6464
}
65-
const createClientStub = sinon.stub(regionProfileManager, '_createQUserClient').resolves(mockClient)
65+
const createClientStub = sinon.stub(regionProfileManager, '_createQClient').resolves(mockClient)
6666

6767
const profileList = await regionProfileManager.listRegionProfile()
6868

@@ -272,11 +272,11 @@ describe('RegionProfileManager', async function () {
272272
})
273273
})
274274

275-
describe('createQUserClient', function () {
275+
describe('createQClient', function () {
276276
it(`should configure the endpoint and region from a profile`, async function () {
277277
await setupConnection('idc')
278278

279-
const iadClient = await regionProfileManager.createQUserClient({
279+
const iadClient = await regionProfileManager.createQClient({
280280
name: 'foo',
281281
region: 'us-east-1',
282282
arn: 'arn',
@@ -286,7 +286,7 @@ describe('RegionProfileManager', async function () {
286286
assert.deepStrictEqual(iadClient.config.region, 'us-east-1')
287287
assert.deepStrictEqual(iadClient.endpoint.href, 'https://q.us-east-1.amazonaws.com/')
288288

289-
const fraClient = await regionProfileManager.createQUserClient({
289+
const fraClient = await regionProfileManager.createQClient({
290290
name: 'bar',
291291
region: 'eu-central-1',
292292
arn: 'arn',
@@ -302,7 +302,7 @@ describe('RegionProfileManager', async function () {
302302

303303
await assert.rejects(
304304
async () => {
305-
await regionProfileManager.createQUserClient({
305+
await regionProfileManager.createQClient({
306306
name: 'foo',
307307
region: 'ap-east-1',
308308
arn: 'arn',
@@ -314,7 +314,7 @@ describe('RegionProfileManager', async function () {
314314

315315
await assert.rejects(
316316
async () => {
317-
await regionProfileManager.createQUserClient({
317+
await regionProfileManager.createQClient({
318318
name: 'foo',
319319
region: 'unknown-somewhere',
320320
arn: 'arn',
@@ -330,7 +330,7 @@ describe('RegionProfileManager', async function () {
330330
await regionProfileManager.switchRegionProfile(profileFoo, 'user')
331331
assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo)
332332

333-
const client = await regionProfileManager._createQUserClient('eu-central-1', 'https://amazon.com/')
333+
const client = await regionProfileManager._createQClient('eu-central-1', 'https://amazon.com/')
334334

335335
assert.deepStrictEqual(client.config.region, 'eu-central-1')
336336
assert.deepStrictEqual(client.endpoint.href, 'https://amazon.com/')

‎packages/core/src/codewhisperer/region/regionProfileManager.ts‎

Lines changed: 33 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,9 @@ import { showConfirmationMessage } from '../../shared/utilities/messages'
1111
import globals from '../../shared/extensionGlobals'
1212
import { once } from '../../shared/utilities/functionUtils'
1313
import CodeWhispererUserClient from '../client/codewhispereruserclient'
14-
import CodeWhispererClient from '../client/codewhispererclient'
1514
import { Credentials, Service } from 'aws-sdk'
1615
import { ServiceOptions } from '../../shared/awsClientBuilder'
1716
import userApiConfig = require('../client/user-service-2.json')
18-
import apiConfig = require('../client/service-2.json')
1917
import { createConstantMap } from '../../shared/utilities/tsUtils'
2018
import { getLogger } from '../../shared/logger/logger'
2119
import { pageableToCollection } from '../../shared/utilities/collectionUtils'
@@ -146,70 +144,37 @@ export class RegionProfileManager {
146144
async listRegionProfile(): Promise<RegionProfile[]> {
147145
this._profiles = []
148146

149-
if (!this.authProvider.isConnected()) {
147+
if (!this.authProvider.isConnected() || !this.authProvider.isSsoSession()) {
150148
return []
151149
}
152150
const availableProfiles: RegionProfile[] = []
153151
const failedRegions: string[] = []
154152

155153
for (const [region, endpoint] of endpoints.entries()) {
154+
const client = await this._createQClient(region, endpoint)
155+
const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) =>
156+
client.listAvailableProfiles(request).promise()
157+
const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {}
156158
try {
157-
// Get region profiles (Q developer profiles) from Q client and authenticate with SSO token
158-
if (this.authProvider.isIdcConnection()) {
159-
const client = await this._createQUserClient(region, endpoint)
160-
const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) => {
161-
return client.listAvailableProfiles(request).promise()
159+
const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles')
160+
.flatten()
161+
.promise()
162+
const mappedPfs = profiles.map((it) => {
163+
let accntId = ''
164+
try {
165+
accntId = parse(it.arn).accountId
166+
} catch (e) {}
167+
168+
return {
169+
name: it.profileName,
170+
region: region,
171+
arn: it.arn,
172+
description: accntId,
162173
}
163-
const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {}
164-
const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles')
165-
.flatten()
166-
.promise()
167-
const mappedPfs = profiles.map((it) => {
168-
let accntId = ''
169-
try {
170-
accntId = parse(it.arn).accountId
171-
} catch (e) {}
172-
173-
return {
174-
name: it.profileName,
175-
region: region,
176-
arn: it.arn,
177-
description: accntId,
178-
}
179-
})
180-
181-
availableProfiles.push(...mappedPfs)
182-
RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`)
183-
}
184-
// Get region profiles (Q developer profiles) from Q client and authenticate with IAM credentials
185-
else if (this.authProvider.isIamSession()) {
186-
const client = await this._createQServiceClient(region, endpoint)
187-
const requester = async (request: CodeWhispererClient.ListProfilesRequest) => {
188-
return client.listProfiles(request).promise()
189-
}
190-
const request: CodeWhispererClient.ListProfilesRequest = {}
191-
const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles')
192-
.flatten()
193-
.promise()
194-
const mappedPfs = profiles.map((it) => {
195-
let accntId = ''
196-
try {
197-
accntId = parse(it.arn).accountId
198-
} catch (e) {}
199-
200-
return {
201-
name: it.profileName,
202-
region: region,
203-
arn: it.arn,
204-
description: accntId,
205-
}
206-
})
174+
})
207175

208-
availableProfiles.push(...mappedPfs)
209-
RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`)
210-
} else {
211-
throw new ToolkitError('Failed to list profiles when signed out of identity center and IAM credentials')
212-
}
176+
availableProfiles.push(...mappedPfs)
177+
RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`)
213178
} catch (e) {
214179
const logMsg = isAwsError(e) ? `requestId=${e.requestId}; message=${e.message}` : (e as Error).message
215180
RegionProfileManager.logger.error(`Failed to list profiles for region ${region}: ${logMsg}`)
@@ -235,7 +200,7 @@ export class RegionProfileManager {
235200
}
236201

237202
async switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) {
238-
if (!this.authProvider.isConnected()) {
203+
if (!this.authProvider.isConnected() || !this.authProvider.isIdcConnection()) {
239204
return
240205
}
241206

@@ -439,40 +404,27 @@ export class RegionProfileManager {
439404
if (this.authProvider.isBuilderIdConnection()) {
440405
return false
441406
}
442-
return (this.authProvider.isIdcConnection() || this.authProvider.isIamSession()) && this.activeRegionProfile === undefined
407+
return this.authProvider.isIdcConnection() && this.activeRegionProfile === undefined
443408
}
444409

445410
async clearCache() {
446411
await this.cache.clearCache()
447412
}
448413

449414
// TODO: Should maintain sdk client in a better way
450-
// Create a Q user client compatible with SSO tokens
451-
async createQUserClient(profile: RegionProfile): Promise<CodeWhispererUserClient> {
452-
if (!this.authProvider.isConnected()) {
453-
throw new Error('No valid connection')
454-
}
455-
const endpoint = endpoints.get(profile.region)
456-
if (!endpoint) {
457-
throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`)
458-
}
459-
return this._createQUserClient(profile.region, endpoint)
460-
}
461-
462-
// Create a Q service client compatible with IAM credentials
463-
async createQServiceClient(profile: RegionProfile): Promise<CodeWhispererClient> {
464-
if (!this.authProvider.isConnected()) {
465-
throw new Error('No valid connection')
415+
async createQClient(profile: RegionProfile): Promise<CodeWhispererUserClient> {
416+
if (!this.authProvider.isConnected() || !this.authProvider.isSsoSession()) {
417+
throw new Error('No valid SSO connection')
466418
}
467419
const endpoint = endpoints.get(profile.region)
468420
if (!endpoint) {
469421
throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`)
470422
}
471-
return this._createQServiceClient(profile.region, endpoint)
423+
return this._createQClient(profile.region, endpoint)
472424
}
473425

474-
// Visible for testing only, do not use this directly, please use createQUserClient(profile)
475-
async _createQUserClient(region: string, endpoint: string): Promise<CodeWhispererUserClient> {
426+
// Visible for testing only, do not use this directly, please use createQClient(profile)
427+
async _createQClient(region: string, endpoint: string): Promise<CodeWhispererUserClient> {
476428
const token = await this.authProvider.getToken()
477429
const serviceOption: ServiceOptions = {
478430
apiConfig: userApiConfig,
@@ -487,32 +439,13 @@ export class RegionProfileManager {
487439
},
488440
],
489441
} as ServiceOptions
490-
491-
return (await globals.sdkClientBuilder.createAwsService(
442+
443+
const c = (await globals.sdkClientBuilder.createAwsService(
492444
Service,
493445
serviceOption,
494446
undefined
495447
)) as CodeWhispererUserClient
496-
}
497-
498-
// Visible for testing only, do not use this directly, please use createQServiceClient(profile)
499-
async _createQServiceClient(region: string, endpoint: string): Promise<CodeWhispererClient> {
500-
const credential = await this.authProvider.getIamCredential()
501-
const serviceOption: ServiceOptions = {
502-
apiConfig: apiConfig,
503-
region: region,
504-
endpoint: endpoint,
505-
credentials: new Credentials({
506-
accessKeyId: credential.accessKeyId,
507-
secretAccessKey: credential.secretAccessKey,
508-
sessionToken: credential.sessionToken,
509-
}),
510-
} as ServiceOptions
511448

512-
return (await globals.sdkClientBuilder.createAwsService(
513-
Service,
514-
serviceOption,
515-
undefined
516-
)) as CodeWhispererClient
449+
return c
517450
}
518451
}

‎packages/core/src/codewhisperer/util/authUtil.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -364,7 +364,7 @@ export class AuthUtil implements IAuthProvider {
364364
private async refreshState(state = this.getAuthState()) {
365365
if (state === 'expired' || state === 'notConnected') {
366366
this.lspAuth.deleteBearerToken()
367-
if (this.isIdcConnection() || this.isIamSession()) {
367+
if (this.isIdcConnection()) {
368368
await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn)
369369
await this.regionProfileManager.clearCache()
370370
}
@@ -379,7 +379,7 @@ export class AuthUtil implements IAuthProvider {
379379
await this.lspAuth.updateIamCredential(params)
380380
}
381381

382-
if (this.isIdcConnection() || this.isIamSession()) {
382+
if (this.isIdcConnection()) {
383383
await this.regionProfileManager.restoreProfileSelection()
384384
}
385385
}

‎packages/core/src/codewhisperer/util/customizationUtil.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ export class CustomizationProvider {
4949
}
5050

5151
static async init(profile: RegionProfile): Promise<CustomizationProvider> {
52-
const client = await AuthUtil.instance.regionProfileManager.createQUserClient(profile)
52+
const client = await AuthUtil.instance.regionProfileManager.createQClient(profile)
5353
return new CustomizationProvider(client, profile)
5454
}
5555
}

‎packages/core/src/test/amazonq/customizationUtil.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ describe('customizationProvider', function () {
4343
regionProfileManager: regionProfileManager,
4444
}
4545
sinon.stub(AuthUtil, 'instance').get(() => mockAuthUtil)
46-
const createClientStub = sinon.stub(regionProfileManager, 'createQUserClient')
46+
const createClientStub = sinon.stub(regionProfileManager, 'createQClient')
4747
const mockProfile = {
4848
name: 'foo',
4949
region: 'us-east-1',

0 commit comments

Comments
 (0)