Skip to content

Commit 41cdbc7

Browse files
dhasani23David Hasanirli
authored
fix(amazonq): retry S3 upload (#5208)
* fix(amazonq): retry S3 upload * fix test * fix detekt * use waitUntil * update changelog * fix detekt * fix detekt again * remove unused import * fix tests * Update .changes/next-release/bugfix-552a4676-3214-4321-a332-344bdf9f1830.json Co-authored-by: Richard Li <742829+rli@users.noreply.github.com> --------- Co-authored-by: David Hasani <davhasan@amazon.com> Co-authored-by: Richard Li <742829+rli@users.noreply.github.com>
1 parent f0cd876 commit 41cdbc7

File tree

4 files changed

+38
-21
lines changed

4 files changed

+38
-21
lines changed

‎.changes/next-release/bugfix-552a4676-3214-4321-a332-344bdf9f1830.json

+4
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "bugfix",
3+
"description" : "Amazon Q Code Transformation: retry initial project upload on failure"
4+
}

‎plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt

+21-7
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import com.intellij.openapi.application.runInEdt
88
import com.intellij.serviceContainer.AlreadyDisposedException
99
import com.intellij.util.io.HttpRequests
1010
import kotlinx.coroutines.delay
11+
import kotlinx.coroutines.withContext
1112
import org.apache.commons.codec.digest.DigestUtils
1213
import software.amazon.awssdk.core.exception.SdkClientException
1314
import software.amazon.awssdk.services.codewhispererruntime.model.ResumeTransformationResponse
@@ -20,11 +21,13 @@ import software.amazon.awssdk.services.codewhispererruntime.model.Transformation
2021
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationUserActionStatus
2122
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
2223
import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
24+
import software.aws.toolkits.core.utils.Waiters.waitUntil
2325
import software.aws.toolkits.core.utils.error
2426
import software.aws.toolkits.core.utils.exists
2527
import software.aws.toolkits.core.utils.getLogger
2628
import software.aws.toolkits.core.utils.info
2729
import software.aws.toolkits.core.utils.warn
30+
import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
2831
import software.aws.toolkits.jetbrains.services.codemodernizer.client.GumbyClient
2932
import software.aws.toolkits.jetbrains.services.codemodernizer.commands.CodeTransformMessageListener
3033
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerException
@@ -55,6 +58,8 @@ import java.io.File
5558
import java.io.FileInputStream
5659
import java.io.IOException
5760
import java.net.ConnectException
61+
import java.net.SocketTimeoutException
62+
import java.net.UnknownHostException
5863
import java.nio.file.Path
5964
import java.time.Instant
6065
import java.util.Base64
@@ -142,7 +147,7 @@ class CodeModernizerSession(
142147
*
143148
* Based on [CodeWhispererCodeScanSession]
144149
*/
145-
fun createModernizationJob(copyResult: MavenCopyCommandsResult?): CodeModernizerStartJobResult {
150+
suspend fun createModernizationJob(copyResult: MavenCopyCommandsResult?): CodeModernizerStartJobResult {
146151
LOG.info { "Compressing local project" }
147152
val payload: File?
148153
var payloadSize = 0
@@ -377,8 +382,12 @@ class CodeModernizerSession(
377382
/**
378383
* Adapted from [CodeWhispererCodeScanSession]
379384
*/
380-
fun uploadPayload(payload: File): String {
381-
val sha256checksum: String = Base64.getEncoder().encodeToString(DigestUtils.sha256(FileInputStream(payload)))
385+
suspend fun uploadPayload(payload: File): String {
386+
val sha256checksum: String = Base64.getEncoder().encodeToString(
387+
withContext(getCoroutineBgContext()) {
388+
DigestUtils.sha256(FileInputStream(payload))
389+
}
390+
)
382391
if (isDisposed.get()) {
383392
throw AlreadyDisposedException("Disposed when about to create upload URL")
384393
}
@@ -394,17 +403,22 @@ class CodeModernizerSession(
394403
throw AlreadyDisposedException("Disposed when about to upload project artifact to s3")
395404
}
396405
val uploadStartTime = Instant.now()
397-
try {
406+
waitUntil(
407+
exceptionsToIgnore = setOf(
408+
UnknownHostException::class,
409+
SocketTimeoutException::class,
410+
HttpRequests.HttpStatusException::class,
411+
ConnectException::class
412+
)
413+
) {
398414
clientAdaptor.uploadArtifactToS3(
399415
createUploadUrlResponse.uploadUrl(),
400416
payload,
401417
sha256checksum,
402418
createUploadUrlResponse.kmsKeyArn().orEmpty(),
403419
) { shouldStop.get() }
404-
} catch (e: Exception) {
405-
LOG.error { "Unexpected error when uploading project artifact to S3: $e" }
406-
throw e // pass along error to callee
407420
}
421+
LOG.info { "Upload to S3 succeeded" }
408422
if (!shouldStop.get()) {
409423
LOG.info { "Uploaded artifact. Latency: ${calculateTotalLatency(uploadStartTime, Instant.now())}ms" }
410424
}

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

+1
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ class GumbyClient(private val project: Project) {
148148
var result: CodeWhispererRuntimeResponse? = null
149149
try {
150150
result = apiCall()
151+
LOG.info { "$apiName request ID: ${result.responseMetadata()?.requestId()}" }
151152
return result
152153
} catch (e: Exception) {
153154
LOG.error(e) { "$apiName failed: ${e.message}" }

‎plugins/amazonq/codetransform/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codemodernizer/CodeWhispererCodeModernizerSessionTest.kt

+12-14
Original file line numberDiff line numberDiff line change
@@ -410,7 +410,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
410410
}
411411

412412
@Test
413-
fun `CodeModernizer can create modernization job`() {
413+
fun `CodeModernizer can create modernization job`() = runTest {
414414
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
415415
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
416416
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
@@ -425,7 +425,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
425425
}
426426

427427
@Test
428-
fun `CodeModernizer cannot upload payload due to already disposed`() {
428+
fun `CodeModernizer cannot upload payload due to already disposed`() = runTest {
429429
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
430430
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
431431
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
@@ -435,7 +435,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
435435
}
436436

437437
@Test
438-
fun `CodeModernizer returns credentials expired when SsoOidcException during upload`() {
438+
fun `CodeModernizer returns credentials expired when SsoOidcException during upload`() = runTest {
439439
setupConnection(BearerTokenAuthState.AUTHORIZED)
440440
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
441441
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
@@ -445,7 +445,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
445445
}
446446

447447
@Test
448-
fun `CodeModernizer returns credentials expired when expired before upload`() {
448+
fun `CodeModernizer returns credentials expired when expired before upload`() = runTest {
449449
listOf(BearerTokenAuthState.NEEDS_REFRESH, BearerTokenAuthState.NOT_AUTHENTICATED).forEach {
450450
setupConnection(it)
451451
val result = testSessionSpy.createModernizationJob(MavenCopyCommandsResult.Success(File("./mock/path/")))
@@ -454,33 +454,31 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
454454
}
455455

456456
@Test
457-
fun `CodeModernizer cannot upload payload due to presigned url issue`() {
457+
fun `CodeModernizer cannot upload payload due to presigned url issue`() = runTest {
458458
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
459459
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
460460
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
461-
doAnswer { throw HttpRequests.HttpStatusException("mock error", 403, "mock url") }
462-
.whenever(clientAdaptorSpy).uploadArtifactToS3(any(), any(), any(), any(), any())
461+
doAnswer { throw HttpRequests.HttpStatusException("mock error", 403, "mock url") }.whenever(testSessionSpy).uploadPayload(any())
463462
val result = testSessionSpy.createModernizationJob(MavenCopyCommandsResult.Success(File("./mock/path/")))
464463
assertEquals(CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.PRESIGNED_URL_EXPIRED), result)
465464
verify(testSessionStateSpy, times(1)).putJobHistory(any(), eq(TransformationStatus.FAILED), any(), any())
466465
assertEquals(testSessionStateSpy.currentJobStatus, TransformationStatus.FAILED)
467466
}
468467

469468
@Test
470-
fun `CodeModernizer cannot upload payload due to other status code`() {
469+
fun `CodeModernizer cannot upload payload due to other status code`() = runTest {
471470
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
472471
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
473472
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
474-
doAnswer { throw HttpRequests.HttpStatusException("mock error", 407, "mock url") }
475-
.whenever(clientAdaptorSpy).uploadArtifactToS3(any(), any(), any(), any(), any())
473+
doAnswer { throw HttpRequests.HttpStatusException("mock error", 407, "mock url") }.whenever(testSessionSpy).uploadPayload(any())
476474
val result = testSessionSpy.createModernizationJob(MavenCopyCommandsResult.Success(File("./mock/path/")))
477475
assertEquals(CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.HTTP_ERROR(407)), result)
478476
verify(testSessionStateSpy, times(1)).putJobHistory(any(), eq(TransformationStatus.FAILED), any(), any())
479477
assertEquals(testSessionStateSpy.currentJobStatus, TransformationStatus.FAILED)
480478
}
481479

482480
@Test
483-
fun `CodeModernizer cannot upload payload due to unknown issue`() {
481+
fun `CodeModernizer cannot upload payload due to unknown issue`() = runTest {
484482
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
485483
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
486484
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
@@ -492,11 +490,11 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
492490
}
493491

494492
@Test
495-
fun `CodeModernizer cannot upload payload due to connection refused`() {
493+
fun `CodeModernizer cannot upload payload due to connection refused`() = runTest {
496494
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
497495
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
498496
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
499-
doAnswer { throw ConnectException("mock exception") }.whenever(clientAdaptorSpy).uploadArtifactToS3(any(), any(), any(), any(), any())
497+
doAnswer { throw ConnectException("mock exception") }.whenever(testSessionSpy).uploadPayload(any())
500498
val result = testSessionSpy.createModernizationJob(MavenCopyCommandsResult.Success(File("./mock/path/")))
501499
assertEquals(CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.CONNECTION_REFUSED), result)
502500
verify(testSessionStateSpy, times(1)).putJobHistory(any(), eq(TransformationStatus.FAILED), any(), any())
@@ -549,7 +547,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
549547
}
550548

551549
@Test
552-
fun `test uploadPayload()`() {
550+
fun `test uploadPayload()`() = runTest {
553551
val s3endpoint = "http://127.0.0.1:${wireMock.port()}"
554552
val gumbyUploadUrlResponse = CreateUploadUrlResponse.builder()
555553
.uploadUrl(s3endpoint)

0 commit comments

Comments
 (0)