Skip to content

Commit f319283

Browse files
evanliu048evanliu048
committed
add an optional profile in getQclient
1 parent 2c7d3fc commit f319283

File tree

5 files changed

+36
-47
lines changed

5 files changed

+36
-47
lines changed

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.UserIntent
3838
import software.aws.toolkits.core.utils.debug
3939
import software.aws.toolkits.core.utils.getLogger
4040
import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext
41+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
4142
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
4243
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
4344
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
@@ -78,7 +79,7 @@ interface CodeWhispererClientAdaptor {
7879

7980
fun getCodeFixJob(request: GetCodeFixJobRequest): GetCodeFixJobResponse
8081

81-
fun listAvailableCustomizations(): List<CodeWhispererCustomization>
82+
fun listAvailableCustomizations(profile: QRegionProfile): List<CodeWhispererCustomization>
8283

8384
fun startTestGeneration(uploadId: String, targetCode: List<TargetCode>, userInput: String): StartTestGenerationResponse
8485

@@ -282,9 +283,9 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
282283
override fun getCodeFixJob(request: GetCodeFixJobRequest): GetCodeFixJobResponse = bearerClient().getCodeFixJob(request)
283284

284285
// DO NOT directly use this method to fetch customizations, use wrapper [CodeWhispererModelConfigurator.listCustomization()] instead
285-
override fun listAvailableCustomizations(): List<CodeWhispererCustomization> =
286-
bearerClient().listAvailableCustomizationsPaginator(
287-
ListAvailableCustomizationsRequest.builder().profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn).build()
286+
override fun listAvailableCustomizations(profile: QRegionProfile): List<CodeWhispererCustomization> =
287+
QRegionProfileManager.getInstance().getQClient<CodeWhispererRuntimeClient>(project, profile).listAvailableCustomizationsPaginator(
288+
ListAvailableCustomizationsRequest.builder().profileArn(profile.arn).build()
288289
)
289290
.stream()
290291
.toList()
@@ -298,7 +299,8 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
298299
CodeWhispererCustomization(
299300
arn = it.arn(),
300301
name = it.name(),
301-
description = it.description()
302+
description = it.description(),
303+
profile = profile
302304
)
303305
}
304306
}

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomizationDialog.kt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2023 Amazon.com, Inc. or it s affiliates. All Rights Reserved.
1+
// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
33

44
package software.aws.toolkits.jetbrains.services.codewhisperer.customization
@@ -36,7 +36,7 @@ import javax.swing.JComponent
3636
import javax.swing.JList
3737

3838
private val NoDataToDisplay = CustomizationUiItem(
39-
CodeWhispererCustomization("", message("codewhisperer.custom.dialog.option.no_data"), "", QRegionProfile("","")),
39+
CodeWhispererCustomization("", message("codewhisperer.custom.dialog.option.no_data"), "", QRegionProfile("", "")),
4040
false,
4141
false
4242
)

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt

Lines changed: 16 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,6 @@ import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererR
2121
import software.aws.toolkits.core.utils.debug
2222
import software.aws.toolkits.core.utils.getLogger
2323
import software.aws.toolkits.core.utils.tryOrNull
24-
import software.aws.toolkits.jetbrains.core.AwsClientManager
25-
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
26-
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
27-
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
2824
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
2925
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
3026
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
@@ -114,32 +110,16 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
114110

115111
@RequiresBackgroundThread
116112
override fun listCustomizations(project: Project, passive: Boolean): List<CustomizationUiItem>? =
117-
calculateIfIamIdentityCenterConnection(project) { it ->
113+
calculateIfIamIdentityCenterConnection(project) {
118114
// 1. invoke API and get result
119115
val listAvailableProfilesResult = tryOrNull { QRegionProfileManager.getInstance().listRegionProfiles(project) } ?: emptyList()
120116

121117
val aggregatedCustomizations = mutableListOf<CodeWhispererCustomization>()
122118

123119
for (profile in listAvailableProfilesResult) {
124120
try {
125-
val setting = AwsRegionProvider.getInstance()[profile.region]?.let { it1 ->
126-
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.getConnectionSettings()
127-
?.withRegion(it1)
128-
}
129-
val client = setting?.let { it1 -> AwsClientManager.getInstance().getClient(CodeWhispererRuntimeClient::class, it1) }
130-
val perProfileCustomizations =
131-
client?.listAvailableCustomizationsPaginator(ListAvailableCustomizationsRequest.builder().profileArn(profile.arn).build())
132-
?.customizations()?.stream()?.toList()?.map {
133-
CodeWhispererCustomization(
134-
arn = it.arn(),
135-
name = it.name(),
136-
description = it.description(),
137-
profile = profile
138-
)
139-
}
140-
if (perProfileCustomizations != null) {
141-
aggregatedCustomizations.addAll(perProfileCustomizations)
142-
}
121+
val perProfileCustomizations = CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations(profile)
122+
aggregatedCustomizations.addAll(perProfileCustomizations)
143123
} catch (e: Exception) {
144124
val requestId = (e as? CodeWhispererRuntimeException)?.requestId()
145125
val logMessage = if (CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e)) {
@@ -148,15 +128,18 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
148128
} else {
149129
"ListAvailableCustomizations: failed due to unknown error ${e.message}, requestId: ${requestId.orEmpty()}"
150130
}
151-
152131
LOG.debug { logMessage }
153132
}
154133
}
155134

156135
// 2. get diff
157136
val previousCustomizationsShapshot = connectionToCustomizationsShownLastTime.getOrElse(it.id) { emptyList() }
158137
// calculate new profile+customization list using profile.arn :: customization.arn as the key
159-
val diff = aggregatedCustomizations.filterNot { customization -> previousCustomizationsShapshot.contains("${customization.profile?.arn}::${customization.arn}") }.toSet()
138+
val diff = aggregatedCustomizations.filterNot { customization ->
139+
previousCustomizationsShapshot.contains(
140+
"${customization.profile?.arn}::${customization.arn}"
141+
)
142+
}.toSet()
160143
// 3 if passive,
161144
// (1) update allowlisting
162145
// (2) prompt "You have New Customizations" toast notification (only show once)
@@ -170,7 +153,10 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
170153
}
171154
} else {
172155
aggregatedCustomizations.let { customizations ->
173-
connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization -> "${customization.profile?.arn}::${customization.arn}"}.toMutableList()
156+
connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization,
157+
->
158+
"${customization.profile?.arn}::${customization.arn}"
159+
}.toMutableList()
174160
}
175161
}
176162

@@ -180,8 +166,10 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
180166
activeCustomization(project)?.let { activeCustom ->
181167
if (aggregatedCustomizations.isEmpty()) {
182168
invalidateSelectedAndNotify(project)
183-
} else if (!aggregatedCustomizations.any { latestCustom -> "${latestCustom.profile?.arn}::${latestCustom.arn}" == "${activeCustom.profile?.arn}::${activeCustom.arn}"}
184-
|| activeCustom.profile!= QRegionProfileManager.getInstance().activeProfile(project)){
169+
} else if (!aggregatedCustomizations.any { latestCustom ->
170+
"${latestCustom.profile?.arn}::${latestCustom.arn}" == "${activeCustom.profile?.arn}::${activeCustom.arn}"
171+
} || activeCustom.profile != QRegionProfileManager.getInstance().activeProfile(project)
172+
) {
185173
invalidateSelectedAndNotify(project)
186174
}
187175
}

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -166,27 +166,26 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
166166
(connectionIdToProfileCount[conn.id] ?: 0) > 1
167167
} ?: false
168168

169-
fun getQClientSettings(project: Project): TokenConnectionSettings {
169+
fun getQClientSettings(project: Project, profile: QRegionProfile?): TokenConnectionSettings {
170170
val conn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
171171
if (conn !is AwsBearerTokenConnection) {
172172
error("not a bearer connection")
173173
}
174174

175175
val settings = conn.getConnectionSettings()
176-
val awsRegion = AwsRegionProvider.getInstance()[QEndpoints.Q_DEFAULT_SERVICE_CONFIG.REGION] ?: error("unknown region from Q default service config")
176+
val defaultRegion = AwsRegionProvider.getInstance()[QEndpoints.Q_DEFAULT_SERVICE_CONFIG.REGION] ?: error("unknown region from Q default service config")
177177

178-
// TODO: different window should be able to select different profile
179-
return activeProfile(project)?.let { profile ->
180-
AwsRegionProvider.getInstance()[profile.region]?.let { region ->
181-
settings.withRegion(region)
182-
}
183-
} ?: settings.withRegion(awsRegion)
178+
val regionId = profile?.region ?: activeProfile(project)?.region
179+
val awsRegion = regionId?.let { AwsRegionProvider.getInstance()[it] } ?: defaultRegion
180+
181+
return settings.withRegion(awsRegion)
184182
}
185183

186-
inline fun <reified T : SdkClient> getQClient(project: Project): T = getQClient(project, T::class)
184+
inline fun <reified T : SdkClient> getQClient(project: Project): T = getQClient(project, null, T::class)
185+
inline fun <reified T : SdkClient> getQClient(project: Project, profile: QRegionProfile?): T = getQClient(project, profile, T::class)
187186

188-
fun <T : SdkClient> getQClient(project: Project, sdkClass: KClass<T>): T {
189-
val settings = getQClientSettings(project)
187+
fun <T : SdkClient> getQClient(project: Project, profile: QRegionProfile?, sdkClass: KClass<T>): T {
188+
val settings = getQClientSettings(project, profile)
190189
val client = AwsClientManager.getInstance().getClient(sdkClass, settings)
191190
return client
192191
}

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,5 @@ data class CodeWhispererCustomization(
1515
var description: String? = null,
1616

1717
@JvmField
18-
var profile: QRegionProfile? = null
18+
var profile: QRegionProfile? = null,
1919
)

0 commit comments

Comments
 (0)