Skip to content

Commit 3805f06

Browse files
author
David Hasani
committed
fix(amazonq): add retries for /transform APIs
1 parent 4685aa0 commit 3805f06

File tree

1 file changed

+31
-6
lines changed
  • plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/client

1 file changed

+31
-6
lines changed

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import com.intellij.openapi.components.Service
77
import com.intellij.openapi.components.service
88
import com.intellij.openapi.project.Project
99
import com.intellij.util.io.HttpRequests
10+
import software.amazon.awssdk.core.exception.SdkException
1011
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
1112
import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeResponse
1213
import software.amazon.awssdk.services.codewhispererruntime.model.ContentChecksumType
@@ -32,15 +33,18 @@ import software.amazon.awssdk.services.codewhispererruntime.model.UploadContext
3233
import software.amazon.awssdk.services.codewhispererruntime.model.UploadIntent
3334
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportContext
3435
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
36+
import software.amazon.awssdk.services.codewhispererstreaming.model.ThrottlingException
3537
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
3638
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationExportContext
39+
import software.amazon.awssdk.services.codewhispererstreaming.model.ValidationException
3740
import software.aws.toolkits.core.utils.error
3841
import software.aws.toolkits.core.utils.getLogger
3942
import software.aws.toolkits.core.utils.info
4043
import software.aws.toolkits.jetbrains.core.AwsClientManager
4144
import software.aws.toolkits.jetbrains.services.amazonq.APPLICATION_ZIP
4245
import software.aws.toolkits.jetbrains.services.amazonq.AWS_KMS
4346
import software.aws.toolkits.jetbrains.services.amazonq.CONTENT_SHA256
47+
import software.aws.toolkits.jetbrains.services.amazonq.RetryableOperation
4448
import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION
4549
import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID
4650
import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient
@@ -52,7 +56,9 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTo
5256
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference
5357
import java.io.File
5458
import java.net.HttpURLConnection
59+
import java.net.SocketTimeoutException
5560
import java.time.Instant
61+
import java.util.concurrent.TimeoutException
5662

5763
@Service(Service.Level.PROJECT)
5864
class GumbyClient(private val project: Project) {
@@ -152,15 +158,34 @@ class GumbyClient(private val project: Project) {
152158
apiCall: () -> T,
153159
apiName: String,
154160
): T {
155-
var result: CodeWhispererRuntimeResponse? = null
161+
var result: T? = null
156162
try {
157-
result = apiCall()
158-
LOG.info { "$apiName request ID: ${result.responseMetadata()?.requestId()}" }
159-
return result
163+
RetryableOperation<Unit>().execute(
164+
operation = {
165+
result = apiCall()
166+
},
167+
isRetryable = { e ->
168+
when (e) {
169+
is ValidationException,
170+
is ThrottlingException,
171+
is SdkException,
172+
is TimeoutException,
173+
is SocketTimeoutException,
174+
-> true
175+
else -> false
176+
}
177+
},
178+
errorHandler = { e, attempts ->
179+
LOG.error(e) { "After $attempts attempts, $apiName failed: ${e.message}" }
180+
throw e
181+
}
182+
)
160183
} catch (e: Exception) {
161-
LOG.error(e) { "$apiName failed: ${e.message}" }
162-
throw e // pass along error to callee
184+
LOG.error(e) { "$apiName failed: ${e.message}; may have been retried up to 3 times" }
185+
throw e
163186
}
187+
LOG.info { "$apiName request ID: ${result?.responseMetadata()?.requestId()}" }
188+
return result ?: error("$apiName failed")
164189
}
165190

166191
suspend fun downloadExportResultArchive(

0 commit comments

Comments
 (0)