Skip to content

Commit b6d1735

Browse files
Will-ShaoHuaWill-ShaoHua
committed
Merge remote-tracking branch 'upstream/master' into telemetry
2 parents 9c8beee + 77f1f47 commit b6d1735

File tree

15 files changed

+963
-134
lines changed

15 files changed

+963
-134
lines changed

‎packages/amazonq/.changes/next-release/Bug Fix-d95abc3d-f40d-481d-bbc2-523aa572f5cd.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Bug Fix",
3+
"description": "Amazon Q support web/container environments running Ubuntu/Linux, even when the host machine is Amazon Linux 2."
4+
}

‎packages/amazonq/src/lsp/auth.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,8 @@ export class AmazonQLspAuth {
118118
)
119119
if (credentials) {
120120
getLogger().info(
121-
`[SageMaker Debug] IAM credentials structure: accessKeyId=${credentials.accessKeyId ? 'present' : 'missing'}, secretAccessKey=${credentials.secretAccessKey ? 'present' : 'missing'}, sessionToken=${credentials.sessionToken ? 'present' : 'missing'}`
121+
`[SageMaker Debug] IAM credentials structure: accessKeyId=${credentials.accessKeyId ? 'present' : 'missing'}, secretAccessKey=${credentials.secretAccessKey ? 'present' : 'missing'}, sessionToken=${credentials.sessionToken ? 'present' : 'missing'}, expiration=${credentials.expiration} ? 'pr
122+
esent' : 'missing'}`
122123
)
123124
}
124125

@@ -163,6 +164,7 @@ export class AmazonQLspAuth {
163164
accessKeyId: credentials.accessKeyId,
164165
secretAccessKey: credentials.secretAccessKey,
165166
sessionToken: credentials.sessionToken,
167+
expiration: credentials.expiration,
166168
}
167169
const payload = new TextEncoder().encode(JSON.stringify({ data: iamCredentials }))
168170

‎packages/core/src/auth/credentials/store.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,14 +95,14 @@ export class CredentialsStore {
9595
credentialsId: CredentialsId,
9696
credentialsProvider: CredentialsProvider
9797
): Promise<CachedCredentials> {
98+
getLogger().debug(`store: Fetch new credentials from provider with id: ${asString(credentialsId)}`)
9899
const credentials = {
99100
credentials: await credentialsProvider.getCredentials(),
100101
credentialsHashCode: credentialsProvider.getHashCode(),
101102
endpointUrl: credentialsProvider.getEndpointUrl?.(),
102103
}
103104

104105
this.credentialsCache[asString(credentialsId)] = credentials
105-
106106
return credentials
107107
}
108108
}

‎packages/core/src/codewhisperer/commands/invokeRecommendation.ts‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,6 @@ export async function invokeRecommendation(
2020
client: DefaultCodeWhispererClient,
2121
config: ConfigurationEntry
2222
) {
23-
// Call report user decisions once to report recommendations leftover from last invocation.
24-
RecommendationHandler.instance.reportUserDecisions(-1)
25-
2623
if (!editor || !config.isManualTriggerEnabled) {
2724
return
2825
}

‎packages/core/src/codewhisperer/service/inlineCompletionService.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,8 @@ export class InlineCompletionService {
8888
}
8989
}
9090

91+
// Call report user decisions once to report recommendations leftover from last invocation.
92+
RecommendationHandler.instance.reportUserDecisions(-1)
9193
TelemetryHelper.instance.setInvokeSuggestionStartTime()
9294
ClassifierTrigger.instance.recordClassifierResultForAutoTrigger(editor, autoTriggerType, event)
9395

‎packages/core/src/codewhisperer/service/keyStrokeHandler.ts‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -87,9 +87,6 @@ export class KeyStrokeHandler {
8787
client: DefaultCodeWhispererClient,
8888
config: ConfigurationEntry
8989
): Promise<void> {
90-
// Call report user decisions once to report recommendations leftover from last invocation.
91-
RecommendationHandler.instance.reportUserDecisions(-1)
92-
9390
try {
9491
if (!config.isAutomatedTriggerEnabled) {
9592
return

‎packages/core/src/sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider.ts‎

Lines changed: 92 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import * as localizedText from '../../../shared/localizedText'
1414
import { ToolkitPromptSettings } from '../../../shared/settings'
1515
import { setContext, getContext } from '../../../shared/vscode/setContext'
1616
import { getLogger } from '../../../shared/logger/logger'
17-
import { SmusUtils, SmusErrorCodes, extractAccountIdFromArn } from '../../shared/smusUtils'
17+
import { SmusUtils, SmusErrorCodes, extractAccountIdFromResourceMetadata } from '../../shared/smusUtils'
1818
import { createSmusProfile, isValidSmusConnection, SmusConnection } from '../model'
1919
import { DomainExecRoleCredentialsProvider } from './domainExecRoleCredentialsProvider'
2020
import { ProjectRoleCredentialsProvider } from './projectRoleCredentialsProvider'
@@ -24,6 +24,7 @@ import { getResourceMetadata } from '../../shared/utils/resourceMetadataUtils'
2424
import { fromIni } from '@aws-sdk/credential-providers'
2525
import { randomUUID } from '../../../shared/crypto'
2626
import { DefaultStsClient } from '../../../shared/clients/stsClient'
27+
import { DataZoneClient } from '../../shared/client/datazoneClient'
2728

2829
/**
2930
* Sets the context variable for SageMaker Unified Studio connection state
@@ -55,6 +56,7 @@ export class SmusAuthenticationProvider {
5556
private projectCredentialProvidersCache = new Map<string, ProjectRoleCredentialsProvider>()
5657
private connectionCredentialProvidersCache = new Map<string, ConnectionCredentialsProvider>()
5758
private cachedDomainAccountId: string | undefined
59+
private cachedProjectAccountIds = new Map<string, string>()
5860

5961
public constructor(
6062
public readonly auth = Auth.instance,
@@ -79,6 +81,8 @@ export class SmusAuthenticationProvider {
7981
this.connectionCredentialProvidersCache.clear()
8082
// Clear cached domain account ID when connection changes
8183
this.cachedDomainAccountId = undefined
84+
// Clear cached project account IDs when connection changes
85+
this.cachedProjectAccountIds.clear()
8286
// Clear all clients in client store when connection changes
8387
ConnectionClientStore.getInstance().clearAll()
8488
await setSmusConnectedContext(this.isConnected())
@@ -445,37 +449,13 @@ export class SmusAuthenticationProvider {
445449

446450
// If in SMUS space environment, extract account ID from resource-metadata file
447451
if (getContext('aws.smus.inSmusSpaceEnvironment')) {
448-
try {
449-
logger.debug('SMUS: Extracting domain account ID from ResourceArn in resource-metadata file')
450-
451-
const resourceMetadata = getResourceMetadata()!
452-
const resourceArn = resourceMetadata.ResourceArn
453-
454-
if (!resourceArn) {
455-
throw new ToolkitError('ResourceArn not found in metadata file', {
456-
code: SmusErrorCodes.AccountIdNotFound,
457-
})
458-
}
452+
const accountId = await extractAccountIdFromResourceMetadata()
459453

460-
// Extract account ID from ResourceArn using SmusUtils
461-
const accountId = extractAccountIdFromArn(resourceArn)
462-
463-
// Cache the account ID
464-
this.cachedDomainAccountId = accountId
465-
466-
logger.debug(
467-
`Successfully extracted and cached domain account ID from resource-metadata file: ${accountId}`
468-
)
469-
470-
return accountId
471-
} catch (err) {
472-
logger.error(`Failed to extract domain account ID from ResourceArn: %s`, err)
454+
// Cache the account ID
455+
this.cachedDomainAccountId = accountId
456+
logger.debug(`Successfully cached domain account ID: ${accountId}`)
473457

474-
throw new ToolkitError('Failed to extract AWS account ID from ResourceArn in SMUS space environment', {
475-
code: SmusErrorCodes.GetDomainAccountIdFailed,
476-
cause: err instanceof Error ? err : undefined,
477-
})
478-
}
458+
return accountId
479459
}
480460

481461
if (!this.activeConnection) {
@@ -520,6 +500,81 @@ export class SmusAuthenticationProvider {
520500
}
521501
}
522502

503+
/**
504+
* Gets the AWS account ID for a specific project using project credentials
505+
* In SMUS space environment, extracts from ResourceArn in metadata (same as domain account)
506+
* Otherwise, makes an STS GetCallerIdentity call using project credentials
507+
* @param projectId The DataZone project ID
508+
* @returns Promise resolving to the project's AWS account ID
509+
*/
510+
public async getProjectAccountId(projectId: string): Promise<string> {
511+
const logger = getLogger()
512+
513+
// Return cached value if available
514+
if (this.cachedProjectAccountIds.has(projectId)) {
515+
logger.debug(`SMUS: Using cached project account ID for project ${projectId}`)
516+
return this.cachedProjectAccountIds.get(projectId)!
517+
}
518+
519+
// If in SMUS space environment, extract account ID from resource-metadata file
520+
if (getContext('aws.smus.inSmusSpaceEnvironment')) {
521+
const accountId = await extractAccountIdFromResourceMetadata()
522+
523+
// Cache the account ID
524+
this.cachedProjectAccountIds.set(projectId, accountId)
525+
logger.debug(`Successfully cached project account ID for project ${projectId}: ${accountId}`)
526+
527+
return accountId
528+
}
529+
530+
if (!this.activeConnection) {
531+
throw new ToolkitError('No active SMUS connection available', { code: SmusErrorCodes.NoActiveConnection })
532+
}
533+
534+
// For non-SMUS space environments, use project credentials with STS
535+
try {
536+
logger.debug('Fetching project account ID via STS GetCallerIdentity with project credentials')
537+
538+
// Get project credentials
539+
const projectCredProvider = await this.getProjectCredentialProvider(projectId)
540+
const projectCreds = await projectCredProvider.getCredentials()
541+
542+
// Get project region from tooling environment
543+
const dzClient = await DataZoneClient.getInstance(this)
544+
const toolingEnv = await dzClient.getToolingEnvironment(projectId)
545+
const projectRegion = toolingEnv.awsAccountRegion
546+
547+
if (!projectRegion) {
548+
throw new ToolkitError('No AWS account region found in tooling environment', {
549+
code: SmusErrorCodes.RegionNotFound,
550+
})
551+
}
552+
553+
// Use STS to get account ID from project credentials
554+
const stsClient = new DefaultStsClient(projectRegion, projectCreds)
555+
const callerIdentity = await stsClient.getCallerIdentity()
556+
557+
if (!callerIdentity.Account) {
558+
throw new ToolkitError('Account ID not found in STS GetCallerIdentity response', {
559+
code: SmusErrorCodes.AccountIdNotFound,
560+
})
561+
}
562+
563+
// Cache the account ID
564+
this.cachedProjectAccountIds.set(projectId, callerIdentity.Account)
565+
logger.debug(
566+
`Successfully retrieved and cached project account ID for project ${projectId}: ${callerIdentity.Account}`
567+
)
568+
569+
return callerIdentity.Account
570+
} catch (err) {
571+
logger.error('Failed to get project account ID: %s', err as Error)
572+
throw new ToolkitError(`Failed to get project account ID: ${(err as Error).message}`, {
573+
code: SmusErrorCodes.GetProjectAccountIdFailed,
574+
})
575+
}
576+
}
577+
523578
public getDomainRegion(): string {
524579
if (getContext('aws.smus.inSmusSpaceEnvironment')) {
525580
const resourceMetadata = getResourceMetadata()!
@@ -617,6 +672,10 @@ export class SmusAuthenticationProvider {
617672
// Clear cached domain account ID
618673
this.cachedDomainAccountId = undefined
619674
logger.debug('SMUS: Cleared cached domain account ID')
675+
676+
// Clear cached project account IDs
677+
this.cachedProjectAccountIds.clear()
678+
logger.debug('SMUS: Cleared cached project account IDs')
620679
}
621680

622681
/**
@@ -665,6 +724,9 @@ export class SmusAuthenticationProvider {
665724
// Clear cached domain account ID
666725
this.cachedDomainAccountId = undefined
667726

727+
// Clear cached project account IDs
728+
this.cachedProjectAccountIds.clear()
729+
668730
this.logger.debug('SMUS Auth: Successfully disposed authentication provider')
669731
}
670732

‎packages/core/src/sagemakerunifiedstudio/explorer/nodes/redshiftStrategy.ts‎

Lines changed: 3 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ import { createPlaceholderItem } from '../../../shared/treeview/utils'
2424
import { ConnectionCredentialsProvider } from '../../auth/providers/connectionCredentialsProvider'
2525
import { GlueCatalog } from '../../shared/client/glueCatalogClient'
2626
import { telemetry } from '../../../shared/telemetry/telemetry'
27-
import { getContext } from '../../../shared/vscode/setContext'
27+
import { recordDataConnectionTelemetry } from '../../shared/telemetry'
2828

2929
/**
3030
* Redshift data node for SageMaker Unified Studio
@@ -119,6 +119,7 @@ export function createRedshiftConnectionNode(
119119
connection: DataZoneConnection,
120120
connectionCredentialsProvider: ConnectionCredentialsProvider
121121
): RedshiftNode {
122+
const logger = getLogger()
122123
return new RedshiftNode(
123124
{
124125
id: connection.connectionId,
@@ -130,19 +131,8 @@ export function createRedshiftConnectionNode(
130131
},
131132
async (node) => {
132133
return telemetry.smus_renderRedshiftNode.run(async (span) => {
133-
const isInSmusSpace = getContext('aws.smus.inSmusSpaceEnvironment')
134-
const accountId = await connectionCredentialsProvider.getDomainAccountId()
135-
span.record({
136-
smusToolkitEnv: isInSmusSpace ? 'smus_space' : 'local',
137-
smusDomainId: connection.domainId,
138-
smusDomainAccountId: accountId,
139-
smusProjectId: connection.projectId,
140-
smusConnectionId: connection.connectionId,
141-
smusConnectionType: connection.type,
142-
smusProjectRegion: connection.location?.awsRegion,
143-
})
144-
const logger = getLogger()
145134
logger.info(`Loading Redshift resources for connection ${connection.name}`)
135+
await recordDataConnectionTelemetry(span, connection, connectionCredentialsProvider)
146136

147137
const connectionParams = extractConnectionParams(connection)
148138
if (!connectionParams) {

‎packages/core/src/sagemakerunifiedstudio/shared/smusUtils.ts‎

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,14 @@ export const SmusErrorCodes = {
5656
UserCancelled: 'UserCancelled',
5757
/** Error code for when domain account Id is missing */
5858
AccountIdNotFound: 'AccountIdNotFound',
59+
/** Error code for when resource ARN is missing */
60+
ResourceArnNotFound: 'ResourceArnNotFound',
5961
/** Error code for when fails to get domain account Id */
6062
GetDomainAccountIdFailed: 'GetDomainAccountIdFailed',
63+
/** Error code for when fails to get project account Id */
64+
GetProjectAccountIdFailed: 'GetProjectAccountIdFailed',
65+
/** Error code for when region is missing */
66+
RegionNotFound: 'RegionNotFound',
6167
} as const
6268

6369
/**
@@ -369,9 +375,9 @@ export class SmusUtils {
369375
* @returns The account ID from the ARN
370376
* @throws If the ARN format is invalid
371377
*/
372-
export function extractAccountIdFromArn(arn: string): string {
378+
export function extractAccountIdFromSageMakerArn(arn: string): string {
373379
// Match the ARN components to extract account ID
374-
const regex = /^arn:aws:sagemaker:(?<region>[^:]+):(?<accountId>\d+):(app|space|domain)\/.+$/i
380+
const regex = /^arn:aws:sagemaker:(?<region>[^:]+):(?<accountId>\d+):(app|space)\/.+$/i
375381
const match = arn.match(regex)
376382

377383
if (!match?.groups) {
@@ -380,3 +386,31 @@ export function extractAccountIdFromArn(arn: string): string {
380386

381387
return match.groups.accountId
382388
}
389+
390+
/**
391+
* Extracts account ID from ResourceArn in SMUS space environment
392+
* @returns Promise resolving to the account ID
393+
* @throws ToolkitError if unable to extract account ID
394+
*/
395+
export async function extractAccountIdFromResourceMetadata(): Promise<string> {
396+
const logger = getLogger()
397+
398+
try {
399+
logger.debug('SMUS: Extracting account ID from ResourceArn in resource-metadata file')
400+
401+
const resourceMetadata = getResourceMetadata()!
402+
const resourceArn = resourceMetadata.ResourceArn
403+
404+
if (!resourceArn) {
405+
throw new Error('ResourceArn not found in metadata file')
406+
}
407+
408+
const accountId = extractAccountIdFromSageMakerArn(resourceArn)
409+
logger.debug(`Successfully extracted account ID from resource-metadata file: ${accountId}`)
410+
411+
return accountId
412+
} catch (err) {
413+
logger.error(`Failed to extract account ID from ResourceArn: %s`, err)
414+
throw new Error('Failed to extract AWS account ID from ResourceArn in SMUS space environment')
415+
}
416+
}

0 commit comments

Comments
 (0)